#!/usr/bin/env python3 """LoRA finetuning for MOSS-TTS Norwegian data.""" import argparse import math import os import time from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import torch import wandb from accelerate import Accelerator from peft import LoraConfig, PeftModel, get_peft_model from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR from torch.utils.data import DataLoader, Subset from transformers import AutoModel, AutoTokenizer from dataset import MossTTSDataset, collate_fn MODEL_ID = "OpenMOSS-Team/MOSS-TTS" MODEL_REVISION = "0c8df9988ab61071cdb06fe40b2bdc3132ac3b7e" LORA_TARGET_MODULES = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ] LM_HEAD_COUNT = 33 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train MOSS-TTS with LoRA adapters") parser.add_argument( "--manifest-train", default="/root/moss-tts-norwegian/moss_tts_train.jsonl", ) parser.add_argument( "--manifest-val", default="/root/moss-tts-norwegian/moss_tts_val.jsonl", ) parser.add_argument( "--tokenized-dir", default="/root/moss-tts-norwegian/tokenized", ) parser.add_argument( "--output-dir", default="/root/moss-tts-norwegian/checkpoints", ) parser.add_argument("--gpu", type=int, default=0) parser.add_argument("--lr", type=float, default=1e-5) parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--grad-accum", type=int, default=16) parser.add_argument("--max-seq-len", type=int, default=2048) parser.add_argument("--save-steps", type=int, default=500) parser.add_argument("--eval-steps", type=int, default=250) parser.add_argument("--log-steps", type=int, default=10) parser.add_argument("--warmup-steps", type=int, default=100) parser.add_argument("--weight-decay", type=float, default=0.01) parser.add_argument("--max-grad-norm", type=float, default=1.0) parser.add_argument("--num-workers", type=int, default=2) parser.add_argument("--smoke-test", action="store_true") parser.add_argument("--wandb-project", default="moss-tts-norwegian") parser.add_argument("--wandb-name", default="lora-r8-no-v1") parser.add_argument("--max-train-steps", type=int, default=0) parser.add_argument("--lora-r", type=int, default=8) parser.add_argument("--lora-alpha", type=int, default=16) parser.add_argument("--lora-dropout", type=float, default=0.05) parser.add_argument("--resume-adapter-path", default="") parser.add_argument( "--trainable-lora-modules", choices=("all", "mlp", "mlp_plus_o"), default="all", ) parser.add_argument( "--lm-heads-mode", choices=("none", "audio", "all"), default="none", ) return parser.parse_args() def load_model_with_attention_fallback( accelerator: Accelerator, ) -> Tuple[torch.nn.Module, str]: common_kwargs = { "trust_remote_code": True, "revision": MODEL_REVISION, "dtype": torch.bfloat16, } # Try flash_attention_2 first, then sdpa, then default for attn_impl in ("flash_attention_2", "sdpa", None): try: kwargs = dict(common_kwargs) if attn_impl is not None: kwargs["attn_implementation"] = attn_impl model = AutoModel.from_pretrained(MODEL_ID, **kwargs) return model, attn_impl or "default" except Exception as exc: if accelerator.is_main_process: label = attn_impl or "default" print(f"{label} loading failed ({exc}), trying next...") raise RuntimeError("Failed to load model with any attention implementation") def load_tokenizer_if_available(accelerator: Accelerator) -> Optional[AutoTokenizer]: try: return AutoTokenizer.from_pretrained( MODEL_ID, trust_remote_code=True, revision=MODEL_REVISION, ) except Exception as exc: if accelerator.is_main_process: print(f"Tokenizer load failed (not required for this training loop): {exc}") return None def build_lora_target_modules(lm_heads_mode: str) -> List[str]: target_modules = list(LORA_TARGET_MODULES) if lm_heads_mode == "audio": target_modules.extend([f"lm_heads.{idx}" for idx in range(1, LM_HEAD_COUNT)]) elif lm_heads_mode == "all": target_modules.extend([f"lm_heads.{idx}" for idx in range(LM_HEAD_COUNT)]) return target_modules def apply_lora_to_language_backbone( model: torch.nn.Module, args: argparse.Namespace, ) -> Tuple[torch.nn.Module, Dict[str, int]]: for param in model.parameters(): param.requires_grad = False target_modules = build_lora_target_modules(args.lm_heads_mode) lora_config = LoraConfig( r=args.lora_r, lora_alpha=args.lora_alpha, target_modules=target_modules, lora_dropout=args.lora_dropout, bias="none", task_type="CAUSAL_LM", ) # Monkey-patch get_input_embeddings for PEFT compatibility. # PEFT calls model.get_input_embeddings() with no args during setup to find # the nn.Embedding layer. But MossTTSDelayModel.get_input_embeddings(input_ids) # requires input_ids. We temporarily patch it, let PEFT do its setup, then restore. original_get_input_embeddings = type(model).get_input_embeddings type(model).get_input_embeddings = lambda self, input_ids=None: ( original_get_input_embeddings(self, input_ids) if input_ids is not None else self.language_model.get_input_embeddings() ) if args.resume_adapter_path: model = PeftModel.from_pretrained( model, args.resume_adapter_path, is_trainable=True, ) else: model = get_peft_model(model, lora_config) # Monkey-patch forward to prevent output_hidden_states duplication. # PEFT passes output_hidden_states via kwargs, but MossTTSDelayModel.forward() # doesn't have it as a named parameter, so it lands in **kwargs. The model then # passes output_hidden_states=True explicitly to self.language_model(**kwargs), # causing 'got multiple values for keyword argument'. Fix: wrap forward to pop it. _original_forward = type( model.get_base_model() if hasattr(model, "get_base_model") else model ).forward def _patched_forward( self, *args, output_hidden_states=None, return_dict=None, **kwargs ): return _original_forward(self, *args, **kwargs) base_cls = type( model.get_base_model() if hasattr(model, "get_base_model") else model ) base_cls.forward = _patched_forward allowed_fragments = ["language_model.layers."] if args.lm_heads_mode in {"audio", "all"}: allowed_fragments.append("lm_heads.") module_substrings = { "all": tuple(LORA_TARGET_MODULES), "mlp": ("gate_proj", "up_proj", "down_proj"), "mlp_plus_o": ("gate_proj", "up_proj", "down_proj", "o_proj"), } allowed_lora_modules = module_substrings[args.trainable_lora_modules] for name, param in model.named_parameters(): is_lora_param = "lora_" in name in_allowed_scope = any(fragment in name for fragment in allowed_fragments) in_allowed_lora_module = any(key in name for key in allowed_lora_modules) param.requires_grad = ( is_lora_param and in_allowed_scope and in_allowed_lora_module ) if args.lm_heads_mode == "audio" and "lm_heads.0." in name: param.requires_grad = False for name, param in model.named_parameters(): if "emb_ext" in name: param.requires_grad = False trainable = { name: param.numel() for name, param in model.named_parameters() if param.requires_grad } if not trainable: raise RuntimeError( "No trainable LoRA parameters found. " "Confirm that target modules exist in language_model.model.layers.*" ) head_trainable = [name for name in trainable if "lm_heads." in name] if args.lm_heads_mode == "none" and head_trainable: raise RuntimeError( f"Unexpected lm_heads LoRA params found: {head_trainable[:3]}" ) if args.lm_heads_mode in {"audio", "all"} and not head_trainable: raise RuntimeError( "lm_heads-mode enabled but no lm_heads LoRA params are trainable" ) if args.lm_heads_mode == "audio" and any( "lm_heads.0." in name for name in head_trainable ): raise RuntimeError("lm_heads.0 must remain frozen in --lm-heads-mode audio") if not any(any(key in name for key in allowed_lora_modules) for name in trainable): raise RuntimeError( "No trainable LoRA parameters matched --trainable-lora-modules=" f"{args.trainable_lora_modules}" ) return model, trainable def enable_gradient_checkpointing(model: torch.nn.Module) -> None: base_model = model.get_base_model() if hasattr(model, "get_base_model") else model language_model = getattr(base_model, "language_model", None) if hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads() elif hasattr(base_model, "enable_input_require_grads"): base_model.enable_input_require_grads() if language_model is not None and hasattr( language_model, "gradient_checkpointing_enable" ): language_model.gradient_checkpointing_enable() elif hasattr(base_model, "gradient_checkpointing_enable"): base_model.gradient_checkpointing_enable() for cfg_obj in ( getattr(model, "config", None), getattr(base_model, "config", None), getattr(language_model, "config", None) if language_model is not None else None, ): if cfg_obj is not None and hasattr(cfg_obj, "use_cache"): cfg_obj.use_cache = False def build_lr_scheduler( optimizer: torch.optim.Optimizer, warmup_steps: int, total_steps: int, ) -> LambdaLR: warmup_steps = max(0, warmup_steps) total_steps = max(1, total_steps) def lr_lambda(current_step: int) -> float: if current_step < warmup_steps: return float(current_step) / float(max(1, warmup_steps)) if total_steps <= warmup_steps: return 1.0 progress = float(current_step - warmup_steps) / float( max(1, total_steps - warmup_steps) ) progress = min(max(progress, 0.0), 1.0) return 0.5 * (1.0 + math.cos(math.pi * progress)) return LambdaLR(optimizer, lr_lambda) def _to_tensor(value: Any, device: torch.device, dtype: torch.dtype) -> torch.Tensor: if torch.is_tensor(value): tensor = value else: tensor = torch.as_tensor(value) return tensor.to(device=device, dtype=dtype, non_blocking=True) def prepare_batch( batch: Dict[str, Any], device: torch.device, default_channelwise_weight: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: if not isinstance(batch, dict): raise TypeError(f"Expected dict batch from collate_fn, got {type(batch)}") for required_key in ("input_ids", "attention_mask", "labels"): if required_key not in batch: raise KeyError(f"Batch missing required key: {required_key}") input_ids = _to_tensor(batch["input_ids"], device=device, dtype=torch.long) attention_mask = _to_tensor( batch["attention_mask"], device=device, dtype=torch.bool ) labels = _to_tensor(batch["labels"], device=device, dtype=torch.long) if "channelwise_loss_weight" in batch: channelwise_loss_weight = _to_tensor( batch["channelwise_loss_weight"], device=device, dtype=torch.float32, ) else: channelwise_loss_weight = default_channelwise_weight if channelwise_loss_weight.ndim == 2: channelwise_loss_weight = channelwise_loss_weight[0] if channelwise_loss_weight.ndim != 1: channelwise_loss_weight = channelwise_loss_weight.reshape(-1) if channelwise_loss_weight.numel() != 33: raise ValueError( "channelwise_loss_weight must have 33 entries, " f"got shape {tuple(channelwise_loss_weight.shape)}" ) return input_ids, attention_mask, labels, channelwise_loss_weight @torch.no_grad() def evaluate( model: torch.nn.Module, val_loader: DataLoader, accelerator: Accelerator, default_channelwise_weight: torch.Tensor, ) -> float: model.eval() total_loss = 0.0 total_batches = 0 for batch in val_loader: input_ids, attention_mask, labels, channelwise_loss_weight = prepare_batch( batch, device=accelerator.device, default_channelwise_weight=default_channelwise_weight, ) with accelerator.autocast(): outputs = model( input_ids=input_ids, attention_mask=attention_mask, labels=labels, channelwise_loss_weight=channelwise_loss_weight, ) loss = outputs.loss.detach().float() gathered = accelerator.gather_for_metrics(loss.unsqueeze(0)) total_loss += gathered.mean().item() total_batches += 1 model.train() if total_batches == 0: return float("nan") return total_loss / total_batches def save_adapter_checkpoint( model: torch.nn.Module, output_dir: str, tag: str, accelerator: Accelerator, ) -> None: accelerator.wait_for_everyone() if not accelerator.is_main_process: return ckpt_dir = Path(output_dir) / tag ckpt_dir.mkdir(parents=True, exist_ok=True) unwrapped_model = accelerator.unwrap_model(model) unwrapped_model.save_pretrained(str(ckpt_dir)) def lora_grad_stats(model: torch.nn.Module) -> Tuple[float, bool]: total_norm_sq = 0.0 has_nonzero_grad = False for name, param in model.named_parameters(): if "lora_" not in name or param.grad is None: continue grad = param.grad.detach().float() grad_norm = grad.norm(2).item() total_norm_sq += grad_norm * grad_norm if not has_nonzero_grad and torch.count_nonzero(grad).item() > 0: has_nonzero_grad = True return math.sqrt(total_norm_sq), has_nonzero_grad def main() -> None: args = parse_args() os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) if torch.cuda.is_available(): torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True accelerator = Accelerator(mixed_precision="bf16") if accelerator.is_main_process: Path(args.output_dir).mkdir(parents=True, exist_ok=True) accelerator.wait_for_everyone() tokenizer = load_tokenizer_if_available(accelerator) model, attn_impl = load_model_with_attention_fallback(accelerator) model, trainable_map = apply_lora_to_language_backbone(model, args) enable_gradient_checkpointing(model) train_dataset = MossTTSDataset( manifest_path=args.manifest_train, tokenized_dir=args.tokenized_dir, max_seq_len=args.max_seq_len, ) val_dataset = MossTTSDataset( manifest_path=args.manifest_val, tokenized_dir=args.tokenized_dir, max_seq_len=args.max_seq_len, ) if args.smoke_test: train_limit = min(128, len(train_dataset)) val_limit = min(128, len(val_dataset)) train_dataset = Subset(train_dataset, range(train_limit)) val_dataset = Subset(val_dataset, range(val_limit)) args.warmup_steps = min(args.warmup_steps, 5) if len(train_dataset) == 0: raise RuntimeError("Training dataset is empty") if len(val_dataset) == 0: raise RuntimeError("Validation dataset is empty") train_loader = DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=torch.cuda.is_available(), persistent_workers=args.num_workers > 0, collate_fn=collate_fn, drop_last=True, ) val_loader = DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=torch.cuda.is_available(), persistent_workers=args.num_workers > 0, collate_fn=collate_fn, drop_last=False, ) num_update_steps_per_epoch = max(1, math.ceil(len(train_loader) / args.grad_accum)) max_train_steps = args.epochs * num_update_steps_per_epoch if args.max_train_steps > 0: max_train_steps = args.max_train_steps if args.smoke_test: max_train_steps = min(max_train_steps, 20) if max_train_steps <= 0: raise RuntimeError("max_train_steps resolved to zero") trainable_params = [param for param in model.parameters() if param.requires_grad] optimizer = AdamW( trainable_params, lr=args.lr, weight_decay=args.weight_decay, ) lr_scheduler = build_lr_scheduler( optimizer, warmup_steps=args.warmup_steps, total_steps=max_train_steps, ) model, optimizer, train_loader, val_loader, lr_scheduler = accelerator.prepare( model, optimizer, train_loader, val_loader, lr_scheduler, ) default_channelwise_weight = torch.ones( 33, dtype=torch.float32, device=accelerator.device ) wandb_run = None if accelerator.is_main_process: wandb_run = wandb.init( project=args.wandb_project, name=args.wandb_name, config={ "model_id": MODEL_ID, "revision": MODEL_REVISION, "attn_implementation": attn_impl, "learning_rate": args.lr, "weight_decay": args.weight_decay, "max_grad_norm": args.max_grad_norm, "warmup_steps": args.warmup_steps, "num_epochs": args.epochs, "batch_size": args.batch_size, "gradient_accumulation_steps": args.grad_accum, "save_steps": args.save_steps, "eval_steps": args.eval_steps, "log_steps": args.log_steps, "max_seq_len": args.max_seq_len, "smoke_test": args.smoke_test, "train_samples": len(train_dataset), "val_samples": len(val_dataset), "max_train_steps": max_train_steps, "lora_r": args.lora_r, "lora_alpha": args.lora_alpha, "lora_dropout": args.lora_dropout, "lm_heads_mode": args.lm_heads_mode, "resume_adapter_path": args.resume_adapter_path, "trainable_lora_modules": args.trainable_lora_modules, }, ) total_trainable = sum(trainable_map.values()) print(f"Using device: {accelerator.device}") print(f"Attention implementation: {attn_impl}") print(f"Tokenizer loaded: {tokenizer is not None}") print( f"Trainable LoRA tensors: {len(trainable_map)}, parameters: {total_trainable}" ) for name in list(trainable_map.keys())[:6]: print(f" trainable: {name}") if len(trainable_map) > 6: print(" ...") model.train() optimizer.zero_grad(set_to_none=True) start_time = time.time() global_step = 0 micro_step = 0 accum_loss_sum = 0.0 log_loss_sum = 0.0 log_count = 0 smoke_losses = [] saw_nonzero_lora_grad = False for epoch in range(args.epochs): if global_step >= max_train_steps: break for batch in train_loader: input_ids, attention_mask, labels, channelwise_loss_weight = prepare_batch( batch, device=accelerator.device, default_channelwise_weight=default_channelwise_weight, ) with accelerator.autocast(): outputs = model( input_ids=input_ids, attention_mask=attention_mask, labels=labels, channelwise_loss_weight=channelwise_loss_weight, ) loss = outputs.loss scaled_loss = loss / args.grad_accum accelerator.backward(scaled_loss) accum_loss_sum += loss.detach().float().item() micro_step += 1 if micro_step % args.grad_accum != 0: continue grad_norm_raw = accelerator.clip_grad_norm_( model.parameters(), args.max_grad_norm ) lora_grad_norm, has_nonzero_grad = lora_grad_stats(model) saw_nonzero_lora_grad = saw_nonzero_lora_grad or has_nonzero_grad optimizer.step() lr_scheduler.step() optimizer.zero_grad(set_to_none=True) global_step += 1 step_loss = accum_loss_sum / args.grad_accum accum_loss_sum = 0.0 log_loss_sum += step_loss log_count += 1 smoke_losses.append(step_loss) grad_norm = ( grad_norm_raw.item() if torch.is_tensor(grad_norm_raw) else float(grad_norm_raw) ) if global_step % args.log_steps == 0: mean_loss = log_loss_sum / max(1, log_count) current_lr = lr_scheduler.get_last_lr()[0] if accelerator.is_main_process: print( f"step={global_step} epoch={epoch + 1} " f"loss={mean_loss:.6f} lr={current_lr:.2e} " f"grad_norm={grad_norm:.4f} lora_grad_norm={lora_grad_norm:.4f}" ) if wandb_run is not None: wandb.log( { "train/loss": mean_loss, "train/learning_rate": current_lr, "train/grad_norm": grad_norm, "train/lora_grad_norm": lora_grad_norm, "train/epoch": epoch + 1, }, step=global_step, ) log_loss_sum = 0.0 log_count = 0 if args.eval_steps > 0 and global_step % args.eval_steps == 0: val_loss = evaluate( model=model, val_loader=val_loader, accelerator=accelerator, default_channelwise_weight=default_channelwise_weight, ) if accelerator.is_main_process: print(f"step={global_step} val_loss={val_loss:.6f}") if wandb_run is not None: wandb.log({"val/loss": val_loss}, step=global_step) if args.save_steps > 0 and global_step % args.save_steps == 0: save_adapter_checkpoint( model=model, output_dir=args.output_dir, tag=f"step_{global_step}", accelerator=accelerator, ) if global_step >= max_train_steps: break save_adapter_checkpoint( model=model, output_dir=args.output_dir, tag="final", accelerator=accelerator, ) elapsed = time.time() - start_time if args.smoke_test: if not saw_nonzero_lora_grad: raise RuntimeError("Smoke test failed: LoRA gradients are all zero") if len(smoke_losses) < 2: raise RuntimeError( "Smoke test failed: fewer than 2 optimizer steps completed" ) window = min(3, len(smoke_losses)) first_loss = sum(smoke_losses[:window]) / window last_loss = sum(smoke_losses[-window:]) / window if last_loss > first_loss * 1.1: raise RuntimeError( "Smoke test failed: loss increased significantly " f"(first={first_loss:.6f}, last={last_loss:.6f}, ratio={last_loss / first_loss:.4f})" ) elif last_loss >= first_loss: if accelerator.is_main_process: print( f"Note: loss did not decrease in {len(smoke_losses)} steps " f"(first={first_loss:.6f}, last={last_loss:.6f}). " f"This is expected for short smoke tests on 8B models." ) if torch.cuda.is_available() and accelerator.is_main_process: peak_allocated = torch.cuda.max_memory_allocated() / (1024**3) peak_reserved = torch.cuda.max_memory_reserved() / (1024**3) current_allocated = torch.cuda.memory_allocated() / (1024**3) print( "Smoke test VRAM " f"current_allocated={current_allocated:.2f}GB " f"peak_allocated={peak_allocated:.2f}GB " f"peak_reserved={peak_reserved:.2f}GB" ) if accelerator.is_main_process: print(f"Training finished at step {global_step} in {elapsed:.1f}s") if wandb_run is not None and accelerator.is_main_process: wandb.finish() if __name__ == "__main__": main()