"""Duo-style hidden-state distillation for AHA dynamic gates. This is the clean static-to-dynamic continuation experiment: loss = MSE(h_full, h_dynamic) on labelled tokens + reg_weight * mean(gate_soft) where h_full is produced by the same frozen checkpoint with every dynamic gate forced to 1.0, and h_dynamic uses the learned per-token gate. For ``aha_mode="duo_dynamic"``, "forced to 1.0" means "reproduce the locked static Duo full-head mask"; Duo streaming heads remain streaming. Backbone weights stay frozen; only q_proj gate rows are trainable. """ import argparse import json import os import random import shutil import sys import numpy as np import torch from torch.utils.data import DataLoader from transformers import AutoTokenizer HERE = os.path.dirname(os.path.abspath(__file__)) if HERE not in sys.path: sys.path.insert(0, HERE) from duo_train import AmDistilledDataset, LongBenchLiteAnswerDataset, collate # noqa: E402 from modeling_aha_qwen3 import ( # noqa: E402 AHA_ROUTER_GRANULARITY, AHAQwen3Config, AHAQwen3ForCausalLM, aha_router_output_size, ) from router_training_utils import configure_gate_only # noqa: E402 def _set_force_gate(model, value): prev = getattr(model.config, "aha_force_gate_value", None) model.config.aha_force_gate_value = value return prev def _restore_force_gate(model, value): model.config.aha_force_gate_value = value def _gate_tensors(out): return [g.float() for g in out.all_gate_soft] def gate_stats_from_output(out) -> dict: with torch.no_grad(): soft = torch.cat([g.reshape(-1) for g in _gate_tensors(out)]) hard = torch.cat([g.float().reshape(-1) for g in out.all_gate_hard]) return { "gate_soft_mean": soft.mean().item(), "gate_soft_std": soft.std().item(), "gate_hard_mean": hard.mean().item(), "gate_min": soft.min().item(), "gate_max": soft.max().item(), } def gate_param_stats(model) -> dict: weights = [] biases = [] num_heads = model.config.num_attention_heads head_dim = getattr(model.config, "head_dim", model.config.hidden_size // num_heads) q_rows = num_heads * head_dim with torch.no_grad(): for layer in model.model.layers: q_proj = layer.self_attn.q_proj weights.append(q_proj.weight[q_rows:].detach().float().reshape(-1).cpu()) if q_proj.bias is not None: biases.append(q_proj.bias[q_rows:].detach().float().cpu()) w = torch.cat(weights) stats = { "gate_weight_l2": w.norm().item(), "gate_weight_abs_mean": w.abs().mean().item(), } if biases: b = torch.cat(biases) alpha = torch.sigmoid(b) stats.update({ "gate_bias_mean": b.mean().item(), "gate_bias_std": b.std().item(), "gate_bias_alpha_mean": alpha.mean().item(), "gate_bias_alpha_gt05": (alpha > 0.5).float().mean().item(), }) if getattr(model.config, "aha_mode", "dynamic") == "duo_dynamic": masks = [] with torch.no_grad(): for layer in model.model.layers: alpha_static = layer.self_attn.full_attention_heads.detach().float() masks.append((alpha_static > 0.5).reshape(-1).cpu()) m = torch.cat(masks).float() stats.update({ "duo_static_full_frac": m.mean().item(), "duo_static_streaming_frac": (1.0 - m).mean().item(), }) return stats def build_reg_head_weights(model, mode: str, power: float): """Optional per-head weights for the sparsity regularizer. ``duo_alpha_margin`` uses Duo's own static alpha as confidence: full heads barely above 0.5 receive little sparsity pressure, while high-alpha full heads receive normal pressure. This tests whether the alpha-low8 inference guard can be moved into training as a smooth objective. """ if mode == "uniform": return None if mode != "duo_alpha_margin": raise ValueError(f"unknown reg head weight mode: {mode}") if getattr(model.config, "aha_mode", "dynamic") != "duo_dynamic": raise ValueError( "--reg_head_weight_mode=duo_alpha_margin requires aha_mode=duo_dynamic" ) if power <= 0.0: raise ValueError("--reg_head_weight_power must be positive") weights = [] with torch.no_grad(): for layer in model.model.layers: alpha = layer.self_attn.full_attention_heads.detach().float().clamp(0.0, 1.0) duo_full = (alpha > 0.5).float() margin = ((alpha - 0.5) / 0.5).clamp(0.0, 1.0).pow(power) weights.append((margin * duo_full).view(1, 1, -1)) flat = torch.cat([w.reshape(-1).cpu() for w in weights]) print( "[dynamic-duo] reg_head_weight_mode=duo_alpha_margin " f"power={power:g} nonzero={int((flat > 0).sum().item())}/{flat.numel()} " f"mean={flat.mean().item():.4f}", flush=True, ) return weights def sparsity_reg_from_gates(out, reg_head_weights): gate_layers = _gate_tensors(out) if reg_head_weights is None: gate_soft = torch.cat([g.reshape(-1) for g in gate_layers]) return gate_soft.mean() weighted = [] for gate, weight in zip(gate_layers, reg_head_weights): w = weight.to(device=gate.device, dtype=gate.dtype) weighted.append((gate * w).reshape(-1)) if not weighted: raise ValueError("empty weighted sparsity regularizer") return torch.cat(weighted).mean() def build_gate_optimizer(model, lr: float): setup = configure_gate_only(model) print( f"[dynamic-duo] router_granularity={model.config.aha_router_granularity} " f"gate_rows/layer={setup.gate_rows} effective trainable gate params: " f"{setup.effective_parameter_count:,}", flush=True, ) return torch.optim.AdamW([{"params": setup.parameters, "lr": lr}], weight_decay=0.0) def save_checkpoint(model, tokenizer, output_dir: str, step: int, stats: dict): sub = os.path.join(output_dir, f"checkpoint-{step}") os.makedirs(sub, exist_ok=True) prev_force = getattr(model.config, "aha_force_gate_value", None) model.config.aha_force_gate_value = None model.save_pretrained(sub, safe_serialization=True) tokenizer.save_pretrained(sub) model.config.aha_force_gate_value = prev_force with open(os.path.join(sub, "dynamic_duo_state.json"), "w") as f: json.dump( { "step": step, "router_granularity": getattr( model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY, ), "native_gate_rows_per_layer": aha_router_output_size(model.config), "effective_sparsity_denominator": "token x KV-head x layer", **stats, **gate_param_stats(model), }, f, indent=2, ) args_manifest = os.path.join(output_dir, "dynamic_duo_train_args.json") if os.path.exists(args_manifest): shutil.copyfile( args_manifest, os.path.join(sub, "dynamic_duo_train_args.json"), ) print(f"[dynamic-duo] saved {sub}", flush=True) def main(): AHAQwen3Config.register_for_auto_class() AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM") p = argparse.ArgumentParser() p.add_argument("--aha_checkpoint", required=True) p.add_argument("--model_path", default="/workspace/AHA/models/Qwen3-0.6B") p.add_argument("--output_dir", required=True) p.add_argument("--am_dataset_path", default="/workspace/Direct-Multitoken-Decoding/am-distilled-8192") p.add_argument("--am_dataset_split", default="train") p.add_argument( "--data_source", default="am_distilled", choices=["am_distilled", "longbench_lite"], help=( "Training data for dynamic gate continuation. The default keeps the " "legacy AM-distilled behavior. longbench_lite is a diagnostic " "target-distribution calibration source, not a paper main protocol." ), ) p.add_argument( "--am_label_mode", default="full", choices=["full", "answer_only"], help=( "Label mask for --data_source=am_distilled. full keeps all-token " "hidden-state distill; answer_only distills only the final answer span." ), ) p.add_argument( "--longbench_tasks", nargs="+", default=["passage_retrieval_en", "multifieldqa_en", "qasper", "2wikimqa"], help="Tasks used when --data_source=longbench_lite.", ) p.add_argument("--longbench_samples_per_task", type=int, default=30) p.add_argument("--longbench_cache_dir", default="/workspace/AHA/AHA-Qwen3/data/longbench_cache") p.add_argument("--max_length", type=int, default=8192) p.add_argument("--num_steps", type=int, default=400) p.add_argument("--warmup_ratio", type=float, default=0.2) p.add_argument("--lr", type=float, default=3e-5) p.add_argument("--reg_weight", type=float, default=0.05) p.add_argument( "--reg_head_weight_mode", default="uniform", choices=["uniform", "duo_alpha_margin"], help=( "Per-head weighting for the sparsity regularizer. uniform keeps " "the original mean(gate_soft). duo_alpha_margin downweights Duo " "static-full heads close to alpha=0.5 so fragile boundary full " "heads are not pushed local as strongly." ), ) p.add_argument( "--reg_head_weight_power", type=float, default=1.0, help=( "Power applied to the Duo alpha margin when " "--reg_head_weight_mode=duo_alpha_margin." ), ) p.add_argument("--ce_weight", type=float, default=0.0) p.add_argument( "--distill_tail_frac", type=float, default=0.0, help=( "If >0, add a tail-aware hidden-state distill term over the top " "fraction of labelled-token MSE values. This keeps the standard " "mean distill objective but prevents rare long-retrieval errors " "from being averaged away." ), ) p.add_argument( "--distill_tail_weight", type=float, default=0.0, help="Weight for the top-token MSE term enabled by --distill_tail_frac.", ) p.add_argument("--batch_size", type=int, default=1) p.add_argument("--grad_accum", type=int, default=1) p.add_argument("--save_steps", type=int, default=100) p.add_argument("--log_steps", type=int, default=10) p.add_argument("--seed", type=int, default=42) p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"]) p.add_argument("--attn_impl", default="sdpa", choices=["sdpa", "eager"]) p.add_argument("--aha_local_kind", default="sink_recent", choices=["sink_recent", "sliding_window"]) p.add_argument( "--router_granularity", default=None, choices=["token", "token_kv_head"], help="Assert that the input checkpoint has this persisted router granularity.", ) args = p.parse_args() torch.manual_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) os.makedirs(args.output_dir, exist_ok=True) with open(os.path.join(args.output_dir, "dynamic_duo_train_args.json"), "w") as f: json.dump(vars(args), f, indent=2) dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype] print(f"[dynamic-duo] checkpoint={args.aha_checkpoint}", flush=True) print( f"[dynamic-duo] loss=hidden_state_distill + {args.reg_weight} * mean(gate_soft), " f"ce_weight={args.ce_weight} tail_frac={args.distill_tail_frac} " f"tail_weight={args.distill_tail_weight}", flush=True, ) if not (0.0 <= args.distill_tail_frac <= 1.0): raise ValueError("--distill_tail_frac must be in [0, 1]") if args.distill_tail_weight < 0.0: raise ValueError("--distill_tail_weight must be non-negative") model = AHAQwen3ForCausalLM.from_pretrained_aha( args.aha_checkpoint, torch_dtype=dtype, attn_implementation=args.attn_impl, ).cuda() if getattr(model.config, "aha_mode", "dynamic") not in ("dynamic", "duo_dynamic"): raise ValueError("dynamic_duo_train.py requires a dynamic or duo_dynamic AHA checkpoint") loaded_granularity = getattr( model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY ) if args.router_granularity and args.router_granularity != loaded_granularity: raise ValueError( "--router_granularity does not match checkpoint architecture: " f"requested={args.router_granularity!r}, checkpoint={loaded_granularity!r}" ) model.config.aha_local_kind = args.aha_local_kind model.config.aha_distill_weight = 0.0 model.config.aha_ce_weight = 0.0 model.config.aha_reg_weight = -1.0 model.config.aha_force_gate_value = None print( f"[dynamic-duo] router_granularity={loaded_granularity} " f"native_gate_rows={aha_router_output_size(model.config)} " f"local_kind={model.config.aha_local_kind} " f"sink={getattr(model.config, 'duo_sink_size', None)} " f"recent={getattr(model.config, 'duo_recent_size', None)}", flush=True, ) tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token optim = build_gate_optimizer(model, args.lr) reg_head_weights = build_reg_head_weights( model, args.reg_head_weight_mode, args.reg_head_weight_power ) model.enable_input_require_grads() model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) if args.data_source == "am_distilled": dataset = AmDistilledDataset( ds_path=args.am_dataset_path, split=args.am_dataset_split, max_length=args.max_length, seed=args.seed, tokenizer=tokenizer, label_mode=args.am_label_mode, ) print( f"[dynamic-duo] data_source=am_distilled path={args.am_dataset_path} " f"split={args.am_dataset_split} n={len(dataset):,} " f"label_mode={args.am_label_mode} max_length={args.max_length}", flush=True, ) elif args.data_source == "longbench_lite": dataset = LongBenchLiteAnswerDataset( tokenizer=tokenizer, tasks=args.longbench_tasks, samples_per_task=args.longbench_samples_per_task, max_length=args.max_length, seed=args.seed, cache_dir=args.longbench_cache_dir, ) print( f"[dynamic-duo] data_source=longbench_lite tasks={args.longbench_tasks} " f"samples_per_task={args.longbench_samples_per_task} " f"rows={len(dataset.rows):,} max_length={args.max_length}", flush=True, ) else: raise ValueError(f"unknown data_source: {args.data_source}") loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate, num_workers=0) data_iter = iter(loader) warm = max(1, int(args.num_steps * args.warmup_ratio)) def lr_at(step): if step < warm: return max(0.1, (step + 1) / warm) if step > args.num_steps - warm: return max(0.1, (args.num_steps - step) / warm) return 1.0 model.train() running = {"distill": 0.0, "reg": 0.0, "ce": 0.0, "loss": 0.0, "gate_soft": 0.0, "gate_hard": 0.0} steps_in_window = 0 for step in range(args.num_steps): try: batch = next(data_iter) except StopIteration: data_iter = iter(loader) batch = next(data_iter) input_ids = batch["input_ids"].cuda() labels = batch["labels"].cuda() label_mask = labels != -100 prev_force = _set_force_gate(model, 1.0) with torch.no_grad(): out_full = model.model(input_ids=input_ids, use_cache=False) h_full = out_full.last_hidden_state _restore_force_gate(model, prev_force) out_mix = model.model(input_ids=input_ids, use_cache=False) h_mix = out_mix.last_hidden_state if label_mask.any(): diff = (h_full.float() - h_mix.float())[label_mask] else: diff = (h_full.float() - h_mix.float()).reshape(-1, h_mix.shape[-1]) token_mse = diff.pow(2).mean(dim=-1) distill_mean = token_mse.mean() if args.distill_tail_frac > 0.0 and args.distill_tail_weight > 0.0: k = max(1, int(np.ceil(token_mse.numel() * args.distill_tail_frac))) distill_tail = torch.topk(token_mse, k=k, largest=True).values.mean() distill = distill_mean + args.distill_tail_weight * distill_tail else: distill_tail = h_mix.new_zeros((), dtype=torch.float32) distill = distill_mean reg = sparsity_reg_from_gates(out_mix, reg_head_weights) if args.ce_weight > 0.0: logits = model.lm_head(h_mix).float() shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() ce = torch.nn.functional.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100, ) else: ce = h_mix.new_zeros((), dtype=torch.float32) loss = distill + args.reg_weight * reg + args.ce_weight * ce (loss / args.grad_accum).backward() if (step + 1) % args.grad_accum == 0: for group in optim.param_groups: group["lr"] = args.lr * lr_at(step) optim.step() optim.zero_grad() stats = gate_stats_from_output(out_mix) running["distill"] += float(distill.detach()) running["reg"] += float(reg.detach()) running["ce"] += float(ce.detach()) running["loss"] += float(loss.detach()) running["gate_soft"] += stats["gate_soft_mean"] running["gate_hard"] += stats["gate_hard_mean"] steps_in_window += 1 if (step + 1) % args.log_steps == 0: denom = float(steps_in_window) print( f"[step {step + 1:4d}/{args.num_steps}] " f"distill={running['distill']/denom:.6f} " f"reg={running['reg']/denom:.6f} " f"loss={running['loss']/denom:.6f} " f"gate_soft={running['gate_soft']/denom:.4f} " f"gate_hard={running['gate_hard']/denom:.4f} " f"lr={optim.param_groups[0]['lr']:.4e} " f"seq_len={input_ids.shape[1]}", flush=True, ) running = {k: 0.0 for k in running} steps_in_window = 0 if (step + 1) % args.save_steps == 0 or (step + 1) == args.num_steps: save_checkpoint(model, tokenizer, args.output_dir, step + 1, stats) print(f"[dynamic-duo] done. Final gate param stats: {gate_param_stats(model)}", flush=True) if __name__ == "__main__": main()