recipe/dynamic_duo_train.py

"""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()