902 lines
42 KiB
Python
902 lines
42 KiB
Python
# train_vibevoice_lora.py
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from datasets import load_dataset, DatasetDict, VerificationMode
|
|
|
|
from transformers import (
|
|
HfArgumentParser,
|
|
Trainer,
|
|
set_seed,
|
|
TrainerCallback,
|
|
)
|
|
from transformers import TrainingArguments as HfTrainingArguments
|
|
|
|
from peft import LoraConfig, get_peft_model, TaskType
|
|
|
|
from vibevoice.modular.modeling_vibevoice import VibeVoiceForConditionalGeneration
|
|
from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
|
|
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
|
|
|
|
from vibevoice.finetune.data_vibevoice import VibeVoiceDataset, VibeVoiceCollator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ================== SAMPLE CALLBACK UTILS ==================
|
|
|
|
import copy
|
|
import torch
|
|
from transformers import TrainerCallback
|
|
|
|
class EmaCallback(TrainerCallback):
|
|
def __init__(self, attr_path="model.prediction_head", decay=0.999, device="cpu"):
|
|
"""
|
|
attr_path: where the head lives under self.model (Trainer wraps your VibeVoiceForConditionalGeneration)
|
|
decay: EMA decay (0.999 ~ stable, 0.9999 ~ very smooth, slower to adapt)
|
|
"""
|
|
self.attr_path = attr_path
|
|
self.decay = float(decay)
|
|
self.device = torch.device(device)
|
|
self.shadow = None
|
|
self._orig = None # store non-EMA weights when we swap
|
|
|
|
def _get_module(self, model):
|
|
# Resolve dotted path like "model.prediction_head"
|
|
mod = model
|
|
for name in self.attr_path.split('.'):
|
|
mod = getattr(mod, name)
|
|
return mod
|
|
|
|
def on_train_begin(self, args, state, control, model=None, **kwargs):
|
|
head = self._get_module(model)
|
|
self.shadow = {k: p.detach().to(self.device).clone()
|
|
for k, p in head.state_dict().items()}
|
|
|
|
def on_step_end(self, args, state, control, model=None, **kwargs):
|
|
if self.shadow is None: return
|
|
head = self._get_module(model)
|
|
with torch.no_grad():
|
|
for k, v in head.state_dict().items():
|
|
self.shadow[k].mul_(self.decay).add_(v.detach().to(self.device), alpha=(1.0 - self.decay))
|
|
|
|
# ---- Swap helpers ----
|
|
def _swap_in_ema(self, model):
|
|
head = self._get_module(model)
|
|
self._orig = copy.deepcopy(head.state_dict())
|
|
head.load_state_dict(self.shadow, strict=False)
|
|
|
|
def _swap_back(self, model):
|
|
if self._orig is None: return
|
|
head = self._get_module(model)
|
|
head.load_state_dict(self._orig, strict=False)
|
|
self._orig = None
|
|
|
|
def on_evaluate(self, args, state, control, model=None, **kwargs):
|
|
# use EMA during eval
|
|
self._swap_in_ema(model)
|
|
|
|
def on_evaluate_end(self, args, state, control, model=None, **kwargs):
|
|
self._swap_back(model)
|
|
|
|
def on_save(self, args, state, control, model=None, **kwargs):
|
|
# temporarily swap to EMA, let Trainer save, then swap back
|
|
self._swap_in_ema(model)
|
|
|
|
def on_save_end(self, args, state, control, model=None, **kwargs):
|
|
self._swap_back(model)
|
|
|
|
def on_train_end(self, args, state, control, model=None, **kwargs):
|
|
# final checkpoint: persist EMA
|
|
self._swap_in_ema(model)
|
|
|
|
|
|
@dataclass
|
|
class ModelArguments:
|
|
model_name_or_path: Optional[str] = field(
|
|
default=None, metadata={"help": "Path to VibeVoice base model with config.json"}
|
|
)
|
|
processor_name_or_path: Optional[str] = field(
|
|
default=None, metadata={"help": "Path to processor dir (preprocessor_config.json). Defaults to model path."}
|
|
)
|
|
cache_dir: Optional[str] = field(default=None)
|
|
freeze_acoustic_tokenizer: bool = field(default=True)
|
|
freeze_semantic_tokenizer: bool = field(default=True)
|
|
lora_r: int = field(default=8)
|
|
lora_alpha: int = field(default=32)
|
|
lora_dropout: float = field(default=0.05)
|
|
lora_target_modules: str = field(
|
|
default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
|
|
metadata={"help": "Comma-separated list of target module names in the LLM blocks"},
|
|
)
|
|
lora_wrap_diffusion_head: bool = field(default=False, metadata={"help": "Wrap diffusion head with PEFT LoRA"})
|
|
train_diffusion_head: bool = field(default=False, metadata={"help": "Train diffusion prediction head (full fine-tune)"})
|
|
train_connectors: bool = field(default=False, metadata={"help": "Train acoustic/semantic connectors (full fine-tune)"})
|
|
layers_to_freeze: Optional[str] = field(
|
|
default=None,
|
|
metadata={"help": "Comma-separated indices of diffusion head layers to freeze (e.g., '0,1,5,7,8')."}
|
|
)
|
|
|
|
@dataclass
|
|
class DataArguments:
|
|
dataset_name: Optional[str] = field(default=None, metadata={"help": "HF dataset name or 'json' with --train_jsonl for local files"})
|
|
dataset_config_name: Optional[str] = field(default=None)
|
|
train_split_name: str = field(default="train")
|
|
eval_split_name: Optional[str] = field(default="validation")
|
|
text_column_name: str = field(default="text")
|
|
audio_column_name: str = field(default="audio")
|
|
voice_prompts_column_name: Optional[str] = field(default="voice_prompts")
|
|
eval_split_size: float = field(default=0.0)
|
|
ignore_verifications: bool = field(default=False)
|
|
max_length: Optional[int] = field(default=None)
|
|
train_jsonl: Optional[str] = field(default=None, metadata={"help": "Path to local train JSONL with {text, audio, [voice_prompts]}"})
|
|
validation_jsonl: Optional[str] = field(default=None, metadata={"help": "Optional path to local validation JSONL"})
|
|
voice_prompt_drop_rate: float = field(
|
|
default=0.0,
|
|
metadata={"help": "Probability to drop conditioning voice prompt during training (0.0 keep always, 1.0 drop always)."},
|
|
)
|
|
|
|
@dataclass
|
|
class CustomTrainingArguments(HfTrainingArguments):
|
|
ddpm_batch_mul: int = field(default=1)
|
|
ce_loss_weight: float = field(default=1.0)
|
|
diffusion_loss_weight: float = field(default=1.0)
|
|
debug_ce_details: bool = field(default=False)
|
|
debug_ce_topk: int = field(default=5)
|
|
debug_ce_max_examples: int = field(default=1)
|
|
debug_ce_every_n_steps: int = field(default=200)
|
|
gradient_clipping: bool = field(
|
|
default=False,
|
|
metadata={"help": "Enable gradient clipping using max_grad_norm (set via --max_grad_norm, default 1.0). When False, disables clipping by forcing max_grad_norm=0.0."},
|
|
)
|
|
debug_save: bool = field(
|
|
default=False,
|
|
metadata={"help": "If set, saves model components BEFORE training starts, into output_dir/debug_initial."},
|
|
)
|
|
|
|
def build_lora_config(args: ModelArguments) -> LoraConfig:
|
|
target_modules = [s.strip() for s in args.lora_target_modules.split(",") if s.strip()]
|
|
return LoraConfig(
|
|
r=args.lora_r,
|
|
lora_alpha=args.lora_alpha,
|
|
lora_dropout=args.lora_dropout,
|
|
bias="none",
|
|
task_type=TaskType.CAUSAL_LM,
|
|
target_modules=target_modules,
|
|
)
|
|
|
|
def build_head_lora_config(args: ModelArguments) -> LoraConfig:
|
|
target_modules = ["noisy_images_proj","cond_proj","gate_proj","up_proj","down_proj","linear"]
|
|
return LoraConfig(
|
|
r=args.lora_r,
|
|
lora_alpha=args.lora_alpha,
|
|
lora_dropout=args.lora_dropout,
|
|
bias="none",
|
|
task_type=TaskType.FEATURE_EXTRACTION,
|
|
target_modules=target_modules,
|
|
)
|
|
|
|
def mask_for_ce(labels: torch.Tensor, attention_mask: torch.Tensor, acoustic_input_mask: torch.Tensor, pad_id: int = -100) -> torch.Tensor:
|
|
shifted = labels[:, 1:].contiguous()
|
|
base_mask = attention_mask[:, 1:].contiguous().eq(1) if (attention_mask is not None and attention_mask.numel() > 0) else torch.ones_like(shifted, dtype=torch.bool)
|
|
label_is_acoustic = acoustic_input_mask[:, 1:].contiguous()
|
|
final_mask = base_mask & (~label_is_acoustic)
|
|
out = shifted.clone()
|
|
out[~final_mask] = pad_id
|
|
return out
|
|
|
|
def _patch_acoustic_encode_for_legacy_indexing(model_obj, logger_):
|
|
try:
|
|
acoustic = getattr(getattr(model_obj, "model", model_obj), "acoustic_tokenizer", None)
|
|
if acoustic is None or not hasattr(acoustic, "encode"):
|
|
logger_.warning("No acoustic_tokenizer.encode() found to patch.")
|
|
return
|
|
base_encode = acoustic.encode
|
|
def encode_wrapped(*args, **kwargs):
|
|
out = base_encode(*args, **kwargs)
|
|
try:
|
|
_ = out[0][0]
|
|
return out
|
|
except Exception:
|
|
pass
|
|
if isinstance(out, dict):
|
|
for k in ("frames", "codes", "tokens", "latents", "hidden_states"):
|
|
if k in out:
|
|
return [[out[k]]]
|
|
if len(out) > 0:
|
|
return [[next(iter(out.values()))]]
|
|
for attr in ("frames", "codes", "tokens", "latents", "hidden_states"):
|
|
if hasattr(out, attr):
|
|
return [[getattr(out, attr)]]
|
|
try:
|
|
if isinstance(out, torch.Tensor):
|
|
return [[out]]
|
|
except Exception:
|
|
pass
|
|
return [[out]]
|
|
acoustic.encode = encode_wrapped
|
|
logger_.info("Patched acoustic_tokenizer.encode() to return [[...]] for legacy indexing.")
|
|
except Exception as e:
|
|
logger_.warning(f"Failed to patch acoustic_tokenizer.encode(): {e}")
|
|
|
|
def main() -> None:
|
|
parser = HfArgumentParser((ModelArguments, DataArguments, CustomTrainingArguments))
|
|
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
|
|
|
logging.basicConfig(
|
|
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
|
datefmt="%m/%d/%Y %H:%M:%S",
|
|
level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
|
|
)
|
|
logger.info("Training/evaluation parameters %s", training_args)
|
|
set_seed(training_args.seed)
|
|
|
|
# Configure gradient clipping
|
|
if not getattr(training_args, "gradient_clipping", False):
|
|
if hasattr(training_args, "max_grad_norm"):
|
|
training_args.max_grad_norm = 0.0
|
|
logger.info("Gradient clipping disabled (set max_grad_norm=0.0). Use --gradient_clipping to enable.")
|
|
else:
|
|
if (not hasattr(training_args, "max_grad_norm")) or training_args.max_grad_norm is None or training_args.max_grad_norm <= 0:
|
|
training_args.max_grad_norm = 1.0
|
|
logger.info(f"Gradient clipping enabled: max_grad_norm={training_args.max_grad_norm}")
|
|
|
|
# Load processor
|
|
processor_path = model_args.processor_name_or_path or model_args.model_name_or_path
|
|
if processor_path is None:
|
|
raise ValueError("--model_name_or_path (or --processor_name_or_path) must be provided")
|
|
processor: VibeVoiceProcessor = VibeVoiceProcessor.from_pretrained(processor_path)
|
|
|
|
# Required special tokens
|
|
tok = processor.tokenizer
|
|
for required in ["speech_start_id", "speech_diffusion_id", "speech_end_id"]:
|
|
if not hasattr(tok, required) or getattr(tok, required) is None:
|
|
raise RuntimeError(f"Tokenizer missing required special id: {required}")
|
|
|
|
# Load model
|
|
if model_args.model_name_or_path is None:
|
|
raise ValueError("--model_name_or_path is required to load VibeVoice base model")
|
|
dtype = torch.float32
|
|
if training_args.bf16:
|
|
dtype = torch.bfloat16
|
|
elif getattr(training_args, "fp16", False):
|
|
dtype = torch.float16
|
|
model = VibeVoiceForConditionalGeneration.from_pretrained(
|
|
model_args.model_name_or_path,
|
|
torch_dtype=dtype,
|
|
)
|
|
_patch_acoustic_encode_for_legacy_indexing(model, logger)
|
|
processor.semantic_tokenizer = getattr(model.model, "semantic_tokenizer", None)
|
|
|
|
# Diagnostics: LM head tie
|
|
try:
|
|
in_emb_mod = model.get_input_embeddings()
|
|
out_emb_mod = model.get_output_embeddings()
|
|
in_w = getattr(in_emb_mod, "weight", None)
|
|
out_w = getattr(out_emb_mod, "weight", None)
|
|
shared_ptr = bool(in_w is not None and out_w is not None and in_w.data_ptr() == out_w.data_ptr())
|
|
values_equal = False
|
|
if in_w is not None and out_w is not None and in_w.shape == out_w.shape:
|
|
try:
|
|
values_equal = bool(torch.allclose(in_w, out_w))
|
|
except Exception:
|
|
values_equal = False
|
|
try:
|
|
tie_cfg = getattr(getattr(model.config, "decoder_config", model.config), "tie_word_embeddings", None)
|
|
except Exception:
|
|
tie_cfg = getattr(model.config, "tie_word_embeddings", None)
|
|
logger.info(f"LM head diagnostics -> shared_params={shared_ptr}, values_equal={values_equal}, tie_word_embeddings={tie_cfg}")
|
|
if out_w is not None:
|
|
logger.info(f"LM head requires_grad before freeze: {bool(out_w.requires_grad)}")
|
|
except Exception as e:
|
|
logger.warning(f"LM head tie diagnostics failed: {e}")
|
|
|
|
# Hard-tie LM head
|
|
try:
|
|
emb_module = model.get_input_embeddings()
|
|
head_module = model.get_output_embeddings()
|
|
if hasattr(emb_module, "weight") and hasattr(head_module, "weight"):
|
|
if emb_module.weight.shape == head_module.weight.shape and emb_module.weight.data_ptr() != head_module.weight.data_ptr():
|
|
with torch.no_grad():
|
|
head_module.weight = emb_module.weight
|
|
logger.info("Force-tied LM head weight to input embeddings (pointer share).")
|
|
except Exception as e:
|
|
logger.warning(f"Force-tie of LM head failed: {e}")
|
|
|
|
# Validate special IDs (info logs only)
|
|
try:
|
|
special_names = ["speech_start_id", "speech_diffusion_id", "speech_end_id"]
|
|
try:
|
|
vocab_size = int(getattr(model.config.decoder_config, "vocab_size", 0))
|
|
except Exception:
|
|
vocab_size = 0
|
|
in_emb_mod = model.get_input_embeddings()
|
|
out_emb_mod = model.get_output_embeddings()
|
|
in_w = getattr(in_emb_mod, "weight", None)
|
|
out_w = getattr(out_emb_mod, "weight", None)
|
|
for name in special_names:
|
|
val = getattr(tok, name, None)
|
|
exists = (val is not None)
|
|
in_range = (exists and isinstance(val, int) and 0 <= val < vocab_size)
|
|
equal_row = None
|
|
if in_range and in_w is not None and out_w is not None and in_w.shape == out_w.shape and in_w.size(0) > val:
|
|
try:
|
|
equal_row = bool(torch.allclose(in_w[val], out_w[val]))
|
|
except Exception:
|
|
equal_row = False
|
|
decoded_str = None
|
|
if exists and isinstance(val, int):
|
|
try:
|
|
decoded_str = tok.decode([val])
|
|
except Exception:
|
|
try:
|
|
decoded_str = tok.convert_ids_to_tokens(val)
|
|
except Exception:
|
|
decoded_str = "<decode_failed>"
|
|
logger.info(f"Special token check -> {name}={val}, decoded='{decoded_str}', exists={exists}, in_vocab_range={in_range}, emb_vs_head_row_equal={equal_row}")
|
|
except Exception as e:
|
|
logger.warning(f"Special token ID/row validation failed: {e}")
|
|
|
|
# Quick tokenizer diagnostics (optional)
|
|
try:
|
|
logger.info("=== TOKENIZER DIAGNOSTICS ===")
|
|
logger.info(f"Tokenizer class: {type(tok).__name__}")
|
|
logger.info(f"Tokenizer vocab_size: {tok.vocab_size}")
|
|
# tiny CE smoke test
|
|
with torch.no_grad():
|
|
simple_text = "The cat sat on the mat."
|
|
simple_ids = torch.tensor([tok.encode(simple_text, add_special_tokens=True)], device=model.device)
|
|
simple_mask = torch.ones_like(simple_ids)
|
|
x = model.get_input_embeddings()(simple_ids)
|
|
outputs = model.model(inputs_embeds=x, attention_mask=simple_mask, return_dict=True)
|
|
logits = model.lm_head(outputs.last_hidden_state)
|
|
shift_logits = logits[:, :-1, :].contiguous()
|
|
shift_labels = simple_ids[:, 1:].contiguous()
|
|
ce_loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), reduction='mean')
|
|
logger.info(f"Simple text CE loss: {ce_loss.item():.4f}")
|
|
except Exception as e:
|
|
logger.warning(f"Tokenizer diagnostics failed: {e}")
|
|
|
|
# Disable cache during training
|
|
if hasattr(model.config, "use_cache") and training_args.do_train:
|
|
model.config.use_cache = False
|
|
|
|
# Freeze tokenizers
|
|
if model_args.freeze_acoustic_tokenizer and hasattr(model.model, "acoustic_tokenizer"):
|
|
for p in model.model.acoustic_tokenizer.parameters():
|
|
p.requires_grad = False
|
|
if model_args.freeze_semantic_tokenizer and hasattr(model.model, "semantic_tokenizer"):
|
|
for p in model.model.semantic_tokenizer.parameters():
|
|
p.requires_grad = False
|
|
|
|
# LoRA wrap LLM (optional)
|
|
lora_cfg = build_lora_config(model_args)
|
|
tm_lower = [s.strip().lower() for s in model_args.lora_target_modules.split(",") if s.strip()]
|
|
skip_lm_lora = (len(tm_lower) == 0) or all(t in ("none", "off", "disable", "disabled") for t in tm_lower)
|
|
if not skip_lm_lora:
|
|
model.model.language_model = get_peft_model(model.model.language_model, lora_cfg)
|
|
else:
|
|
logger.info("Skipping LLM LoRA wrapping (lora_target_modules indicates none).")
|
|
|
|
try:
|
|
model.tie_weights()
|
|
except Exception:
|
|
pass
|
|
|
|
# Freeze all then enable trainable subsets
|
|
for _, p in model.named_parameters():
|
|
p.requires_grad = False
|
|
|
|
try:
|
|
for n, p in model.model.language_model.named_parameters():
|
|
if "lora_A" in n or "lora_B" in n:
|
|
p.requires_grad = True
|
|
except Exception:
|
|
logger.warning("Could not re-enable LoRA params on language_model.")
|
|
|
|
# Diffusion head LoRA wrapping (optional)
|
|
if getattr(model_args, "lora_wrap_diffusion_head", False) and hasattr(model.model, "prediction_head"):
|
|
class _HeadForwardShim(nn.Module):
|
|
def __init__(self, base: nn.Module): super().__init__(); self.base = base
|
|
def forward(self, *args, **kwargs):
|
|
if len(args) >= 3:
|
|
noisy_images, timesteps, condition = args[:3]
|
|
else:
|
|
noisy_images = kwargs.get("noisy_images")
|
|
timesteps = kwargs.get("timesteps")
|
|
condition = kwargs.get("condition")
|
|
return self.base(noisy_images, timesteps, condition)
|
|
try:
|
|
shim = _HeadForwardShim(model.model.prediction_head)
|
|
model.model.prediction_head = get_peft_model(shim, build_head_lora_config(model_args))
|
|
for n, p in model.model.prediction_head.named_parameters():
|
|
if "lora_A" in n or "lora_B" in n:
|
|
p.requires_grad = True
|
|
except Exception as e:
|
|
logger.warning(f"Could not LoRA-wrap diffusion head: {e}")
|
|
|
|
# Train full diffusion head (optional)
|
|
if getattr(model_args, "train_diffusion_head", False) and hasattr(model.model, "prediction_head"):
|
|
for p in model.model.prediction_head.parameters():
|
|
p.requires_grad = True
|
|
|
|
# Freeze diffusion head layers (optional)
|
|
if model_args.layers_to_freeze is not None and hasattr(model.model, "prediction_head"):
|
|
head_params = list(model.model.prediction_head.named_parameters())
|
|
try:
|
|
indices_to_freeze = {int(x.strip()) for x in model_args.layers_to_freeze.split(',') if x.strip()}
|
|
frozen_count = 0
|
|
for i, (name, param) in enumerate(head_params):
|
|
if i in indices_to_freeze:
|
|
param.requires_grad = False
|
|
frozen_count += 1
|
|
logger.info(f"Froze layer [{i}]: {name}")
|
|
logger.info(f"Successfully froze {frozen_count} parameter groups in the diffusion head.")
|
|
except Exception as e:
|
|
logger.error(f"Could not parse --layers_to_freeze: {e}")
|
|
raise
|
|
|
|
# Connectors
|
|
if getattr(model_args, "train_connectors", False):
|
|
if hasattr(model.model, "acoustic_connector"):
|
|
for p in model.model.acoustic_connector.parameters():
|
|
p.requires_grad = True
|
|
if hasattr(model.model, "semantic_connector"):
|
|
for p in model.model.semantic_connector.parameters():
|
|
p.requires_grad = True
|
|
else:
|
|
if hasattr(model.model, "acoustic_connector"):
|
|
for p in model.model.acoustic_connector.parameters():
|
|
p.requires_grad = False
|
|
if hasattr(model.model, "semantic_connector"):
|
|
for p in model.model.semantic_connector.parameters():
|
|
p.requires_grad = False
|
|
|
|
# Freeze embedding + head
|
|
try:
|
|
emb = model.get_input_embeddings()
|
|
if hasattr(emb, "weight"):
|
|
emb.weight.requires_grad_(False)
|
|
head = model.get_output_embeddings()
|
|
if head is not None and hasattr(head, "weight"):
|
|
head.weight.requires_grad_(False)
|
|
except Exception:
|
|
pass
|
|
|
|
# Diagnostics
|
|
def _sum_params(named_iter):
|
|
return sum(p.numel() for _, p in named_iter if p.requires_grad)
|
|
try:
|
|
lm_lora = _sum_params(model.model.language_model.named_parameters()) if hasattr(model.model, "language_model") else 0
|
|
pred_head_train = _sum_params(model.model.prediction_head.named_parameters()) if hasattr(model.model, "prediction_head") else 0
|
|
ac_conn_train = _sum_params(model.model.acoustic_connector.named_parameters()) if hasattr(model.model, "acoustic_connector") else 0
|
|
se_conn_train = _sum_params(model.model.semantic_connector.named_parameters()) if hasattr(model.model, "semantic_connector") else 0
|
|
total_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
logger.info(f"Trainable by block -> LLM-LoRA: {lm_lora:,} | diff_head: {pred_head_train:,} | ac_conn: {ac_conn_train:,} | se_conn: {se_conn_train:,}")
|
|
logger.info("TOTAL trainable: %s", f"{total_trainable:,}")
|
|
except Exception:
|
|
pass
|
|
|
|
# Datasets
|
|
verification_mode = VerificationMode.NO_CHECKS if data_args.ignore_verifications else VerificationMode.BASIC_CHECKS
|
|
if data_args.train_jsonl is not None:
|
|
data_files: Dict[str, str] = {"train": data_args.train_jsonl}
|
|
if data_args.validation_jsonl is not None:
|
|
data_files["validation"] = data_args.validation_jsonl
|
|
raw = load_dataset("json", data_files=data_files, verification_mode=verification_mode, cache_dir=model_args.cache_dir)
|
|
else:
|
|
if data_args.dataset_name is None:
|
|
raise ValueError("Provide --dataset_name (HF datasets) or use --train_jsonl/--validation_jsonl for local files.")
|
|
raw = load_dataset(
|
|
data_args.dataset_name,
|
|
data_args.dataset_config_name,
|
|
verification_mode=verification_mode,
|
|
cache_dir=model_args.cache_dir,
|
|
)
|
|
train_ds = raw[data_args.train_split_name]
|
|
eval_ds = None
|
|
if training_args.do_eval:
|
|
if data_args.eval_split_name and data_args.eval_split_name in raw:
|
|
eval_ds = raw[data_args.eval_split_name]
|
|
elif data_args.eval_split_size and data_args.eval_split_size > 0 and len(train_ds) > 1:
|
|
split = train_ds.train_test_split(test_size=data_args.eval_split_size, seed=training_args.seed)
|
|
train_ds, eval_ds = split["train"], split["test"]
|
|
|
|
train_dataset = VibeVoiceDataset(
|
|
train_ds,
|
|
text_column=data_args.text_column_name,
|
|
audio_column=data_args.audio_column_name,
|
|
voice_prompts_column=data_args.voice_prompts_column_name,
|
|
)
|
|
eval_dataset = None
|
|
if eval_ds is not None:
|
|
eval_dataset = VibeVoiceDataset(
|
|
eval_ds,
|
|
text_column=data_args.text_column_name,
|
|
audio_column=data_args.audio_column_name,
|
|
voice_prompts_column=data_args.voice_prompts_column_name,
|
|
)
|
|
|
|
# Ratios/dims from processor+model
|
|
speech_compress_ratio = getattr(processor, "speech_tok_compress_ratio", 3200)
|
|
semantic_dim = getattr(model.config, "semantic_vae_dim", None)
|
|
if semantic_dim is None:
|
|
try:
|
|
semantic_dim = int(getattr(model.config.semantic_tokenizer_config, "vae_dim", 128))
|
|
except Exception:
|
|
semantic_dim = 128
|
|
|
|
compute_semantics_flag = hasattr(processor, "semantic_tokenizer") and processor.semantic_tokenizer is not None
|
|
|
|
data_collator = VibeVoiceCollator(
|
|
processor=processor,
|
|
max_length=data_args.max_length,
|
|
speech_compress_ratio=speech_compress_ratio,
|
|
semantic_vae_dim=semantic_dim,
|
|
compute_semantics=compute_semantics_flag,
|
|
debug_checks=False,
|
|
voice_prompt_drop_rate=data_args.voice_prompt_drop_rate,
|
|
)
|
|
|
|
class LoRADebugCallback(TrainerCallback):
|
|
def __init__(self, log_every_n_steps: int = 50):
|
|
self.log_every_n_steps = max(1, int(log_every_n_steps))
|
|
self.prev_param_norms: Dict[str, float] = {}
|
|
self.lora_param_names: List[str] = []
|
|
|
|
def on_train_begin(self, args, state, control, model=None, **kwargs):
|
|
try:
|
|
if model is None:
|
|
return
|
|
named: Dict[str, torch.nn.Parameter] = dict(model.named_parameters())
|
|
self.lora_param_names = [n for n in named.keys() if ("lora_A" in n or "lora_B" in n)]
|
|
for n in self.lora_param_names:
|
|
p = named[n]
|
|
self.prev_param_norms[n] = float(p.data.norm().item())
|
|
total = len(self.lora_param_names)
|
|
req_grad = sum(1 for n in self.lora_param_names if named[n].requires_grad)
|
|
num_A = sum(1 for n in self.lora_param_names if "lora_A" in n)
|
|
num_B = sum(1 for n in self.lora_param_names if "lora_B" in n)
|
|
zero_B = sum(1 for n in self.lora_param_names if ("lora_B" in n and float(named[n].data.norm().item()) == 0.0))
|
|
logger.info(f"LoRA debug: found {total} LoRA params (A={num_A}, B={num_B}); trainable={req_grad}. Initial lora_B_zero={zero_B}.")
|
|
if total == 0:
|
|
logger.warning("LoRA debug: No LoRA parameters found. Check lora_target_modules.")
|
|
if req_grad != total:
|
|
logger.warning("LoRA debug: Some LoRA params are frozen. They should be trainable.")
|
|
except Exception as e:
|
|
logger.warning(f"LoRA debug (on_train_begin) failed: {e}")
|
|
|
|
def on_step_end(self, args, state, control, model=None, **kwargs):
|
|
try:
|
|
if model is None or len(self.lora_param_names) == 0:
|
|
return
|
|
step = int(getattr(state, "global_step", 0) or 0)
|
|
if step % self.log_every_n_steps != 0 and step != 1:
|
|
return
|
|
named: Dict[str, torch.nn.Parameter] = dict(model.named_parameters())
|
|
changed_A = 0
|
|
changed_B = 0
|
|
zero_B = 0
|
|
eps = 1e-12
|
|
for n in self.lora_param_names:
|
|
p = named.get(n, None)
|
|
if p is None:
|
|
continue
|
|
prev = self.prev_param_norms.get(n, 0.0)
|
|
curr = float(p.data.norm().item())
|
|
if "lora_A" in n and abs(curr - prev) > eps:
|
|
changed_A += 1
|
|
if "lora_B" in n:
|
|
if abs(curr - prev) > eps:
|
|
changed_B += 1
|
|
if curr == 0.0:
|
|
zero_B += 1
|
|
self.prev_param_norms[n] = curr
|
|
total_A = sum(1 for n in self.lora_param_names if "lora_A" in n)
|
|
total_B = sum(1 for n in self.lora_param_names if "lora_B" in n)
|
|
logger.info(f"LoRA debug step {step}: changed A {changed_A}/{total_A}, changed B {changed_B}/{total_B}, lora_B_zero_now={zero_B}.")
|
|
except Exception as e:
|
|
logger.warning(f"LoRA debug (on_step_end) failed: {e}")
|
|
|
|
class VibeVoiceTrainer(Trainer):
|
|
def compute_loss(self, model: VibeVoiceForConditionalGeneration, inputs: Dict[str, Any], return_outputs=False, num_items_in_batch: Optional[int] = None):
|
|
labels = inputs.get("input_ids")
|
|
attention_mask = inputs.get("attention_mask")
|
|
acoustic_input_mask = inputs.get("acoustic_input_mask")
|
|
|
|
# Ensure semantic tensors exist and have correct dtype/device
|
|
sem = inputs.get("speech_semantic_tensors", None)
|
|
try:
|
|
target_dtype = next(model.model.semantic_connector.parameters()).dtype
|
|
except Exception:
|
|
target_dtype = model.get_input_embeddings().weight.dtype
|
|
|
|
if sem is None:
|
|
sm = inputs.get("speech_masks")
|
|
if sm is not None:
|
|
zeros = torch.zeros(
|
|
sm.size(0), sm.size(1),
|
|
getattr(model.config, "semantic_vae_dim", 128),
|
|
dtype=target_dtype,
|
|
device=sm.device,
|
|
)
|
|
inputs["speech_semantic_tensors"] = zeros
|
|
else:
|
|
if isinstance(sem, torch.Tensor):
|
|
inputs["speech_semantic_tensors"] = sem.to(dtype=target_dtype)
|
|
|
|
outputs = model(
|
|
input_ids=inputs.get("input_ids"),
|
|
attention_mask=attention_mask,
|
|
speech_tensors=inputs.get("speech_tensors"),
|
|
speech_masks=inputs.get("speech_masks"),
|
|
speech_semantic_tensors=inputs.get("speech_semantic_tensors"),
|
|
acoustic_input_mask=acoustic_input_mask,
|
|
acoustic_loss_mask=inputs.get("acoustic_loss_mask"),
|
|
speeches_loss_input=inputs.get("speeches_loss_input"),
|
|
ddpm_batch_mul=training_args.ddpm_batch_mul,
|
|
)
|
|
|
|
# Invariants: token/latent selection equality across views (warn, don't assert)
|
|
try:
|
|
al_mask = inputs.get("acoustic_loss_mask")
|
|
sp_masks = inputs.get("speech_masks")
|
|
sp_loss_sel = inputs.get("speeches_loss_input")
|
|
num_tok_total = int(acoustic_input_mask.sum().item()) if acoustic_input_mask is not None else 0
|
|
num_tok_loss = int(al_mask.sum().item()) if al_mask is not None else 0
|
|
num_lat_total = int(sp_masks.sum().item()) if sp_masks is not None else 0
|
|
num_lat_loss = int(((sp_loss_sel & sp_masks).sum().item())) if (sp_loss_sel is not None and sp_masks is not None) else 0
|
|
self.log({
|
|
"debug/num_tok_total": float(num_tok_total),
|
|
"debug/num_tok_loss": float(num_tok_loss),
|
|
"debug/num_lat_total": float(num_lat_total),
|
|
"debug/num_lat_loss": float(num_lat_loss),
|
|
})
|
|
if sp_loss_sel is not None and sp_masks is not None and al_mask is not None:
|
|
if num_tok_loss != num_lat_loss:
|
|
logger.warning(f"Loss selection mismatch: acoustic_loss_mask={num_tok_loss} vs speeches_loss_input={num_lat_loss}")
|
|
except Exception:
|
|
pass
|
|
|
|
# CE Loss
|
|
logits = outputs.logits
|
|
ce_labels = mask_for_ce(labels, attention_mask, acoustic_input_mask, pad_id=-100)
|
|
shift_logits = logits[:, :-1, :].contiguous()
|
|
loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
|
|
ce_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), ce_labels.view(-1))
|
|
|
|
# Optional CE diagnostics
|
|
try:
|
|
self._debug_ce(shift_logits, ce_labels, attention_mask, acoustic_input_mask)
|
|
except Exception as e:
|
|
logger.warning(f"Failed invoking CE debug: {e}")
|
|
|
|
# Diffusion loss
|
|
diffusion_loss = outputs.diffusion_loss if outputs.diffusion_loss is not None else torch.tensor(0.0, device=ce_loss.device)
|
|
total = training_args.ce_loss_weight * ce_loss + training_args.diffusion_loss_weight * diffusion_loss
|
|
|
|
# Logs
|
|
try:
|
|
prefix = "train" if model.training else "eval"
|
|
self.log({
|
|
f"{prefix}/ce_loss": ce_loss.detach().item(),
|
|
f"{prefix}/diffusion_loss": diffusion_loss.detach().item() if isinstance(diffusion_loss, torch.Tensor) else float(diffusion_loss),
|
|
})
|
|
if hasattr(self, "optimizer") and self.optimizer is not None and len(self.optimizer.param_groups) > 0:
|
|
lr_val = self.optimizer.param_groups[0].get("lr", None)
|
|
if lr_val is not None:
|
|
self.log({"train/learning_rate_real": float(lr_val)})
|
|
except Exception:
|
|
pass
|
|
|
|
return (total, outputs) if return_outputs else total
|
|
|
|
def _debug_ce(self, shift_logits: torch.Tensor, ce_labels: torch.Tensor, attention_mask: Optional[torch.Tensor], acoustic_input_mask: Optional[torch.Tensor]):
|
|
try:
|
|
if not getattr(training_args, "debug_ce_details", False):
|
|
return
|
|
step = int(getattr(self.state, "global_step", 0) or 0)
|
|
every_n = max(1, int(getattr(training_args, "debug_ce_every_n_steps", 200) or 200))
|
|
if not (step <= 1 or (step % every_n == 0)):
|
|
return
|
|
|
|
with torch.no_grad():
|
|
vocab = shift_logits.size(-1)
|
|
per_token_loss = F.cross_entropy(
|
|
shift_logits.view(-1, vocab),
|
|
ce_labels.view(-1),
|
|
reduction="none",
|
|
ignore_index=-100,
|
|
).view_as(ce_labels)
|
|
|
|
valid_mask = ce_labels.ne(-100)
|
|
num_valid = int(valid_mask.sum().item())
|
|
avg_loss = float((per_token_loss[valid_mask].mean().item())) if num_valid > 0 else float("nan")
|
|
|
|
per_ex_avgs = []
|
|
max_examples = max(1, int(getattr(training_args, "debug_ce_max_examples", 1) or 1))
|
|
B = ce_labels.size(0)
|
|
for b in range(min(B, max_examples)):
|
|
vb = valid_mask[b]
|
|
if int(vb.sum().item()) > 0:
|
|
per_ex_avgs.append(float(per_token_loss[b][vb].mean().item()))
|
|
else:
|
|
per_ex_avgs.append(float("nan"))
|
|
logger.info(f"CE debug: tokens_in_loss={num_valid}, avg_loss={avg_loss:.4f}, per_example_avgs={[round(x,4) if x==x else None for x in per_ex_avgs]}")
|
|
except Exception as e:
|
|
logger.warning(f"CE detailed debug failed: {e}")
|
|
|
|
# --------- CRITICAL SAVE OVERRIDES: also dump FULL head/connectors for inference ---------
|
|
|
|
|
|
def _save(self, output_dir: Optional[str] = None, state_dict=None) -> None:
|
|
try:
|
|
target_dir = output_dir or self.args.output_dir
|
|
lora_out = os.path.join(target_dir, "lora")
|
|
os.makedirs(lora_out, exist_ok=True)
|
|
|
|
# --- LLM PEFT adapters (if LoRA-wrapped) ---
|
|
language_model = getattr(self.model.model, "language_model", None)
|
|
if hasattr(language_model, "save_pretrained"):
|
|
language_model.save_pretrained(lora_out)
|
|
|
|
# --- Diffusion head PEFT adapters (if LoRA-wrapped) ---
|
|
pred_head = getattr(self.model.model, "prediction_head", None)
|
|
if hasattr(pred_head, "save_pretrained"):
|
|
ph_dir = os.path.join(lora_out, "diffusion_head")
|
|
os.makedirs(ph_dir, exist_ok=True)
|
|
pred_head.save_pretrained(ph_dir)
|
|
|
|
# --- ALWAYS save FULL diffusion head state_dict for fallback ---
|
|
if pred_head is not None and hasattr(pred_head, "state_dict"):
|
|
sd = pred_head.state_dict()
|
|
torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
|
|
ph_dir = os.path.join(lora_out, "diffusion_head")
|
|
os.makedirs(ph_dir, exist_ok=True)
|
|
torch.save(sd, os.path.join(ph_dir, "diffusion_head_full.bin"))
|
|
|
|
# --- Connectors (plain state_dicts) ---
|
|
ac = getattr(self.model.model, "acoustic_connector", None)
|
|
if ac is not None:
|
|
ac_dir = os.path.join(lora_out, "acoustic_connector")
|
|
os.makedirs(ac_dir, exist_ok=True)
|
|
torch.save(ac.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
|
|
|
|
se = getattr(self.model.model, "semantic_connector", None)
|
|
if se is not None:
|
|
se_dir = os.path.join(lora_out, "semantic_connector")
|
|
os.makedirs(se_dir, exist_ok=True)
|
|
torch.save(se.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to save LoRA assets: {e}")
|
|
|
|
|
|
# ------------- Build the Trainer -------------
|
|
|
|
# Resolve which adapters to apply in samples
|
|
|
|
ema_cb = EmaCallback(attr_path="model.prediction_head", decay=0.999, device="cpu")
|
|
|
|
trainer = VibeVoiceTrainer(
|
|
model=model,
|
|
args=training_args,
|
|
train_dataset=train_dataset,
|
|
eval_dataset=eval_dataset,
|
|
data_collator=data_collator,
|
|
callbacks=[ema_cb, LoRADebugCallback(log_every_n_steps=(int(getattr(training_args, "logging_steps", 50) or 50)))],
|
|
)
|
|
|
|
# Optional debug pre-training save
|
|
if getattr(training_args, "debug_save", False):
|
|
try:
|
|
debug_dir = os.path.join(training_args.output_dir, "debug_initial")
|
|
lora_out = os.path.join(debug_dir, "lora")
|
|
os.makedirs(lora_out, exist_ok=True)
|
|
logger.info(f"[debug_save] Saving initial (pre-training) model components to: {debug_dir}")
|
|
# language model adapters / base
|
|
try:
|
|
if hasattr(model.model.language_model, "save_pretrained"):
|
|
model.model.language_model.save_pretrained(lora_out)
|
|
except Exception as e_lm:
|
|
logger.warning(f"[debug_save] Failed to save language_model: {e_lm}")
|
|
# diffusion head
|
|
try:
|
|
if hasattr(model.model, "prediction_head") and hasattr(model.model.prediction_head, "save_pretrained"):
|
|
model.model.prediction_head.save_pretrained(os.path.join(lora_out, "diffusion_head"))
|
|
except Exception as e_head:
|
|
logger.warning(f"[debug_save] Failed to save prediction_head: {e_head}")
|
|
# NEW: full diffusion head state_dict as fallback
|
|
try:
|
|
ph = getattr(model.model, "prediction_head", None)
|
|
if ph is not None and hasattr(ph, "state_dict"):
|
|
sd = ph.state_dict()
|
|
torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
|
|
os.makedirs(os.path.join(lora_out, "diffusion_head"), exist_ok=True)
|
|
torch.save(sd, os.path.join(lora_out, "diffusion_head", "diffusion_head_full.bin"))
|
|
except Exception as e:
|
|
logger.warning(f"[debug_save] Failed to save FULL diffusion head: {e}")
|
|
# connectors
|
|
try:
|
|
ac_conn = getattr(model.model, "acoustic_connector", None)
|
|
if ac_conn is not None:
|
|
ac_dir = os.path.join(lora_out, "acoustic_connector")
|
|
os.makedirs(ac_dir, exist_ok=True)
|
|
torch.save(ac_conn.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
|
|
except Exception as e_ac:
|
|
logger.warning(f"[debug_save] Failed to save acoustic_connector: {e_ac}")
|
|
try:
|
|
se_conn = getattr(model.model, "semantic_connector", None)
|
|
if se_conn is not None:
|
|
se_dir = os.path.join(lora_out, "semantic_connector")
|
|
os.makedirs(se_dir, exist_ok=True)
|
|
torch.save(se_conn.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
|
|
except Exception as e_se:
|
|
logger.warning(f"[debug_save] Failed to save semantic_connector: {e_se}")
|
|
except Exception as e:
|
|
logger.warning(f"[debug_save] Unexpected failure saving initial components: {e}")
|
|
|
|
if getattr(training_args, "gradient_checkpointing", False):
|
|
try:
|
|
model.gradient_checkpointing_enable()
|
|
except Exception:
|
|
logger.warning("Failed to enable gradient checkpointing on the model.")
|
|
|
|
if training_args.do_train:
|
|
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
|
|
|
lora_out = os.path.join(training_args.output_dir, "lora")
|
|
os.makedirs(lora_out, exist_ok=True)
|
|
|
|
# LLM PEFT (if any)
|
|
lm = getattr(model.model, "language_model", None)
|
|
if hasattr(lm, "save_pretrained"):
|
|
lm.save_pretrained(lora_out)
|
|
|
|
# Diffusion head PEFT (if any)
|
|
ph = getattr(model.model, "prediction_head", None)
|
|
if hasattr(ph, "save_pretrained"):
|
|
ph_dir = os.path.join(lora_out, "diffusion_head")
|
|
os.makedirs(ph_dir, exist_ok=True)
|
|
ph.save_pretrained(ph_dir)
|
|
|
|
# ALWAYS: full diffusion head state_dict fallback
|
|
try:
|
|
if ph is not None and hasattr(ph, "state_dict"):
|
|
sd = ph.state_dict()
|
|
torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
|
|
ph_dir = os.path.join(lora_out, "diffusion_head")
|
|
os.makedirs(ph_dir, exist_ok=True)
|
|
torch.save(sd, os.path.join(ph_dir, "diffusion_head_full.bin"))
|
|
except Exception as e:
|
|
logger.warning(f"Failed to save FULL diffusion head at end: {e}")
|
|
|
|
# Connectors (if trained)
|
|
try:
|
|
ac = getattr(model.model, "acoustic_connector", None)
|
|
if ac is not None:
|
|
ac_dir = os.path.join(lora_out, "acoustic_connector")
|
|
os.makedirs(ac_dir, exist_ok=True)
|
|
torch.save(ac.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
|
|
except Exception as e:
|
|
logger.warning(f"Failed to save acoustic_connector: {e}")
|
|
|
|
try:
|
|
se = getattr(model.model, "semantic_connector", None)
|
|
if se is not None:
|
|
se_dir = os.path.join(lora_out, "semantic_connector")
|
|
os.makedirs(se_dir, exist_ok=True)
|
|
torch.save(se.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
|
|
except Exception as e:
|
|
logger.warning(f"Failed to save semantic_connector: {e}")
|
|
|
|
if training_args.do_eval and eval_dataset is not None:
|
|
trainer.evaluate()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |