import gc import re import string import sys from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR from datasets import load_dataset from tqdm import tqdm from jiwer import wer # Tee stdout/stderr to log file immediately on import. from utils import Logger, set_seed, env_flag, mode_label from config import get_config, TrainingConfig from losses import _fmt, trajectory_regularization_loss from data import prepare_dataset, DataCollatorQwenASR from peft_utils import attach_peft_adapter, save_peft_adapter_checkpoint from model import LatentQwenASR sys.stdout = Logger() sys.stderr = sys.stdout from qwen_asr import Qwen3ASRModel # type: ignore from transformers import GenerationConfig # --------------------------------------------------------------------------- # Evaluation utilities # --------------------------------------------------------------------------- @torch.no_grad() def evaluate_model( model: LatentQwenASR, processor: Any, eval_dataset: Any, num_samples: int, use_baseline: bool, ) -> Tuple[float, int]: """Compute the word error rate (WER) on a subset of evaluation data. Returns: wer_value: WER over valid samples processed: number of samples processed """ model.eval() preds: List[str] = [] refs: List[str] = [] it = iter(eval_dataset) prompt_active = bool(getattr(model, "use_soft_prompt", False)) effective_baseline = use_baseline or (not model.use_latent and not prompt_active) if effective_baseline: desc = "Baseline" elif model.use_latent: desc = f"Latent (N={model.n_latent})" else: desc = f"Prompt (N={model.n_latent})" print(f"\n--- {desc} ---") processed = 0 total = num_samples try: ds_len = len(eval_dataset) total = min(num_samples, ds_len) except Exception: total = num_samples for _ in tqdm(range(total), desc="Evaluating"): try: sample = next(it) except StopIteration: if processed == 0: print("Warning: eval dataset exhausted before any samples were read.") else: print(f"Warning: eval dataset exhausted early at {processed} samples.") break processed += 1 target_dtype = model.thinker.dtype if hasattr(model.thinker, "dtype") else torch.float32 feats = torch.tensor(sample["input_features"], dtype=target_dtype).unsqueeze(0).to(model.base_model.device) B, F, T = feats.shape if "feature_attention_mask" in sample: fam = sample["feature_attention_mask"] if not isinstance(fam, torch.Tensor): fam = torch.tensor(fam, dtype=torch.long) if fam.size(-1) < T: fam = torch.cat([fam, torch.zeros(T - fam.size(-1), dtype=torch.long)]) elif fam.size(-1) > T: fam = fam[:T] feature_attention_mask = fam.unsqueeze(0).to(feats.device) else: feature_attention_mask = torch.ones((B, T), dtype=torch.long, device=feats.device) if effective_baseline: gen_ids = model.generate( feats, feature_attention_mask=feature_attention_mask, max_new_tokens=128, use_baseline=True, return_thoughts=False, do_sample=False, eos_token_id=model.stop_ids, num_beams=1, ) thoughts = None else: if model.use_latent: gen_ids, thoughts = model.generate( feats, feature_attention_mask=feature_attention_mask, max_new_tokens=128, use_baseline=False, return_thoughts=True, do_sample=False, eos_token_id=model.stop_ids, num_beams=1, ) else: gen_ids = model.generate( feats, feature_attention_mask=feature_attention_mask, max_new_tokens=128, use_baseline=False, return_thoughts=False, do_sample=False, eos_token_id=model.stop_ids, num_beams=1, ) thoughts = None stop_ids = getattr(model, "stop_ids", [processor.tokenizer.eos_token_id]) ids = gen_ids[0] earliest_stop = ids.numel() for sid in stop_ids: if sid is None: continue matches = (ids == sid).nonzero(as_tuple=True)[0] if matches.numel() > 0: pos = int(matches[0].item()) if pos < earliest_stop: earliest_stop = pos if earliest_stop < ids.numel(): ids = ids[:earliest_stop] pred_text_raw = processor.tokenizer.decode(ids, skip_special_tokens=True) pred_text = re.sub(r"language\s+\w+", "", pred_text_raw, flags=re.IGNORECASE) if "" in pred_text: pred_text = pred_text.split("")[1] pred_text = pred_text.strip() def _normalize(t: str) -> str: t = t.lower() t = t.translate(str.maketrans("", "", string.punctuation)) return " ".join(t.split()) ref_text = sample["reference_text"] preds.append(_normalize(pred_text)) refs.append(_normalize(ref_text)) if len(preds) <= 5: print(f"\n[Sample {len(preds)}]") print(f" Ref: {ref_text.strip()}") print(f" Pred: {pred_text.strip()}") print(f" Norm Pred: {_normalize(pred_text)}") print(f" Ids: {gen_ids[0].tolist()[:20]}...") if thoughts is not None: thought_text = "" try: t_vecs = thoughts[0].float() t_norm = t_vecs / (t_vecs.norm(dim=-1, keepdim=True) + 1e-8) emb_weight = model.embed_tokens.weight.float() vocab_size = emb_weight.size(0) if vocab_size > 10000: idx = torch.randperm(vocab_size, device=emb_weight.device)[:10000] emb_sub = emb_weight[idx] emb_norm = emb_sub / (emb_sub.norm(dim=-1, keepdim=True) + 1e-8) sims = torch.matmul(t_norm, emb_norm.t()) top_vals, top_ids = sims.topk(3, dim=-1) thought_lines = [] for i in range(len(top_ids)): sub_ids = idx[top_ids[i]].tolist() toks = processor.tokenizer.convert_ids_to_tokens(sub_ids) thought_lines.append(f"T{i}:{toks}") thought_text = " | ".join(thought_lines) else: emb_norm = emb_weight / (emb_weight.norm(dim=-1, keepdim=True) + 1e-8) sims = torch.matmul(t_norm, emb_norm.t()) top_vals, top_ids = sims.topk(3, dim=-1) thought_lines = [] for i in range(len(top_ids)): toks = processor.tokenizer.convert_ids_to_tokens(top_ids[i].tolist()) thought_lines.append(f"T{i}:{toks}") thought_text = " | ".join(thought_lines) except Exception as e: thought_text = f"Error decoding: {e}" print(f" Thoughts: {thought_text}") valid = [(p, r) for p, r in zip(preds, refs) if r] if not valid: return 1.0, processed vp, vr = zip(*valid) return wer(list(vr), list(vp)), processed def _run_eval_pair( model: LatentQwenASR, processor: Any, eval_ds_clean: Any, eval_ds_other: Any, cfg: TrainingConfig, num_samples: int, label: str, primary_eval_name: str, primary_use_baseline: bool, train_mode_tag: str, ) -> Tuple[float, int, float, int, float, int, float, int]: """Run baseline + latent evaluation on both test-clean and test-other. Returns: (wer_base_clean, n_base_clean, wer_lat_clean, n_lat_clean, wer_base_other, n_base_other, wer_lat_other, n_lat_other) """ print("\n[test-clean]") wer_base_clean, n_base_clean = evaluate_model( model, processor, eval_ds_clean, num_samples=num_samples, use_baseline=primary_use_baseline, ) print( f">>> {primary_eval_name} WER ({label}, test-clean): " f"{wer_base_clean:.4f} ({wer_base_clean * 100:.2f}%) | n={n_base_clean}" ) if model.use_latent: wer_lat_clean, n_lat_clean = evaluate_model( model, processor, eval_ds_clean, num_samples=num_samples, use_baseline=False, ) print( f">>> Latent WER ({label}, test-clean): " f"{wer_lat_clean:.4f} ({wer_lat_clean * 100:.2f}%) | n={n_lat_clean}" ) else: wer_lat_clean, n_lat_clean = wer_base_clean, n_base_clean print(f">>> Latent WER ({label}, test-clean): skipped ({train_mode_tag} mode)") n_other_total = len(list(eval_ds_other)) if hasattr(eval_ds_other, '__iter__') else len(eval_ds_other) if n_other_total > 0: print("\n[test-other]") wer_base_other, n_base_other = evaluate_model( model, processor, eval_ds_other, num_samples=num_samples, use_baseline=primary_use_baseline, ) print( f">>> {primary_eval_name} WER ({label}, test-other): " f"{wer_base_other:.4f} ({wer_base_other * 100:.2f}%) | n={n_base_other}" ) if model.use_latent: wer_lat_other, n_lat_other = evaluate_model( model, processor, eval_ds_other, num_samples=num_samples, use_baseline=False, ) print( f">>> Latent WER ({label}, test-other): " f"{wer_lat_other:.4f} ({wer_lat_other * 100:.2f}%) | n={n_lat_other}" ) else: wer_lat_other, n_lat_other = wer_base_other, n_base_other print(f">>> Latent WER ({label}, test-other): skipped ({train_mode_tag} mode)") else: wer_base_other, n_base_other = 1.0, 0 wer_lat_other, n_lat_other = 1.0, 0 return ( wer_base_clean, n_base_clean, wer_lat_clean, n_lat_clean, wer_base_other, n_base_other, wer_lat_other, n_lat_other, ) # --------------------------------------------------------------------------- # Main training loop # --------------------------------------------------------------------------- def main() -> None: cfg = get_config() set_seed() print(f"CUDA Available: {torch.cuda.is_available()}") device = "cuda" if torch.cuda.is_available() else "cpu" use_bf16 = torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8 asr_wrapper = Qwen3ASRModel.from_pretrained( cfg.model_id, dtype=torch.bfloat16 if use_bf16 else torch.float16, device_map=device if device == "cuda" else None, ) asr_model = asr_wrapper.model processor = asr_wrapper.processor print(f"Using device: {device}") if cfg.use_latent_reasoning: active_n_latent = cfg.n_latent elif cfg.use_prompt_tuning: active_n_latent = cfg.prompt_tuning_num_virtual_tokens else: active_n_latent = 0 if (cfg.use_latent_reasoning or cfg.use_prompt_tuning) and active_n_latent <= 0: raise ValueError( f"{mode_label(cfg.train_mode)} requires a positive front-token count, " f"got {active_n_latent}." ) freeze_base_default = cfg.use_latent_reasoning or cfg.use_prompt_tuning if cfg.use_peft_mode: freeze_base_default = False freeze_base = env_flag("FREEZE_BASE", default=freeze_base_default) if cfg.use_peft_mode and freeze_base: print("[warn] FREEZE_BASE=1 is incompatible with PEFT adapters. Forcing FREEZE_BASE=0.") freeze_base = False freeze_audio_stack_default = cfg.use_latent_reasoning or cfg.use_prompt_tuning or cfg.use_peft_mode freeze_audio_stack = env_flag("FREEZE_AUDIO_STACK", default=freeze_audio_stack_default) print( f"Training mode: {cfg.train_mode} ({mode_label(cfg.train_mode)}) | " f"use_latent={cfg.use_latent_reasoning} | n_latent={active_n_latent} | " f"freeze_base={freeze_base} | freeze_audio_stack={freeze_audio_stack}" ) peft_metadata: Dict[str, Any] = {} bos_id = processor.tokenizer.bos_token_id eos_id = processor.tokenizer.eos_token_id pad_id = processor.tokenizer.pad_token_id print("=== Tokenizer Special Tokens ===") print(f" bos_token_id: {bos_id} ({processor.tokenizer.bos_token!r})") print(f" eos_token_id: {eos_id} ({processor.tokenizer.eos_token!r})") _im_end_check = processor.tokenizer.convert_tokens_to_ids("<|im_end|>") print(f" im_end_id (by literal): {_im_end_check}") if _im_end_check != eos_id: print(f" [WARN] im_end_id ({_im_end_check}) != eos_token_id ({eos_id}) — using literal im_end_id") print(f" pad_token_id: {pad_id}") added_tokens = list(processor.tokenizer.added_tokens_encoder.keys())[:20] print(f" First 20 added tokens: {added_tokens}") im_start_id = processor.tokenizer.convert_tokens_to_ids("<|im_start|>") if im_start_id is None or im_start_id == processor.tokenizer.unk_token_id: im_start_id = eos_id print(f" im_start_id: {im_start_id}") start_id = im_start_id if im_start_id is not None else (eos_id if eos_id is not None else 0) lang_id = start_id transcribe_id = start_id if active_n_latent > 0: print(" Adding special <|latent|> token...") special_tokens_dict = {"additional_special_tokens": ["<|latent|>"]} num_added_toks = processor.tokenizer.add_special_tokens(special_tokens_dict) if num_added_toks > 0: print(f" Resizing model embeddings to {len(processor.tokenizer)}...") asr_model.thinker.resize_token_embeddings(len(processor.tokenizer)) nt_id = processor.tokenizer.convert_tokens_to_ids("<|latent|>") print(f" nt_id: {nt_id} ({processor.tokenizer.convert_ids_to_tokens(nt_id)!r})") else: nt_id = -1 print(f" Front prompt token disabled in {mode_label(cfg.train_mode)} mode.") print(f"LANG_ID={lang_id}, TRANSCRIBE_ID={transcribe_id}, NT_ID={nt_id}") if cfg.use_peft_mode: print(f"Attaching PEFT adapter for mode={mode_label(cfg.train_mode)}...") peft_metadata = attach_peft_adapter(asr_model=asr_model, cfg=cfg) print(f"PEFT metadata: {peft_metadata}") # Load datasets thought_mode = cfg.thought_mode thought_group_size = cfg.thought_group_size print(f"Thought mode: {thought_mode} (group_size={thought_group_size})") def _make_dataset_fn(split_nt_id: int, split_n_latent: int) -> Any: """Return a prepare_dataset lambda bound to the current thought config.""" return lambda batch: prepare_dataset( processor, lang_id, transcribe_id, split_nt_id, split_n_latent, batch, thought_mode=thought_mode, thought_group_size=thought_group_size, ) print("Loading train dataset...") train_split = "train" if "extreme_asr_pony" in cfg.dataset_name else "train.100" train_ds = load_dataset(cfg.dataset_name, cfg.dataset_config, split=train_split) if cfg.train_max_samples > 0: max_train = min(int(cfg.train_max_samples), len(train_ds)) train_ds = train_ds.shuffle(seed=42).select(range(max_train)) print( f"Train dataset subsampled: {max_train} samples " f"(TRAIN_MAX_SAMPLES={cfg.train_max_samples})" ) # Dynamically remove columns that exist in the loaded dataset cols_to_remove = ["audio", "file", "id", "chapter_id", "speaker_id"] train_remove = [c for c in cols_to_remove if c in train_ds.column_names] train_ds = train_ds.map( _make_dataset_fn(nt_id, active_n_latent), remove_columns=train_remove, ) print(f"Train dataset size: {len(train_ds)} samples") try: if "extreme_asr_pony" in cfg.dataset_name: # Pony doesn't have test splits, so we just take a small validation slice from train # To avoid dropping train samples, we reload a separate train slice just for eval eval_ds_clean = load_dataset(cfg.dataset_name, cfg.dataset_config, split="train[:5%]") else: eval_ds_clean = load_dataset(cfg.dataset_name, "clean", split="test") clean_remove = [c for c in cols_to_remove if c in eval_ds_clean.column_names] eval_ds_clean = eval_ds_clean.map( _make_dataset_fn(nt_id, active_n_latent), remove_columns=clean_remove, ) except Exception as e: print(f"Warning: Could not load test-clean split: {e}") eval_ds_clean = [] try: if "extreme_asr_pony" in cfg.dataset_name: # No test-other for pony eval_ds_other = [] else: eval_ds_other = load_dataset(cfg.dataset_name, "other", split="test") other_remove = [c for c in cols_to_remove if c in eval_ds_other.column_names] eval_ds_other = eval_ds_other.map( _make_dataset_fn(nt_id, active_n_latent), remove_columns=other_remove, ) except Exception as e: print(f"Warning: Could not load test-other split: {e}") eval_ds_other = [] # Safely get lengths n_clean = len(list(eval_ds_clean)) if hasattr(eval_ds_clean, '__iter__') else len(eval_ds_clean) n_other = len(list(eval_ds_other)) if hasattr(eval_ds_other, '__iter__') else len(eval_ds_other) print(f"Eval dataset sizes: test-clean={n_clean}, test-other={n_other}") collator = DataCollatorQwenASR(processor) dl_kwargs: Dict[str, Any] = {"batch_size": cfg.batch_size, "shuffle": True, "collate_fn": collator} if device == "cuda": dl_kwargs["num_workers"] = 2 dl_kwargs["pin_memory"] = True train_loader = DataLoader(train_ds, **dl_kwargs) model = LatentQwenASR( asr_model, processor, n_latent=active_n_latent, nt_token_id=nt_id, lang_token_id=lang_id, transcribe_token_id=transcribe_id, freeze_base=freeze_base, use_latent=cfg.use_latent_reasoning, use_soft_prompt=cfg.use_prompt_tuning, soft_prompt_init_mode=cfg.prompt_tuning_init_mode, soft_prompt_init_text=cfg.prompt_tuning_init_text, user_prompt_text=cfg.user_prompt_text, delta_tanh_c=cfg.delta_tanh_c, scale_max=cfg.scale_max, scale_init=cfg.scale_init, thought_mode=cfg.thought_mode, thought_group_size=cfg.thought_group_size, halt_threshold=cfg.halt_threshold, latent_drop_prob=cfg.latent_drop_prob, latent_input_noise_std=cfg.latent_input_noise_std, latent_use_bounded_delta=cfg.latent_use_bounded_delta, latent_use_injection_gate=cfg.latent_use_injection_gate, latent_use_embedding_anchor=cfg.latent_use_embedding_anchor, freeze_audio_stack=freeze_audio_stack, ).to(device) train_mode_tag = mode_label(cfg.train_mode) # Optimizer if model.use_latent: optim_params = [ { "params": ( list(model.init_proj.parameters()) + list(model.delta_proj.parameters()) + list(model.step_proj.parameters()) + [model.step_embed] ), "lr": cfg.lr_adapter, }, {"params": [model.log_scale], "lr": cfg.lr_scale}, { "params": ( list(model.value_head.parameters()) + list(model.thought_ln.parameters()) + list(model.injection_gate.parameters()) ), "lr": cfg.lr_adapter, }, ] optimizer = torch.optim.AdamW(optim_params, weight_decay=0.01) elif model.use_soft_prompt: if not model.soft_prompt_embed.requires_grad: raise RuntimeError("Prompt-tuning active but soft_prompt_embed is frozen.") optimizer = torch.optim.AdamW( [model.soft_prompt_embed], lr=cfg.lr_prompt_tuning, weight_decay=0.01, ) print( f"Optimizer ({train_mode_tag}): AdamW lr={cfg.lr_prompt_tuning} " f"trainable_params={model.soft_prompt_embed.numel()}" ) else: trainable_params = [p for p in model.parameters() if p.requires_grad] if not trainable_params: raise RuntimeError(f"No trainable parameters found in {train_mode_tag} mode.") non_latent_lr = cfg.lr_lora_r16 if cfg.use_lora_r16 else cfg.lr_baseline_ft optimizer = torch.optim.AdamW(trainable_params, lr=non_latent_lr, weight_decay=0.01) print( f"Optimizer ({train_mode_tag}): AdamW lr={non_latent_lr} " f"trainable_params={sum(p.numel() for p in trainable_params)}" ) loss_fct = nn.CrossEntropyLoss(ignore_index=-100) grad_accum_steps = max(1, int(cfg.grad_accum_steps)) effective_batch = cfg.batch_size * grad_accum_steps print( f"Gradient accumulation: micro_batch={cfg.batch_size} × accum={grad_accum_steps} " f"= effective_batch={effective_batch}" ) batches_per_epoch = len(train_loader) optim_steps_per_epoch = max(1, batches_per_epoch // grad_accum_steps) total_optim_steps = optim_steps_per_epoch * cfg.num_epochs use_lr_schedule = cfg.use_prompt_tuning or cfg.use_lora_r16 scheduler = None if use_lr_schedule: warmup_steps = min(100, max(1, total_optim_steps // 10)) warmup_scheduler = LinearLR(optimizer, start_factor=0.1, total_iters=warmup_steps) cosine_scheduler = CosineAnnealingLR(optimizer, T_max=max(1, total_optim_steps - warmup_steps)) scheduler = SequentialLR( optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_steps], ) print( f"LR schedule: warmup={warmup_steps} optim-steps -> cosine decay " f"(total={total_optim_steps} optim-steps)" ) primary_eval_name = "Prompt" if model.use_soft_prompt else "Baseline" primary_use_baseline = not model.use_soft_prompt # Pre-training evaluation print("\n" + "=" * 50) print("Pre-Training Evaluation") print("=" * 50) print(f"Pre-training eval samples: {cfg.pretrain_eval_samples}") ( wer_base_pre_clean, n_base_pre_clean, wer_lat_pre_clean, n_lat_pre_clean, wer_base_pre_other, n_base_pre_other, wer_lat_pre_other, n_lat_pre_other, ) = _run_eval_pair( model, processor, eval_ds_clean, eval_ds_other, cfg, num_samples=cfg.pretrain_eval_samples, label="pre", primary_eval_name=primary_eval_name, primary_use_baseline=primary_use_baseline, train_mode_tag=train_mode_tag, ) print("\n" + "=" * 50) print("Starting Training...") print("=" * 50) global_step = 0 best_wer = float("inf") if model.use_latent: ckpt_prefix = cfg.checkpoint_prefix or "latent_qwen_asr" best_metric_name = "latent" elif cfg.use_prompt_tuning: ckpt_prefix = cfg.checkpoint_prefix or "prompt_tuning_qwen_asr" best_metric_name = "prompt_tuning" elif cfg.use_lora_r16: ckpt_prefix = cfg.checkpoint_prefix or "lora_r16_qwen_asr" best_metric_name = "lora_r16" else: ckpt_prefix = cfg.checkpoint_prefix or "baseline_qwen_asr" best_metric_name = "baseline" for epoch in range(1, cfg.num_epochs + 1): print(f"\n{'='*50}") print(f"Epoch {epoch}/{cfg.num_epochs}") print(f"{'='*50}") model.train() running = 0.0 epoch_loss = 0.0 num_batches = 0 accum_count = 0 optimizer.zero_grad(set_to_none=True) pbar = tqdm(train_loader, desc=f"Epoch {epoch}") for batch in pbar: target_dtype = model.thinker.dtype if hasattr(model.thinker, "dtype") else torch.float32 input_features = batch["input_features"].to(device, dtype=target_dtype) feature_attention_mask = batch["feature_attention_mask"].to(device) labels = batch["labels"].to(device) logits, stats, deltas, states, labels_for_loss, initial_state = model( input_features, labels, feature_attention_mask=feature_attention_mask, global_step=global_step, ) # Bug fix: guard against all-masked batches that would produce NaN CE loss. if (labels_for_loss != -100).sum() == 0: optimizer.zero_grad(set_to_none=True) continue eos_counts = (labels_for_loss == int(model.im_end_id)).sum(dim=1) if not torch.all(eos_counts >= 1): raise AssertionError( f"Found sample(s) without <|im_end|> supervision in batch: {eos_counts.tolist()}" ) if global_step == 0: print( "[sanity] <|im_end|> targets per sample (first batch): " f"min={int(eos_counts.min().item())} max={int(eos_counts.max().item())}" ) # Global average CE Loss for Text Decoder ce = loss_fct(logits.reshape(-1, logits.size(-1)), labels_for_loss.reshape(-1)) if model.use_latent: deltas_f = deltas.float() states_f = states.float() # ---- Value Head: Tanh Delta CE Loss (Continuous Impact) ---- # Delta_CE = cl_baseline - cl_lr # if CE_lr < CE_baseline (LR helps), delta is positive # if CE_lr > CE_baseline (LR hurts), delta is negative baseline_ce = stats.get("baseline_ce", None) # (B,) if baseline_ce is not None and "predicted_value" in stats and "lr_ce" in stats: predicted_value = stats["predicted_value"].float() predicted_value_flat = predicted_value.reshape(-1) lr_ce_per_sample = stats["lr_ce"] # (B,) lr_fixes = stats.get("lr_fixes", baseline_ce.new_zeros(baseline_ce.size(0))) lr_breaks = stats.get("lr_breaks", baseline_ce.new_zeros(baseline_ce.size(0))) baseline_acc = stats.get("baseline_acc", baseline_ce.new_zeros(baseline_ce.size(0))) lr_acc = stats.get("lr_acc", baseline_ce.new_zeros(baseline_ce.size(0))) with torch.no_grad(): batch_bl_ce = baseline_ce.mean() batch_lr_ce = lr_ce_per_sample.mean() # ---- Negative Sampling: force "LR is harmful" examples ---- # During training, LR CE is directly optimized so lr_acc >= bl_acc # and lr_ce < bl_ce almost always → target always positive → overconfidence. # Permuting lr metrics across batch doesn't help on OOD data # because ALL samples have lr_ce < bl_ce. # Fix: with probability p_neg, compute target normally then negate it: # target = -|target| (forced negative) # This tells the Value Head "this delta is wrong" ~30% of the time. p_neg = float(cfg.value_forced_neg_prob) is_neg_sample = torch.rand(1).item() < p_neg acc_diff = lr_acc - baseline_acc # (B,), range ~[-0.3, 0.3] # Fallback: when both accuracies are 0 (extreme OOD), acc_diff is # uninformative for that utterance. Use CE difference as the # per-utterance surrogate signal instead of waiting for the # whole minibatch to be degenerate. both_zero = (baseline_acc.abs() < 1e-6) & (lr_acc.abs() < 1e-6) # (B,) # Scale factor 3 (reduced from 10 to avoid saturation): # a ±0.1 accuracy diff maps to tanh(±0.3) ≈ ±0.29 # a ±0.3 accuracy diff maps to tanh(±0.9) ≈ ±0.72 acc_target = torch.tanh(acc_diff * 3.0).view(-1) ce_diff = (baseline_ce - lr_ce_per_sample).clamp(-2.0, 2.0) # (B,) ce_target = torch.tanh(ce_diff * 0.5).view(-1) target_value = torch.where(both_zero.view(-1), ce_target, acc_target) # Label smoothing: shrink toward 0 to prevent target saturation target_value = target_value * 0.9 # Force negative: flip target to -|target| for negative samples if is_neg_sample: target_value = -target_value.abs() if target_value.numel() != predicted_value_flat.numel(): if predicted_value_flat.numel() % target_value.numel() != 0: raise RuntimeError( "Value-head prediction/target size mismatch: " f"pred={predicted_value_flat.numel()} target={target_value.numel()}" ) steps_per_sample = predicted_value_flat.numel() // target_value.numel() target_value = ( target_value.unsqueeze(1) .expand(-1, steps_per_sample) .reshape(-1) ) l_value = F.mse_loss(predicted_value_flat, target_value) if global_step % 5 == 0: pos_ratio = (target_value > 0).float().mean().item() neg_ratio = (target_value < 0).float().mean().item() fixes_mean = lr_fixes.mean().item() breaks_mean = lr_breaks.mean().item() bl_ce_mean = batch_bl_ce.item() lr_ce_mean = batch_lr_ce.item() bl_acc_mean = baseline_acc.mean().item() if baseline_acc is not None else 0.0 lr_acc_mean = lr_acc.mean().item() if lr_acc is not None else 0.0 acc_diff_mean = acc_diff.mean().item() print(f"\n[Value-DEBUG] step={global_step} " f"bl_acc={bl_acc_mean:.4f} lr_acc={lr_acc_mean:.4f} " f"acc_diff={acc_diff_mean:.4f} " f"fixes={fixes_mean:.1f} breaks={breaks_mean:.1f} " f"bl_ce={bl_ce_mean:.4f} lr_ce={lr_ce_mean:.4f} " f"pos_rate={pos_ratio:.2f} neg_rate={neg_ratio:.2f} " f"neg_sample={'Y' if is_neg_sample else 'N'} " f"tgt={target_value.mean().item():.4f} " f"pred={predicted_value_flat.mean().item():.4f} l_value={l_value.item():.4f}") else: l_value = ce.new_tensor(0.0).float() # Decay w_cycle linearly over training to encourage deviation from baseline progress = global_step / max(1, float(total_optim_steps)) current_w_cycle = cfg.w_cycle * max(0.0, 1.0 - progress) # Cycle consistency loss traj = torch.cat([initial_state.float().unsqueeze(1), states_f], dim=1) l_cycle = trajectory_regularization_loss(traj, alpha=0.3) if current_w_cycle > 0 else states_f.new_tensor(0.0) loss = ( ce + current_w_cycle * l_cycle + cfg.w_value * l_value # Value Head MSE Loss ) if torch.isnan(loss): print(f"\n[FATAL] NaN Loss Detected at step {global_step}!") print(f" ce: {ce.item()} | l_cycle: {l_cycle.item()} | l_value: {l_value.item() if isinstance(l_value, torch.Tensor) else l_value}") if isinstance(l_value, torch.Tensor) and torch.isnan(l_value): print(f" --- Value Head Breakdown ---") print(f" predicted_value: {predicted_value.item():.4f} | target_value: {target_value.item():.4f}") else: loss = ce if torch.isnan(loss): print(f"\n[FATAL] NaN Baseline Loss Detected at step {global_step}!") print(f" ce: {ce.item()}") scaled_loss = loss / grad_accum_steps if scaled_loss.requires_grad: scaled_loss.backward() accum_count += 1 if accum_count % grad_accum_steps == 0: if global_step % cfg.grad_log_every == 0: with torch.no_grad(): if model.use_latent: scale_grad_norm = ( 0.0 if model.log_scale.grad is None else model.log_scale.grad.norm().item() ) delta_proj_grad_norm = ( 0.0 if model.delta_proj.weight.grad is None else model.delta_proj.weight.grad.norm().item() ) step_embed_grad_norm = ( 0.0 if model.step_embed.grad is None else model.step_embed.grad.norm().item() ) if delta_proj_grad_norm == 0.0: print(f"[WARNING] step={global_step}: delta_proj gradients are ZERO!") init_grad_norm = ( 0.0 if model.init_proj.weight.grad is None else model.init_proj.weight.grad.norm().item() ) delta_grad_norm = ( 0.0 if model.delta_proj.weight.grad is None else model.delta_proj.weight.grad.norm().item() ) scales = model.step_scales().detach().cpu() v_loss_val = l_value.item() if isinstance(l_value, torch.Tensor) else 0.0 print( f"[grad] scale_g_norm={scale_grad_norm:.6f} | " f"init_g_norm={init_grad_norm:.4f} | delta_g_norm={delta_grad_norm:.4f} | " f"step_embed_g={step_embed_grad_norm:.4f} | " f"ce={ce.item():.4f} | v_mse={v_loss_val:.4f} | v_pred={stats.get('predicted_value', torch.tensor(0.0)).mean().item():.4f} | " f"scales={_fmt(scales)}" ) else: grad_norms = [ p.grad.norm().item() for p in model.parameters() if p.requires_grad and p.grad is not None ] mean_grad = float(np.mean(grad_norms)) if grad_norms else 0.0 print(f"[grad] {train_mode_tag}_grad_norm_mean={mean_grad:.6f}") torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if scheduler is not None: scheduler.step() optimizer.zero_grad(set_to_none=True) global_step += 1 running += loss.item() epoch_loss += loss.item() num_batches += 1 if num_batches % 10 == 0: avg = running / 10.0 running = 0.0 with torch.no_grad(): if model.use_latent: scale_now = model.step_scales().mean().item() val_info = "" try: if "predicted_value" in stats: v_pred = stats["predicted_value"].mean().item() if hasattr(stats["predicted_value"], 'mean') else float(stats["predicted_value"]) val_info = f" | v_pred:{v_pred:.3f}" except: pass pbar.set_description( f"Epoch {epoch} | loss {avg:.4f} | scale {scale_now:.4f}{val_info}" ) else: pbar.set_description(f"Epoch {epoch} | loss {avg:.4f} | {train_mode_tag}") if num_batches % cfg.log_every == 0: if model.use_latent: with torch.no_grad(): rn_m = stats["raw_norm_mean"].cpu() rn_s = stats["raw_norm_std"].cpu() sn_m = stats["scaled_norm_mean"].cpu() sn_s = stats["scaled_norm_std"].cpu() cm = stats["cos_mean"].cpu() cs = stats["cos_std"].cpu() sc = stats["scales"].cpu() dn = stats["diff_norm"].item() st_c = stats["step_cos"].item() deq_i = stats.get("deq_iters", torch.tensor(0.0)).item() v_loss = l_value.item() if "l_value" in locals() else 0.0 thought_text = "" try: t_vecs = states[0] t_norm = t_vecs / (t_vecs.norm(dim=-1, keepdim=True) + 1e-8) emb_weight = model.embed_tokens.weight emb_norm = emb_weight / (emb_weight.norm(dim=-1, keepdim=True) + 1e-8) sims = torch.matmul(t_norm, emb_norm.t()) _, top_ids = sims.topk(3, dim=-1) thought_lines = [] for i in range(len(top_ids)): toks = processor.tokenizer.convert_ids_to_tokens(top_ids[i].tolist()) thought_lines.append(f"T{i}:{toks}") thought_text = " | ".join(thought_lines) except Exception as e: thought_text = f"Error decoding: {e}" print( f"[loss] ce={ce.item():.4f} | " f"traj={l_cycle.item():.4f} | v_loss={v_loss:.4f}\n" f"[thoughts] {thought_text}\n" f"[latent-metrics]\n" f" deq_iters: {deq_i:.1f}\n" f" raw_norm_mean: {_fmt(rn_m)}\n" f" raw_norm_std: {_fmt(rn_s)}\n" f" scaled_norm_mean: {_fmt(sn_m)}\n" f" scaled_norm_std: {_fmt(sn_s)}\n" f" cos_mean: {_fmt(cm)}\n" f" cos_std: {_fmt(cs)}\n" f" scales: {_fmt(sc)}\n" f" smoothness: diff_norm={dn:.4f} | step_cos={st_c:.4f}" ) else: print(f"[loss] ce={ce.item():.4f} | {train_mode_tag}") # Flush remaining accumulated gradients at end of epoch. if accum_count % grad_accum_steps != 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if scheduler is not None: scheduler.step() optimizer.zero_grad(set_to_none=True) global_step += 1 avg_epoch_loss = epoch_loss / num_batches if num_batches > 0 else 0.0 print(f"\n[Epoch {epoch}] Average Loss: {avg_epoch_loss:.4f} | Total Optim Steps: {global_step}") # Epoch evaluation print(f"\n--- Epoch {epoch} Evaluation ---") ( wer_clean_base, n_clean_base, wer_clean_lat, n_clean_lat, wer_other_base, n_other_base, wer_other_lat, n_other_lat, ) = _run_eval_pair( model, processor, eval_ds_clean, eval_ds_other, cfg, num_samples=cfg.eval_samples, label=f"epoch {epoch}", primary_eval_name=primary_eval_name, primary_use_baseline=primary_use_baseline, train_mode_tag=train_mode_tag, ) current_metric = wer_clean_lat if model.use_latent else wer_clean_base # Save checkpoint per epoch ckpt_payload: Dict[str, Any] = { "epoch": epoch, "global_step": global_step, "model_id": cfg.model_id, "train_mode": cfg.train_mode, "n_latent": active_n_latent, "freeze_base": freeze_base, "freeze_audio_stack": freeze_audio_stack, "wer_clean_baseline": wer_clean_base, "wer_other_baseline": wer_other_base, } if peft_metadata: ckpt_payload["peft"] = peft_metadata if model.use_latent: ckpt_payload.update( { "delta_tanh_c": cfg.delta_tanh_c, "value_forced_neg_prob": cfg.value_forced_neg_prob, "latent_use_bounded_delta": cfg.latent_use_bounded_delta, "latent_use_injection_gate": cfg.latent_use_injection_gate, "latent_use_embedding_anchor": cfg.latent_use_embedding_anchor, "train_max_samples": cfg.train_max_samples, "init_proj": model.init_proj.state_dict(), "delta_proj": model.delta_proj.state_dict(), "step_proj": model.step_proj.state_dict(), "step_embed": model.step_embed.detach().cpu(), "log_scale": model.log_scale.detach().cpu(), "value_head": model.value_head.state_dict(), "injection_gate": model.injection_gate.state_dict(), "wer_clean_latent": wer_clean_lat, "wer_other_latent": wer_other_lat, } ) ckpt_path = f"{ckpt_prefix}_epoch{epoch}.pth" torch.save(ckpt_payload, ckpt_path) elif model.use_soft_prompt: ckpt_payload.update( { "soft_prompt_embed": model.soft_prompt_embed.detach().cpu(), "prompt_tuning_init_mode": cfg.prompt_tuning_init_mode, "prompt_tuning_init_text": cfg.prompt_tuning_init_text, } ) ckpt_path = f"{ckpt_prefix}_epoch{epoch}.pth" torch.save(ckpt_payload, ckpt_path) elif cfg.use_peft_mode: ckpt_path = f"{ckpt_prefix}_epoch{epoch}" save_peft_adapter_checkpoint(model=model, ckpt_dir=ckpt_path, metadata=ckpt_payload) else: ckpt_payload.update({"model_state_dict": model.state_dict()}) ckpt_path = f"{ckpt_prefix}_epoch{epoch}.pth" torch.save(ckpt_payload, ckpt_path) print(f"Saved checkpoint: {ckpt_path}") if current_metric < best_wer: best_wer = current_metric if cfg.use_peft_mode and (not model.use_latent): best_ckpt_path = f"{ckpt_prefix}_best" save_peft_adapter_checkpoint(model=model, ckpt_dir=best_ckpt_path, metadata=ckpt_payload) else: best_ckpt_path = f"{ckpt_prefix}_best.pth" torch.save(ckpt_payload, best_ckpt_path) print(f"New best model! WER={best_wer:.4f} -> Saved to {best_ckpt_path}") model.train() print("\n" + "=" * 50) print("Training Complete!") print("=" * 50) print(f"Total epochs: {cfg.num_epochs}") print(f"Total steps: {global_step}") print(f"Best WER (test-clean, {best_metric_name}): {best_wer:.4f}") print(f"\nCheckpoints saved:") epoch_suffix = "" if (cfg.use_peft_mode and (not model.use_latent)) else ".pth" best_suffix = "" if (cfg.use_peft_mode and (not model.use_latent)) else ".pth" for e in range(1, cfg.num_epochs + 1): print(f" - {ckpt_prefix}_epoch{e}{epoch_suffix}") print(f" - {ckpt_prefix}_best{best_suffix} (best model)") print("\nDone.") if __name__ == "__main__": main()