Add LoRA fine-tuning scripts (#5)
This commit is contained in:
mrfakename
@@ -0,0 +1,902 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user