recipe/duo_train.py

"""Faithful DuoAttention-style training for Qwen3 in the AHA framework.

Reproduces the paper's training objective closely:
    L = distill + reg_weight * L1(alpha) + ce_weight * CE(student_logits, labels)
    distill = MSE(h_full, h_mix) on label positions
    h_full  = forward with alpha = 1 everywhere     (pure global attention)
    h_mix   = forward with the currently learned alpha (blend global + streaming)

The CE anchor term defaults to 0 (paper-grade DuoAttention). It is required
when --unfreeze_attn_proj is set, because once the backbone is trainable the
distill objective degenerates: the teacher (alpha=1 forward) is rebuilt from
the same drifting backbone, so distill becomes self-distillation against a
moving target and admits collapse solutions where h_full ≡ h_mix but both
generate garbled tokens. CE pins the backbone to "predicting labels under
mix attention" and prevents that drift; see docs §9.4.4.

Data: synthetic multi-passkey retrieval on PaulGrahamEssays haystack,
ported from duo-attention/duo_attn/data.py::MultiplePasskeyRetrievalDataset.

Trainable params:
    - default              : 224 alpha scalars (28 layers * 8 kv_heads on Qwen3-0.6B).
    - --unfreeze_attn_proj : alphas + q/k/v/o_proj weights of every layer
                             (Setting B; sink-ablation experiment).

Use `attn_implementation="eager"` or "sdpa" — no flash-attn dependency.

Usage (paper-grade, frozen backbone):
    python duo_train.py \\
        --model_path /workspace/AHA/models/Qwen3-0.6B \\
        --haystack_dir /workspace/AHA/third_party/duo-attention/eval/needle/PaulGrahamEssays \\
        --output_dir ckpts/duo_paper_s64_r256 \\
        --max_length 8192 --context_length_min 2000 --context_length_max 8000 \\
        --num_steps 800 --lr 0.02 --reg_weight 0.05 \\
        --sink_size 64 --recent_size 256

Usage (Setting B sink ablation, unfrozen backbone + CE anchor):
    python duo_train.py \\
        --model_path /workspace/AHA/models/Qwen3-0.6B \\
        --haystack_dir /workspace/AHA/third_party/duo-attention/eval/needle/PaulGrahamEssays \\
        --output_dir ckpts/duo_sinkabl_Bv3_ce \\
        --max_length 8192 --context_length_min 2000 --context_length_max 8000 \\
        --num_steps 400 --lr 0.02 --reg_weight 0.1 \\
        --sink_size 0 --recent_size 256 \\
        --unfreeze_attn_proj --backbone_lr 1e-5 --ce_weight 1.0
"""

import argparse
import json
import math
import os
import random
import sys
from typing import List

import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from transformers import AutoTokenizer

HERE = os.path.dirname(os.path.abspath(__file__))
if HERE not in sys.path:
    sys.path.insert(0, HERE)

from modeling_aha_qwen3 import AHAQwen3ForCausalLM  # noqa: E402


LONG_BENCH_PROMPT_TEMPLATES = {
    "qasper": (
        "You are given a scientific article and a question. "
        "Answer the question as concisely as you can, using a single phrase or sentence if possible. "
        "If the question cannot be answered based on the information in the article, write \"unanswerable\". "
        "If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". "
        "Do not provide any explanation.\n\n"
        "Article: {context}\n\n"
        "Answer the question based on the above article as concisely as you can, using a single phrase or sentence if possible. "
        "If the question cannot be answered based on the information in the article, write \"unanswerable\". "
        "If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". "
        "Do not provide any explanation.\n\nQuestion: {question}\nAnswer:"
    ),
    "multifieldqa_en": (
        "Read the following text and answer briefly.\n\n"
        "{context}\n\n"
        "Now, answer the following question based on the above text, only give me the answer and do not output any other words.\n\n"
        "Question: {question}\nAnswer:"
    ),
    "2wikimqa": (
        "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\n"
        "The following are given passages.\n{context}\n\n"
        "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\n"
        "Question: {question}\nAnswer:"
    ),
    "passage_retrieval_en": (
        "Here are 30 paragraphs from Wikipedia, along with an abstract. "
        "Please determine which paragraph the abstract is from.\n\n"
        "{context}\n\n"
        "The following is an abstract.\n\n"
        "{question}\n\n"
        "Please enter the number of the paragraph that the abstract is from. "
        "The answer format must be like \"Paragraph 1\", \"Paragraph 2\", etc.\n\n"
        "The answer is: "
    ),
}


# -----------------------------------------------------------------------------
# Dataset (direct port of duo_attn/data.py::MultiplePasskeyRetrievalDataset)
# -----------------------------------------------------------------------------
PASSKEY_ALPHABET = [
    "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel",
    "india", "juliett", "kilo", "lima", "mike", "november", "oscar", "papa",
    "quebec", "romeo", "sierra", "tango", "uniform", "victor", "whiskey",
    "xray", "yankee", "zulu",
]
ORDINAL_NUMBERS = [
    "first", "second", "third", "fourth", "fifth", "sixth", "seventh",
    "eighth", "ninth", "tenth", "eleventh", "twelfth", "thirteenth",
    "fourteenth", "fifteenth", "sixteenth", "seventeenth", "eighteenth",
    "nineteenth", "twentieth",
]


def _load_haystack_text(haystack_dir: str) -> str:
    parts = []
    for fname in sorted(os.listdir(haystack_dir)):
        if not fname.endswith(".txt"):
            continue
        with open(os.path.join(haystack_dir, fname), "r", encoding="utf-8", errors="ignore") as f:
            parts.append(f.read())
    return "\n\n".join(parts)


class MultiPasskeyDataset(Dataset):
    def __init__(
        self,
        tokenizer,
        haystack_text: str,
        context_length_min: int,
        context_length_max: int,
        context_lengths_num_intervals: int,
        depth_ratio_num_intervals: int,
        min_depth_ratio: float,
        max_depth_ratio: float,
        num_passkeys: int,
        passkey_length: int,
        pad_multiple: int = 16,
        buffer_size: int = 300,
        needle: str = "Remember this sequence of words, it's the {ordinal_number} passkey to the vault: ",
        retrieval_question: str = "Based on the content of the book, what is the {ordinal_number} passkey to the vault?\nPasskey: ",
        prompt1: str = "<|im_start|> This is a very long story book: <book> ",
        prompt2: str = " </book>.\n\n",
        seperator: str = "\n\n",
    ):
        self.tokenizer = tokenizer
        self.num_passkeys = num_passkeys
        self.passkey_length = passkey_length
        self.pad_multiple = pad_multiple

        self.context_length_intervals = torch.linspace(
            context_length_min, context_length_max,
            context_lengths_num_intervals, dtype=torch.int,
        ).tolist()
        self.depth_ratio_intervals = torch.linspace(
            min_depth_ratio, max_depth_ratio, depth_ratio_num_intervals,
        ).tolist()

        self.needle_tokens_list = [
            tokenizer.encode(
                needle.format(ordinal_number=ord_), add_special_tokens=False
            ) for ord_ in ORDINAL_NUMBERS[:num_passkeys]
        ]
        self.retrieval_question_tokens_list = [
            tokenizer.encode(
                retrieval_question.format(ordinal_number=ord_), add_special_tokens=False
            ) for ord_ in ORDINAL_NUMBERS[:num_passkeys]
        ]

        self.haystack_tokens = tokenizer.encode(haystack_text, add_special_tokens=False)
        if len(self.haystack_tokens) < context_length_max:
            # tile the corpus until long enough
            repeats = context_length_max // max(1, len(self.haystack_tokens)) + 2
            self.haystack_tokens = self.haystack_tokens * repeats
        self.haystack_tokens = self.haystack_tokens[: context_length_max + 200]

        self.seperator_tokens = tokenizer.encode(seperator, add_special_tokens=False)
        self.prompt1_tokens = tokenizer.encode(prompt1, add_special_tokens=True)
        self.prompt2_tokens = tokenizer.encode(prompt2, add_special_tokens=False)
        self.buffer_size = buffer_size

    def __len__(self):
        return 10 ** 9  # effectively infinite; trainer slices by num_steps

    def _gen_passkey(self):
        seq = torch.randint(0, len(PASSKEY_ALPHABET), (self.passkey_length,))
        return " ".join(PASSKEY_ALPHABET[i] for i in seq)

    def __getitem__(self, idx):
        rng = random.Random(idx)
        context_length = int(rng.choice(self.context_length_intervals))
        depths = sorted(rng.sample(self.depth_ratio_intervals, self.num_passkeys))
        passkey_tokens_list = [
            self.tokenizer.encode(self._gen_passkey(), add_special_tokens=False)
            for _ in range(self.num_passkeys)
        ]

        haystack = self.haystack_tokens[:context_length]
        context = []
        last = 0
        for i, (d, pk) in enumerate(zip(depths, passkey_tokens_list)):
            ip = int(len(haystack) * d)
            needle = self.needle_tokens_list[i] + pk
            context += haystack[last:ip] + self.seperator_tokens + needle + self.seperator_tokens
            last = ip
        context += haystack[last:]

        qa = []
        for i, pk in enumerate(passkey_tokens_list):
            qa += self.retrieval_question_tokens_list[i] + pk + self.seperator_tokens

        ctx = self.prompt1_tokens + context + self.prompt2_tokens
        ids = ctx + qa
        # pad to multiple of 16
        pad = (-len(ids)) % self.pad_multiple
        if pad:
            ids = ids + self.haystack_tokens[-pad:]
        labels = [-100] * (len(ids) - len(qa)) + qa
        # clip pad-extension off labels
        labels = labels[: len(ids)]
        assert len(ids) == len(labels)
        return {"input_ids": torch.tensor(ids), "labels": torch.tensor(labels)}


def collate(batch):
    return {
        "input_ids": torch.stack([b["input_ids"] for b in batch]),
        "labels": torch.stack([b["labels"] for b in batch]),
    }


def _find_subsequence(haystack: List[int], needle: List[int]) -> int:
    if not needle or len(needle) > len(haystack):
        return -1
    last = len(haystack) - len(needle)
    for i in range(last + 1):
        if haystack[i:i + len(needle)] == needle:
            return i
    return -1


class AmDistilledDataset(Dataset):
    """Wraps a pre-tokenized HF dataset (e.g. /workspace/...am-distilled-8192).

    Expects each sample to have an `input_ids` field already produced by
    `tokenize-am_distill.py`. By default labels = input_ids (full-token
    distill). With label_mode="answer_only", labels before the final answer
    span are masked to -100, matching DuoAttention's answer-only distill
    pressure more closely.

    Used when `--data_source am_distilled` is set, as a drop-in replacement
    for the passkey synthetic dataset. Distill loss is then computed over
    real reasoning data instead of haystack passkey retrieval, which avoids
    the in-distribution overfitting documented in docs §9.4.7-8.
    """

    def __init__(
        self,
        ds_path: str,
        split: str,
        max_length: int,
        seed: int = 42,
        tokenizer=None,
        label_mode: str = "full",
    ):
        import datasets as hf_datasets
        loaded = hf_datasets.load_from_disk(ds_path)
        if hasattr(loaded, "keys"):
            self.ds = loaded[split]
        else:
            self.ds = loaded
        self.max_length = max_length
        self.label_mode = label_mode
        self.answer_marker_ids = []
        self.think_end_ids = []
        self.assistant_marker_ids = []
        if tokenizer is not None:
            self.answer_marker_ids = tokenizer.encode("<answer>", add_special_tokens=False)
            self.think_end_ids = tokenizer.encode("</think>", add_special_tokens=False)
            self.assistant_marker_ids = tokenizer.encode("<|im_start|>assistant", add_special_tokens=False)
        self._order = list(range(len(self.ds)))
        random.Random(seed).shuffle(self._order)

    def __len__(self):
        return len(self._order)

    def __getitem__(self, idx):
        real_idx = self._order[idx % len(self._order)]
        sample = self.ds[real_idx]
        ids = list(sample["input_ids"])[: self.max_length]
        labels = list(ids)
        if self.label_mode == "answer_only":
            labels = [-100] * len(ids)
            start = _find_subsequence(ids, self.answer_marker_ids)
            if start >= 0:
                start = start + len(self.answer_marker_ids)
            else:
                start = _find_subsequence(ids, self.think_end_ids)
                if start >= 0:
                    start = start + len(self.think_end_ids)
            if start < 0:
                # Fallback for traces without explicit <answer> inside the
                # truncation window: supervise only assistant-side tokens.
                start = _find_subsequence(ids, self.assistant_marker_ids)
                if start >= 0:
                    start = start + len(self.assistant_marker_ids)
            if 0 <= start < len(ids):
                labels[start:] = ids[start:]
        elif self.label_mode != "full":
            raise ValueError(f"unknown AM label_mode: {self.label_mode}")
        return {
            "input_ids": torch.tensor(ids, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


class LongBenchLiteAnswerDataset(Dataset):
    """Small target-distribution calibration set for alpha-only diagnostics.

    Builds LongBench-lite prompts with the same templates as eval, appends one
    gold answer, and masks labels to answer tokens only. This is intentionally
    a diagnostic data source: it answers whether a high-sparsity static Duo mask
    exists on the target distribution.
    """

    def __init__(
        self,
        tokenizer,
        tasks: List[str],
        samples_per_task: int,
        max_length: int,
        seed: int = 42,
        cache_dir: str = "/workspace/AHA/AHA-Qwen3/data/longbench_cache",
        pad_multiple: int = 16,
    ):
        from datasets import load_dataset

        self.tokenizer = tokenizer
        self.max_length = max_length
        self.pad_multiple = pad_multiple
        self.rows = []
        for task in tasks:
            template = LONG_BENCH_PROMPT_TEMPLATES[task]
            ds = load_dataset("Xnhyacinth/LongBench", task, split="test", cache_dir=cache_dir)
            n = min(samples_per_task, len(ds))
            for idx in range(n):
                sample = ds[idx]
                user_content = template.format(context=sample["context"], question=sample["question"])
                prompt = tokenizer.apply_chat_template(
                    [{"role": "user", "content": user_content}],
                    tokenize=False,
                    add_generation_prompt=True,
                    enable_thinking=False,
                )
                answers = sample["answers"] if isinstance(sample["answers"], list) else [sample["answers"]]
                answer = str(answers[0])
                prompt_ids = tokenizer(prompt, truncation=False, add_special_tokens=False)["input_ids"]
                answer_ids = tokenizer(answer, truncation=False, add_special_tokens=False)["input_ids"]
                budget = max_length - len(answer_ids) - 1
                if len(prompt_ids) > budget:
                    half = max(1, budget // 2)
                    prompt_ids = prompt_ids[:half] + prompt_ids[-(budget - half):]
                ids = prompt_ids + answer_ids
                pad = (-len(ids)) % pad_multiple
                if pad:
                    ids = ids + [tokenizer.pad_token_id] * pad
                labels = [-100] * len(prompt_ids) + answer_ids + [-100] * pad
                self.rows.append({"input_ids": ids, "labels": labels, "task": task, "idx": idx})

        random.Random(seed).shuffle(self.rows)

    def __len__(self):
        return 10 ** 9

    def __getitem__(self, idx):
        row = self.rows[idx % len(self.rows)]
        return {
            "input_ids": torch.tensor(row["input_ids"], dtype=torch.long),
            "labels": torch.tensor(row["labels"], dtype=torch.long),
        }


# -----------------------------------------------------------------------------
# Training
# -----------------------------------------------------------------------------
@torch.no_grad()
def _set_alpha_full(model, value: float = 1.0):
    """Temporarily overwrite `full_attention_heads` to a constant.

    Used to compute the 'teacher' forward pass (full attention everywhere).
    Call `_restore_alpha` with the saved tensors afterwards.
    """
    saved = []
    for layer in model.model.layers:
        p = layer.self_attn.full_attention_heads
        saved.append(p.data.clone())
        p.data.fill_(value)
    return saved


@torch.no_grad()
def _restore_alpha(model, saved):
    for layer, s in zip(model.model.layers, saved):
        layer.self_attn.full_attention_heads.data.copy_(s)


def log_alpha_stats(model) -> dict:
    with torch.no_grad():
        alphas = torch.stack([
            layer.self_attn.full_attention_heads.detach().float().clamp(0, 1)
            for layer in model.model.layers
        ], dim=0)
    m = alphas.mean().item()
    return {
        "alpha_mean": m,
        "alpha_std": alphas.std().item(),
        "alpha_gt05": (alphas > 0.5).float().mean().item(),
        "alpha_min": alphas.min().item(),
        "alpha_max": alphas.max().item(),
    }


def save_alpha_matrix(model, path: str):
    with torch.no_grad():
        alphas = torch.stack([
            layer.self_attn.full_attention_heads.detach().float().clamp(0, 1).cpu()
            for layer in model.model.layers
        ], dim=0).numpy()
    np.savetxt(path, alphas, delimiter="\t")


def _init_distributed():
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size <= 1:
        return False, 0, 0, 1, torch.device("cuda")
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    rank = int(os.environ.get("RANK", "0"))
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    return True, local_rank, rank, world_size, torch.device("cuda", local_rank)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model_path", required=True)
    p.add_argument("--aha_checkpoint_path", default="",
                   help="Optional AHA/Duo checkpoint to continue training from. "
                        "When set, model weights and alpha scalars are loaded from "
                        "this checkpoint instead of converting --model_path from base Qwen3.")
    p.add_argument("--haystack_dir", default="",
                   help="Path to PaulGraham essays for the synthetic passkey dataset. "
                        "Required when --data_source=passkey, ignored otherwise.")
    p.add_argument("--output_dir", required=True)
    p.add_argument("--max_length", type=int, default=8192)
    p.add_argument("--context_length_min", type=int, default=2000)
    p.add_argument("--context_length_max", type=int, default=8000)
    p.add_argument("--context_lengths_num_intervals", type=int, default=20)
    p.add_argument("--depth_ratio_num_intervals", type=int, default=1000)
    p.add_argument("--min_depth_ratio", type=float, default=0.05)
    p.add_argument("--max_depth_ratio", type=float, default=0.95)
    p.add_argument("--num_passkeys", type=int, default=10)
    p.add_argument("--passkey_length", type=int, default=32)
    p.add_argument("--num_steps", type=int, default=800)
    p.add_argument("--warmup_ratio", type=float, default=0.2)
    p.add_argument("--lr", type=float, default=0.02)
    p.add_argument("--reg_weight", type=float, default=0.05)
    p.add_argument("--sink_size", type=int, default=64)
    p.add_argument("--recent_size", type=int, default=256)
    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=200)
    p.add_argument("--log_steps", type=int, default=10)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--dtype", default="bfloat16")
    p.add_argument("--attn_impl", default="sdpa", choices=["sdpa", "eager"])
    # Optional: unfreeze attention projection weights (q/k/v/o_proj) together
    # with the alpha scalars — this reproduces the senior-student experiment
    # where "retraining" the model removes the need for the attention sink.
    p.add_argument("--unfreeze_attn_proj", action="store_true",
                   help="Also train q/k/v/o_proj weights alongside alpha scalars.")
    p.add_argument("--backbone_lr", type=float, default=1e-5,
                   help="Learning rate for unfrozen backbone params (alpha keeps --lr).")
    # CE anchor: required to prevent self-distill collapse when backbone is unfrozen.
    # When backbone is frozen (paper-grade DuoAttention), distill alone is well-defined
    # because the teacher (alpha=1 forward) is a fixed pretrained reference; CE is
    # redundant. When backbone is trainable, the teacher itself drifts together with
    # the student, so distill becomes self-distillation against a moving target and
    # admits degenerate solutions (h_full ≡ h_mix but both wrong → garbled output).
    # CE on the student forward pins backbone to the "predicting labels correctly"
    # manifold, blocking that failure mode.
    p.add_argument("--ce_weight", type=float, default=0.0,
                   help="Weight for cross-entropy anchor loss on labels (student/mix forward). "
                        "0 disables (paper-grade DuoAttention). Recommended >0 when "
                        "--unfreeze_attn_proj is set, to prevent backbone drift.")
    # Data source: passkey (DuoAttention legacy synthetic) or am_distilled
    # (real reasoning SFT data, see docs §9.4.7-8 for why we may want this).
    p.add_argument("--data_source", default="passkey",
                   choices=["passkey", "am_distilled", "longbench_lite"],
                   help="Training data: 'passkey' replicates DuoAttention's "
                        "synthetic haystack retrieval; 'am_distilled' uses a "
                        "pre-tokenized multi-task SFT dataset (e.g. AM-Thinking "
                        "or AM-Qwen3-Distilled). 'longbench_lite' is a diagnostic "
                        "target-distribution calibration source.")
    p.add_argument("--am_dataset_path", default="/workspace/Direct-Multitoken-Decoding/am-distilled-8192",
                   help="Path to a `datasets.load_from_disk`-compatible dataset.")
    p.add_argument("--am_dataset_split", default="train")
    p.add_argument("--am_label_mode", default="full", choices=["full", "answer_only"],
                   help="Label mask for --data_source=am_distilled. 'full' keeps the legacy "
                        "all-token hidden-state distill; 'answer_only' masks tokens before "
                        "the final <answer> span, closer to DuoAttention's QA-only objective.")
    p.add_argument("--longbench_tasks", nargs="+",
                   default=["passage_retrieval_en", "multifieldqa_en", "qasper", "2wikimqa"])
    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")
    args = p.parse_args()

    distributed, local_rank, rank, world_size, device = _init_distributed()
    is_main = rank == 0

    def log(*log_args, **log_kwargs):
        if is_main:
            print(*log_args, **log_kwargs)

    torch.manual_seed(args.seed + rank)
    random.seed(args.seed + rank)
    np.random.seed(args.seed + rank)
    os.makedirs(args.output_dir, exist_ok=True)
    if is_main:
        with open(os.path.join(args.output_dir, "duo_train_args.json"), "w") as f:
            saved_args = vars(args).copy()
            saved_args.update({"distributed": distributed, "world_size": world_size})
            json.dump(saved_args, f, indent=2)

    log(f"[duo-train] model={args.model_path}  ctx=[{args.context_length_min},{args.context_length_max}]")
    log(f"[duo-train] sink={args.sink_size} recent={args.recent_size} passkeys={args.num_passkeys}")
    log(f"[duo-train] lr={args.lr} reg_weight={args.reg_weight} num_steps={args.num_steps}")
    if distributed:
        log(f"[duo-train] distributed=torchrun world_size={world_size}")

    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    if args.data_source == "passkey":
        if not args.haystack_dir:
            raise ValueError("--haystack_dir is required when --data_source=passkey")
        haystack_text = _load_haystack_text(args.haystack_dir)
        log(f"[duo-train] haystack char length: {len(haystack_text):,}")
    else:
        haystack_text = ""

    # Build model in DUO mode. Default is initialising from base Qwen3 weights;
    # --aha_checkpoint_path is used for static-alpha continuation controls.
    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
    if args.aha_checkpoint_path:
        log(f"[duo-train] continuing from aha_checkpoint_path={args.aha_checkpoint_path}")
        model = AHAQwen3ForCausalLM.from_pretrained_aha(
            args.aha_checkpoint_path,
            torch_dtype=dtype,
            attn_implementation=args.attn_impl,
        ).to(device)
        if getattr(model.config, "aha_mode", "") != "duo":
            raise ValueError("--aha_checkpoint_path for duo_train.py must have aha_mode='duo'")
        model.config.duo_sink_size = args.sink_size
        model.config.duo_recent_size = args.recent_size
        model.config.aha_distill_weight = 0.0
        model.config.aha_ce_weight = 0.0
        model.config.aha_lambda = 0.0
        model.config.aha_gate_target = 0.0
        model.config.aha_reg_weight = -1.0
    else:
        model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
            args.model_path,
            aha_mode="duo",
            duo_sink_size=args.sink_size,
            duo_recent_size=args.recent_size,
            duo_alpha_init=1.0,
            aha_distill_weight=0.0,   # we compute distill externally
            aha_ce_weight=0.0,        # no CE in DuoAttention objective
            aha_lambda=0.0,           # we compute L1 externally
            aha_gate_target=0.0,
            torch_dtype=dtype,
            attn_implementation=args.attn_impl,
        ).to(device)

    # Freeze everything except full_attention_heads
    for param in model.parameters():
        param.requires_grad = False
    alpha_params, backbone_params = [], []
    for layer in model.model.layers:
        layer.self_attn.full_attention_heads.requires_grad = True
        alpha_params.append(layer.self_attn.full_attention_heads)
        if args.unfreeze_attn_proj:
            # Sink-ablation setting B: also retrain attention projections so the
            # model can learn attention patterns that do not rely on sink tokens.
            for proj in ("q_proj", "k_proj", "v_proj", "o_proj"):
                mod = getattr(layer.self_attn, proj, None)
                if mod is None:
                    continue
                for pname, param in mod.named_parameters():
                    param.requires_grad = True
                    backbone_params.append(param)

    # Gradient checkpointing requires *some* input to require grad. In the
    # pure-alpha setting all backbone weights are frozen so we must manually
    # enable input grads; when `--unfreeze_attn_proj` is on the projections
    # themselves already require grad so this is still harmless but optional.
    model.enable_input_require_grads()
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    core_model = model.model
    if distributed:
        core_model = DDP(core_model, device_ids=[local_rank], output_device=local_rank)
    n_alpha = sum(p.numel() for p in alpha_params)
    n_backbone = sum(p.numel() for p in backbone_params)
    log(f"[duo-train] trainable alpha scalars: {n_alpha}")
    if args.unfreeze_attn_proj:
        log(f"[duo-train] trainable backbone params (q/k/v/o_proj): {n_backbone:,}")
    else:
        log(f"[duo-train] backbone: frozen (paper-grade DuoAttention protocol)")

    if args.data_source == "passkey":
        log(f"[duo-train] data_source=passkey, haystack_dir={args.haystack_dir}")
        dataset = MultiPasskeyDataset(
            tokenizer=tokenizer,
            haystack_text=haystack_text,
            context_length_min=args.context_length_min,
            context_length_max=args.context_length_max,
            context_lengths_num_intervals=args.context_lengths_num_intervals,
            depth_ratio_num_intervals=args.depth_ratio_num_intervals,
            min_depth_ratio=args.min_depth_ratio,
            max_depth_ratio=args.max_depth_ratio,
            num_passkeys=args.num_passkeys,
            passkey_length=args.passkey_length,
        )
    elif args.data_source == "am_distilled":
        log(f"[duo-train] data_source=am_distilled, path={args.am_dataset_path} split={args.am_dataset_split}")
        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,
        )
        log(f"[duo-train] am_distilled dataset n={len(dataset):,}, max_length={args.max_length}, "
            f"label_mode={args.am_label_mode}")
    elif args.data_source == "longbench_lite":
        log(f"[duo-train] data_source=longbench_lite tasks={args.longbench_tasks} "
            f"samples_per_task={args.longbench_samples_per_task}")
        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,
        )
    else:
        raise ValueError(f"unknown data_source: {args.data_source}")
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False) if distributed else None
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        sampler=sampler,
        collate_fn=collate,
        num_workers=0,
    )
    data_iter = iter(loader)

    # Two parameter groups with independent LR multipliers. Group 0 = alpha
    # scalars (base lr = args.lr, e.g. 0.02); group 1 = backbone (base lr =
    # args.backbone_lr, e.g. 1e-5). The trapezoidal schedule multiplies both.
    param_groups = [{"params": alpha_params, "base_lr": args.lr, "lr": args.lr}]
    if backbone_params:
        param_groups.append({"params": backbone_params, "base_lr": args.backbone_lr, "lr": args.backbone_lr})
    optim = torch.optim.AdamW(param_groups, weight_decay=0.0)
    warm = max(1, int(args.num_steps * args.warmup_ratio))

    def lr_at(step):
        # trapezoidal schedule, same as DuoAttention's: ramp up over warm, hold, ramp down over warm
        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 = running_reg = running_ce = 0.0
    steps_in_window = 0
    data_epoch = 0
    for step in range(args.num_steps):
        try:
            batch = next(data_iter)
        except StopIteration:
            data_epoch += 1
            if sampler is not None:
                sampler.set_epoch(data_epoch)
            data_iter = iter(loader)
            batch = next(data_iter)
        input_ids = batch["input_ids"].to(device)
        labels = batch["labels"].to(device)
        label_mask = labels != -100

        # --- Teacher forward: force alpha = 1 everywhere (hidden_states) ----
        saved = _set_alpha_full(model, 1.0)
        old_teacher_fastpath = getattr(model.config, "_aha_teacher_full_fastpath", False)
        model.config._aha_teacher_full_fastpath = True
        try:
            with torch.no_grad():
                out_full = core_model(input_ids=input_ids, use_cache=False)
            h_full = out_full.last_hidden_state
        finally:
            model.config._aha_teacher_full_fastpath = old_teacher_fastpath
            _restore_alpha(model, saved)

        # --- Student forward: current alpha --------------------------------
        out_mix = core_model(input_ids=input_ids, use_cache=False)
        h_mix = out_mix.last_hidden_state

        # DuoAttention's exact distill: mean over hidden_dim, then mean over labelled tokens
        if label_mask.any():
            diff = (h_full.float() - h_mix.float())[label_mask]  # [N_tok, d_model]
            distill = diff.pow(2).mean(dim=-1).mean()
        else:
            distill = (h_full.float() - h_mix.float()).pow(2).mean(dim=-1).mean()

        # L1 on alpha (clamped)
        alpha_all = torch.cat([
            layer.self_attn.full_attention_heads.clamp(0.0, 1.0)
            for layer in model.model.layers
        ])
        # DuoAttention uses sum/numel == mean; kept explicit for clarity.
        reg = alpha_all.abs().sum() / alpha_all.numel()

        # CE anchor on the student (mix) forward. Only computed when ce_weight > 0
        # to keep paper-grade DuoAttention runs bit-identical to before.
        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 g in optim.param_groups:
                g["lr"] = g["base_lr"] * lr_at(step)
            optim.step()
            optim.zero_grad()
            # hard clamp alpha into [0, 1]
            with torch.no_grad():
                for layer in model.model.layers:
                    layer.self_attn.full_attention_heads.data.clamp_(0.0, 1.0)

        running_distill += float(distill.detach())
        running_reg += float(reg.detach())
        running_ce += float(ce.detach())
        steps_in_window += 1
        if (step + 1) % args.log_steps == 0:
            stats = log_alpha_stats(model)
            lr_str = f"lr_alpha={optim.param_groups[0]['lr']:.4e}"
            if len(optim.param_groups) > 1:
                lr_str += f" lr_bb={optim.param_groups[1]['lr']:.2e}"
            ce_str = f"ce={running_ce/steps_in_window:.4f} " if args.ce_weight > 0.0 else ""
            log(
                f"[step {step+1:4d}/{args.num_steps}] "
                f"distill={running_distill/steps_in_window:.4f} "
                f"reg={running_reg/steps_in_window:.4f} "
                f"{ce_str}"
                f"alpha_mean={stats['alpha_mean']:.3f} "
                f"alpha_std={stats['alpha_std']:.3f} "
                f"alpha>0.5_frac={stats['alpha_gt05']:.3f} "
                f"{lr_str} "
                f"seq_len={input_ids.shape[1]}",
                flush=True,
            )
            running_distill = running_reg = running_ce = 0.0
            steps_in_window = 0

        if (step + 1) % args.save_steps == 0 or (step + 1) == args.num_steps:
            sub = os.path.join(args.output_dir, f"checkpoint-{step+1}")
            if is_main:
                os.makedirs(sub, exist_ok=True)
                model.save_pretrained(sub, safe_serialization=True)
                tokenizer.save_pretrained(sub)
                save_alpha_matrix(model, os.path.join(sub, "full_attention_heads.tsv"))
                stats = log_alpha_stats(model)
                with open(os.path.join(sub, "duo_state.json"), "w") as f:
                    json.dump({"step": step + 1, **stats}, f, indent=2)
                log(f"[duo-train] saved {sub}")
            if distributed:
                dist.barrier()

    log(f"[duo-train] done. Final alpha stats: {log_alpha_stats(model)}")
    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()