初始化项目，由ModelHub XC社区提供模型

Model: jiamingshan/AHA-L2A-Qwen3-1.7B-repro
Source: Original Platform

.gitattributes

@@ -0,0 +1,4 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.jsonl filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+outputs/
+.cache/
+__pycache__/
+*.pyc

README.md

@@ -0,0 +1,123 @@
+---
+language:
+- en
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+base_model: Qwen/Qwen3-1.7B
+tags:
+- qwen3
+- long-context
+- sparse-attention
+- aha
+- l2a-style
+- reproducibility
+---
+
+# AHA vs L2A-style shared-gate: Qwen3-1.7B reproducibility package
+
+This repository is a self-contained package for training and evaluating two matched sparse-attention variants from the same tuned Qwen3-1.7B checkpoint. No separate source repository, starting checkpoint, or training dataset is required.
+
+## Start here
+
+An **arm** means one independent experimental variant. It does not refer to an ARM processor or a group of GPUs.
+
+| Variant | Router granularity | Selected output checkpoint |
+|---|---|---|
+| AHA | one gate per token, KV head, and layer (`token_kv_head`) | `outputs/aha/stage2/checkpoint-25` |
+| L2A-style shared-gate | one gate per token and layer, shared by every head (`token`) | `outputs/l2a_style/stage2/checkpoint-25` |
+
+The L2A-style variant is not an official L2A implementation or a claim of exact L2A reproduction. Both variants use the same AHA code, local attention, losses, data, optimizer settings, and two-stage protocol. Router granularity is the only experimental variable.
+
+The downloaded repository contains the exact tuned-vanilla starting checkpoint at its root. It does **not** contain pre-trained AHA or L2A-style output checkpoints; the training command below creates them.
+
+## Hardware and GPU behavior
+
+- Linux with Docker and NVIDIA Container Toolkit.
+- At least one NVIDIA GPU with 80 GB of memory. The protocol has been run on an A100 80 GB. H100 is expected to be compatible, but this release has not completed a separate end-to-end H100 validation.
+- At least 120 GB of free disk space is recommended for the Docker image, optimizer states, and intermediate checkpoints.
+- Internet access is needed for the initial Docker build and for public evaluation task dependencies.
+
+Each variant is a **single-GPU, world-size-1** run with global batch size 1:
+
+- With one GPU, AHA and L2A-style run sequentially on GPU 0.
+- With two or more GPUs, AHA runs on GPU 0 and L2A-style runs on GPU 1 concurrently.
+- On an eight-GPU machine, the default launcher still uses only GPUs 0 and 1; the other six GPUs remain unused.
+- The launcher does not use `torchrun`, DDP, FSDP, or DeepSpeed.
+
+Using eight-way DDP for one variant would change the declared global batch size and optimization protocol, so it would no longer reproduce this experiment.
+
+## Train both variants
+
+```bash
+python -m pip install 'huggingface-hub==0.36.2'
+
+hf download jiamingshan/AHA-L2A-Qwen3-1.7B-repro \
+  --local-dir AHA-L2A-Qwen3-1.7B-repro
+
+cd AHA-L2A-Qwen3-1.7B-repro
+bash recipe/scripts/reproduce.sh
+```
+
+The host-side entry point builds the pinned Docker environment, verifies file hashes and dataset sizes, runs ten implementation tests, and then executes hot start, Stage 1, and Stage 2 for both variants. Do not run the Python training programs directly in the host environment.
+
+The selected checkpoints are:
+
+```text
+outputs/aha/stage2/checkpoint-25
+outputs/l2a_style/stage2/checkpoint-25
+```
+
+Logs are written to `outputs/aha/logs/` and `outputs/l2a_style/logs/`. The recovery stage continues through step 75 for diagnostics, but step 25 is the predeclared selected checkpoint. The launcher is restart-safe at completed checkpoint boundaries.
+
+## Exact training protocol
+
+Both arms start from the tuned-vanilla checkpoint at repository root. Dynamic gate weights are zero and gate bias is initialized to `logit(0.9)`, so the initial hard route is full attention.
+
+| Stage | Data | Length / batch | Steps | Trainable | LR | Objective | Seed |
+|---|---|---|---:|---|---|---|---:|
+| Hot start | tuned vanilla | — | 0 | none | — | 90% full gate initialization | — |
+| 1. Gate distillation | bundled 1,024-row AM-distilled long mix | 8,192 / 1 | 300 (0.293 epoch) | native gate rows only | `3e-5` | hidden-state distillation + `0.1 × mean(gate_soft)`, CE=0 | 42 |
+| 2. Grouped-LR recovery | same rows | 8,192 / 1 | 75; select step 25 | gate + backbone, embedding/LM head frozen | gate `3e-6`, backbone `3e-7` | CE `1.0` + attention distill `0.5` + regularizer `0.01` | 47 |
+
+Both arms use `64 sink + 256 recent` local attention. The Stage-2 train threshold is `0.58`. The effective sparsity denominator is always `token × KV-head × layer`, including after broadcasting the shared token gate.
+
+Machine-readable details are in [`recipe/protocol.json`](recipe/protocol.json). Every checkpoint receives a training manifest containing the resolved granularity, native gate parameter count, data path, learning rates, threshold, seed, and source checkpoint.
+
+## Strict evaluation
+
+After training:
+
+```bash
+bash recipe/scripts/evaluate.sh
+```
+
+This evaluates aligned tuned vanilla, AHA, and L2A-style checkpoints on:
+
+- local RULER 13-config set × 20 samples;
+- BabiLong qa1–qa5 × 50;
+- HELMET-ICL 5 configs × 50;
+- MRCR 2/4/8-needle × 10.
+
+It scans inference thresholds `0.45, 0.50, 0.525, 0.55, 0.575, 0.60, 0.625, 0.65`. Prefill and decode both use sparse routing. Full-decode fallback, force-full heads, head protection, and prefill-tail fallback are disabled.
+
+The summarizer emits per-sample logs, per-config CSV, total JSON/CSV, prompt/target hash audits, and quality–sparsity plots under `outputs/summary/`. The headline point is the highest measured effective sparsity whose macro score is at least 95% of tuned vanilla on **every** suite. If no threshold passes, the output explicitly reports that no passing point was found.
+
+The local RULER 13-config set is not claimed to be the complete upstream RULER suite. Quality–sparsity measurements are not latency measurements; this package makes no speedup claim.
+
+## Integrity and environment
+
+The Docker image starts from `pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime`. Exact Python dependencies are pinned in [`recipe/requirements.txt`](recipe/requirements.txt).
+
+Before training, the launcher automatically:
+
+1. verifies SHA-256 for the starting weights, training Arrow file, frozen evaluation inputs, and core training sources;
+2. checks that the bundled dataset has exactly 1,024 rows;
+3. checks the 250 HELMET and 30 MRCR frozen rows;
+4. runs ten tiny-config tests covering router shapes, broadcast, full/local parity, regularizer equality, gate-only updates, true 10× grouped LR, save/load, embedded modeling code, one-step training, and native/effective sparsity accounting.
+
+The source implementation corresponds to Git commit [`b47a549`](https://github.com/shanjiaming/AHA/commit/b47a549) on `codex/l2a-router-granularity`. File hashes are recorded in [`recipe/manifest.json`](recipe/manifest.json).
+
+## Data and model note
+
+The starting weights are derived from `Qwen/Qwen3-1.7B`. The bundled training and evaluation snapshots are provided to freeze the exact experimental rows. Users remain responsible for complying with the terms of the respective upstream model and datasets.

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,60 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 6144,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 28,
+  "model_type": "qwen3",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.54.0",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "temperature": 0.6,
+  "top_k": 20,
+  "top_p": 0.95,
+  "transformers_version": "4.54.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22c971a65f7f1835b0e2a38b45f2f92191f5d8b63b30969f2f61a69d30afbc05
+size 3441185608

recipe/Dockerfile

@@ -0,0 +1,16 @@
+FROM pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    TOKENIZERS_PARALLELISM=false \
+    HF_HOME=/workspace/.cache/huggingface \
+    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends git ca-certificates \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt /tmp/requirements.txt
+RUN python -m pip install --no-cache-dir -r /tmp/requirements.txt
+
+WORKDIR /workspace/recipe

recipe/data/am_distilled_long_mix/build_summary.json

@@ -0,0 +1,18 @@
+{
+  "source_dataset": "/workspace/AHA/AHA-OLMO2/experiments/datasets/am-distilled-32768",
+  "output_dataset": "/workspace/AHA/AHA-Qwen3/experiments/datasets/am-distilled-32768-mix-long12k512-rand512",
+  "seed": 124,
+  "long_n": 512,
+  "rand_n": 512,
+  "min_long": 12288,
+  "rows": 1024,
+  "min_tokens": 369,
+  "p50": 12580,
+  "p90": 21189,
+  "p95": 24547,
+  "p99": 30902,
+  "max_tokens": 32618,
+  "mean_tokens": 10600.341796875,
+  "frac_ge_8192": 0.5517578125,
+  "frac_ge_12288": 0.5302734375
+}

recipe/data/am_distilled_long_mix/dataset_dict.json

@@ -0,0 +1 @@
+{"splits": ["train"]}

recipe/data/am_distilled_long_mix/train/cache-e6708cbf54082730.arrow

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0135bf6e96f7a4f452894727b9be1ce92ff27cc65330b94faeab11e993f5aa3c
+size 61801216

recipe/data/am_distilled_long_mix/train/data-00000-of-00001.arrow

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27f158ba22512269d33e8fb14db97a84e541ca9b701fba73a98e70c73687f45c
+size 87171000

recipe/data/am_distilled_long_mix/train/dataset_info.json

@@ -0,0 +1,152 @@
+{
+  "builder_name": "generator",
+  "citation": "",
+  "config_name": "default",
+  "dataset_name": "generator",
+  "dataset_size": 45895360821,
+  "description": "",
+  "download_checksums": {},
+  "download_size": 0,
+  "features": {
+    "messages": {
+      "feature": {
+        "content": {
+          "dtype": "string",
+          "_type": "Value"
+        },
+        "role": {
+          "dtype": "string",
+          "_type": "Value"
+        }
+      },
+      "_type": "List"
+    },
+    "input_ids": {
+      "feature": {
+        "dtype": "int32",
+        "_type": "Value"
+      },
+      "_type": "List"
+    },
+    "attention_mask": {
+      "feature": {
+        "dtype": "int8",
+        "_type": "Value"
+      },
+      "_type": "List"
+    },
+    "num_tokens": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": "",
+  "size_in_bytes": 45895360821,
+  "splits": {
+    "train": {
+      "name": "train",
+      "num_bytes": 45895360821,
+      "num_examples": 1890613,
+      "shard_lengths": [
+        15000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        15000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        14000,
+        25000,
+        26000,
+        27000,
+        26000,
+        26000,
+        25000,
+        26000,
+        25000,
+        26000,
+        25000,
+        26000,
+        25000,
+        26000,
+        26000,
+        25000,
+        25000,
+        25000,
+        24000,
+        25000,
+        24000,
+        24000,
+        24000,
+        24000,
+        23000,
+        19000,
+        18000,
+        18000,
+        18000,
+        18000,
+        18000,
+        18000,
+        18000,
+        19000,
+        22000,
+        23000,
+        23000,
+        23000,
+        23000,
+        23000,
+        23000,
+        23000,
+        23000,
+        23000,
+        23000,
+        22000,
+        22000,
+        23000,
+        22000,
+        22000,
+        22000,
+        23000,
+        22000,
+        22000,
+        22000,
+        22000,
+        23000,
+        23000,
+        23000,
+        23000,
+        22000,
+        22000,
+        28000,
+        30000,
+        30000,
+        29000,
+        30000,
+        20613
+      ],
+      "dataset_name": "generator"
+    }
+  },
+  "version": {
+    "version_str": "0.0.0",
+    "major": 0,
+    "minor": 0,
+    "patch": 0
+  }
+}

recipe/data/am_distilled_long_mix/train/state.json

@@ -0,0 +1,18 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "8470c8174c1f6ae9",
+  "_format_columns": [
+    "attention_mask",
+    "input_ids",
+    "messages",
+    "num_tokens"
+  ],
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

recipe/data/eval_inputs/input_manifest.json

@@ -0,0 +1,34 @@
+{
+  "tokenizer_path": "/data/sjm/AHA/models/Qwen3-1.7B",
+  "thinking": false,
+  "truncation": false,
+  "target_context_k": 8,
+  "helmet": {
+    "path": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/inputs/helmet_icl_8k_n50_per_config.jsonl",
+    "source": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fivebench_n100_20260712/helmet8k_rows_n100.jsonl",
+    "configs": [
+      "trec_coarse",
+      "trec_fine",
+      "banking77",
+      "clinic150",
+      "nlu"
+    ],
+    "samples_per_config": 50,
+    "total_samples": 250,
+    "min_prompt_tokens": 6344,
+    "max_prompt_tokens": 7413
+  },
+  "mrcr": {
+    "path": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl",
+    "source": "/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_broad_20260713/inputs/mrcr_8k_16k_2_4_8needle_n10_per_config.jsonl",
+    "configs": [
+      "8k_2needle",
+      "8k_4needle",
+      "8k_8needle"
+    ],
+    "samples_per_config": 10,
+    "total_samples": 30,
+    "min_prompt_tokens": 4312,
+    "max_prompt_tokens": 8028
+  }
+}

recipe/duo_train.py

@@ -0,0 +1,818 @@
+"""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()

recipe/dynamic_duo_train.py

@@ -0,0 +1,494 @@
+"""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()

recipe/manifest.json

@@ -0,0 +1,13 @@
+{
+  "schema": "aha-l2a-qwen3-repro-v1",
+  "source_commit": "b47a549",
+  "files": {
+    "../model.safetensors": "22c971a65f7f1835b0e2a38b45f2f92191f5d8b63b30969f2f61a69d30afbc05",
+    "data/am_distilled_long_mix/train/data-00000-of-00001.arrow": "27f158ba22512269d33e8fb14db97a84e541ca9b701fba73a98e70c73687f45c",
+    "data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl": "76e4c8115014d11e2dbf47ca8a2e300c490fd8e454326d5e589b12682b7f140e",
+    "data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl": "7074281289fa051278e213ecea523cedf3f609e89b6ae13ab575c075f386ddc4",
+    "modeling_aha_qwen3.py": "4b10c633ab62dbbacbd2a9ac0c69f8225876e97f70bd6ddc4b37367053e2fc1b",
+    "dynamic_duo_train.py": "14188f72ae418831ccda997d0e4cae71c808a59b96aabd3be1e289536a1e073e",
+    "sft.py": "d1f7ae5818639b25ed130433073f4717dcad310dc941bf4a13a9063730b4e555"
+  }
+}

recipe/protocol.json

@@ -0,0 +1,83 @@
+{
+  "release": "jiamingshan/AHA-L2A-Qwen3-1.7B-repro",
+  "source_git": {
+    "repository": "https://github.com/shanjiaming/AHA",
+    "branch": "codex/l2a-router-granularity",
+    "commit": "b47a549"
+  },
+  "base_checkpoint": {
+    "location": "repository root",
+    "description": "Qwen3-1.7B tuned vanilla, global step 6911",
+    "sha256_model_safetensors": "22c971a65f7f1835b0e2a38b45f2f92191f5d8b63b30969f2f61a69d30afbc05"
+  },
+  "execution": {
+    "distributed_training": false,
+    "world_size_per_arm": 1,
+    "gpu_count_per_arm": 1,
+    "one_gpu_machine": "run AHA then L2A-style sequentially on GPU 0",
+    "two_or_more_gpu_machine": "run AHA on GPU 0 and L2A-style on GPU 1 concurrently; remaining GPUs are unused",
+    "arm_definition": "one independent experimental variant, not a processor architecture or GPU group"
+  },
+  "arms": {
+    "aha": {
+      "router_granularity": "token_kv_head",
+      "description": "one native gate per token, KV head, and layer"
+    },
+    "l2a_style": {
+      "router_granularity": "token",
+      "description": "one native gate per token and layer, shared across all heads",
+      "claim_boundary": "L2A-style shared-gate; not an official L2A reproduction"
+    }
+  },
+  "common": {
+    "local_attention": {
+      "kind": "sink_recent",
+      "sink_tokens": 64,
+      "recent_tokens": 256
+    },
+    "gate_initialization": {
+      "weight": "zeros",
+      "full_probability": 0.9
+    },
+    "effective_sparsity_denominator": "token x KV-head x layer"
+  },
+  "stage_1": {
+    "dataset": "recipe/data/am_distilled_long_mix",
+    "rows": 1024,
+    "max_length": 8192,
+    "batch_size": 1,
+    "optimizer_steps": 300,
+    "epochs": 0.29296875,
+    "trainable": "native gate rows only",
+    "learning_rate": 0.00003,
+    "warmup_ratio": 0.1,
+    "hidden_state_distillation": 1.0,
+    "regularizer": 0.1,
+    "ce": 0.0,
+    "train_threshold": 0.5,
+    "seed": 42
+  },
+  "stage_2": {
+    "dataset": "same 1024 rows",
+    "max_length": 8192,
+    "batch_size": 1,
+    "optimizer_steps_run": 75,
+    "selected_checkpoint_step": 25,
+    "selected_epochs": 0.0244140625,
+    "gate_learning_rate": 0.000003,
+    "backbone_learning_rate": 0.0000003,
+    "ce": 1.0,
+    "attention_distillation": 0.5,
+    "regularizer": 0.01,
+    "train_threshold": 0.58,
+    "freeze_embeddings_and_lm_head": true,
+    "seed": 47
+  },
+  "evaluation": {
+    "thresholds": [0.45, 0.5, 0.525, 0.55, 0.575, 0.6, 0.625, 0.65],
+    "strict_prefill_and_decode_sparse_routing": true,
+    "force_full_heads": false,
+    "full_decode": false,
+    "headline": "highest measured effective sparsity whose score is at least 95% of tuned vanilla on every suite"
+  }
+}

recipe/requirements.txt

@@ -0,0 +1,12 @@
+transformers==4.54.0
+trl==0.19.1
+datasets==4.8.4
+accelerate==1.13.0
+safetensors==0.7.0
+huggingface-hub==0.36.2
+tokenizers==0.21.4
+pyarrow==24.0.0
+pandas==3.0.3
+numpy==2.3.4
+matplotlib==3.10.9
+lm_eval==0.4.11

recipe/router_training_utils.py

@@ -0,0 +1,86 @@
+"""Shared training helpers for AHA q_proj router rows."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from modeling_aha_qwen3 import aha_router_output_size
+
+
+@dataclass
+class GateOnlySetup:
+    parameters: list[torch.nn.Parameter]
+    effective_parameter_count: int
+    q_rows: int
+    gate_rows: int
+
+
+class RowWiseAdamW(torch.optim.AdamW):
+    """AdamW with an exact lower LR on prefixes of selected tensors."""
+
+    def __init__(self, params, *, row_scales, **kwargs):
+        super().__init__(params, **kwargs)
+        self._row_scales = row_scales
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        before = [p[:n_rows].detach().clone() for p, n_rows, _ in self._row_scales]
+        loss = super().step(closure=closure)
+        for (parameter, n_rows, scale), old in zip(self._row_scales, before):
+            if scale != 1.0:
+                new = parameter[:n_rows]
+                new.copy_(old + scale * (new - old))
+        return loss
+
+
+def q_projection_rows(config) -> int:
+    head_dim = getattr(
+        config,
+        "head_dim",
+        config.hidden_size // config.num_attention_heads,
+    )
+    return int(config.num_attention_heads * head_dim)
+
+
+def configure_gate_only(model: torch.nn.Module) -> GateOnlySetup:
+    """Freeze a model and expose only the appended q_proj gate rows.
+
+    PyTorch cannot mark only a slice of a Parameter trainable, so each q_proj
+    tensor remains trainable while a hook zeros the ordinary Q-row gradient.
+    The returned parameter count is the effective native gate parameter count,
+    not the full q_proj tensor size seen by the optimizer.
+    """
+
+    for parameter in model.parameters():
+        parameter.requires_grad = False
+
+    q_rows = q_projection_rows(model.config)
+    gate_rows = aha_router_output_size(model.config)
+
+    def mask_q_rows(gradient: torch.Tensor) -> torch.Tensor:
+        masked = gradient.clone()
+        masked[:q_rows] = 0.0
+        return masked
+
+    parameters: list[torch.nn.Parameter] = []
+    effective = 0
+    for layer in model.model.layers:
+        q_proj = layer.self_attn.q_proj
+        q_proj.weight.requires_grad = True
+        q_proj.weight.register_hook(mask_q_rows)
+        parameters.append(q_proj.weight)
+        effective += gate_rows * q_proj.in_features
+        if q_proj.bias is not None:
+            q_proj.bias.requires_grad = True
+            q_proj.bias.register_hook(mask_q_rows)
+            parameters.append(q_proj.bias)
+            effective += gate_rows
+
+    return GateOnlySetup(
+        parameters=parameters,
+        effective_parameter_count=effective,
+        q_rows=q_rows,
+        gate_rows=gate_rows,
+    )

recipe/scripts/eval_cell.sh

@@ -0,0 +1,82 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# -ne 4 ]]; then
+  echo "Usage: $0 METHOD MODEL_PATH GPU BENCHMARK" >&2
+  exit 2
+fi
+
+METHOD="$1"
+MODEL_PATH="$2"
+GPU="$3"
+BENCHMARK="$4"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+REPO="$(cd "$RECIPE/.." && pwd)"
+EVAL_ROOT="${EVAL_ROOT:-$REPO/outputs/eval}"
+OUT="$EVAL_ROOT/$METHOD"
+LOG="$EVAL_ROOT/logs/$METHOD"
+TOKENIZER="${VANILLA_DIR:-$REPO}"
+HELMET_ROWS="$RECIPE/data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl"
+MRCR_ROWS="$RECIPE/data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl"
+RULER_LIMIT="${RULER_LIMIT:-20}"
+BABILONG_LIMIT="${BABILONG_LIMIT:-50}"
+AHA_GATE_HARD_THRESHOLD="${AHA_GATE_HARD_THRESHOLD:-0.5}"
+mkdir -p "$OUT" "$LOG"
+
+if [[ -f "$OUT/${BENCHMARK}.DONE" ]]; then
+  exit 0
+fi
+
+export CUDA_VISIBLE_DEVICES="$GPU"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export AHA_SPARSITY_STATS_PATH="$OUT/sparsity_${BENCHMARK}.json"
+export AHA_GATE_HARD_THRESHOLD
+export AHA_FORCE_FULL_DECODE=0
+export AHA_FORCE_FULL_PREFILL_TAIL=0
+export AHA_FORCE_FULL_HEADS=
+export AHA_FORCE_LOW_ALPHA_FULL_HEADS=0
+export AHA_DUO_PREFILL_FULL=0
+
+model_args="pretrained=${MODEL_PATH},trust_remote_code=True,dtype=bfloat16,max_length=16384,attn_implementation=sdpa"
+split_a="niah_single_1,niah_single_3,niah_multikey_2,niah_multiquery,ruler_vt,ruler_fwe,ruler_qa_hotpot"
+split_b="niah_single_2,niah_multikey_1,niah_multikey_3,niah_multivalue,ruler_cwe,ruler_qa_squad"
+
+case "$BENCHMARK" in
+  ruler_a|ruler_b)
+    tasks="$split_a"
+    [[ "$BENCHMARK" == ruler_b ]] && tasks="$split_b"
+    python -m lm_eval \
+      --model hf --model_args "$model_args" \
+      --tasks "$tasks" --metadata '{"max_seq_lengths":[8192]}' \
+      --batch_size 1 --limit "$RULER_LIMIT" --log_samples \
+      --output_path "$OUT/ruler8k_splits/${BENCHMARK#ruler_}/lm_eval" \
+      >"$LOG/${BENCHMARK}.log" 2>&1
+    ;;
+  babilong)
+    python -m lm_eval \
+      --model hf --model_args "$model_args" \
+      --tasks babilong_longctx --metadata '{"max_seq_lengths":"8k"}' \
+      --num_fewshot 2 --batch_size 1 --limit "$BABILONG_LIMIT" --log_samples \
+      --output_path "$OUT/babilong8k_qa1_qa5_n50" \
+      >"$LOG/babilong.log" 2>&1
+    ;;
+  helmet)
+    python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
+      --benchmark helmet_icl --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
+      --attn-implementation sdpa --rows "$HELMET_ROWS" \
+      --method "$METHOD" --output "$OUT/helmet_icl8k_n50.jsonl" \
+      >"$LOG/helmet.log" 2>&1
+    ;;
+  mrcr)
+    python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
+      --benchmark mrcr --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
+      --attn-implementation sdpa --rows "$MRCR_ROWS" \
+      --method "$METHOD" --output "$OUT/mrcr_8k_2_4_8needle_n10.jsonl" \
+      >"$LOG/mrcr.log" 2>&1
+    ;;
+  *) echo "Unknown benchmark: $BENCHMARK" >&2; exit 2 ;;
+esac
+
+touch "$OUT/${BENCHMARK}.DONE"

recipe/scripts/eval_qwen_external_longctx_pilot.py

@@ -0,0 +1,275 @@
+#!/usr/bin/env python3
+"""Aligned 8K HELMET-ICL and OpenAI-MRCR pilot for Qwen checkpoints.
+
+This is deliberately a small diagnostic runner.  HELMET prompt construction
+matches the ICL schedule used by the local OmniServe reproduction; MRCR uses
+the official OpenAI rows selected by that reproduction, but renders the
+messages with the Qwen chat template.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import math
+import random
+import re
+import string
+from collections import defaultdict
+from difflib import SequenceMatcher
+from pathlib import Path
+
+import pandas as pd
+import torch
+from datasets import load_dataset
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+HELMET_SPECS = {
+    "trec_coarse": ("icl_trec_coarse_400shot_balance", 6),
+    "trec_fine": ("icl_trec_fine_400shot_balance", 50),
+    "banking77": ("icl_banking77_360shot_balance", 77),
+    "clinic150": ("icl_clinic150_440shot_balance", 151),
+    "nlu": ("icl_nlu_510shot_balance", 68),
+}
+
+
+def balanced(data, shots: int, label_field: str, seed: int):
+    rng = random.Random(seed)
+    by_label = defaultdict(list)
+    for item in data:
+        by_label[item[label_field]].append(item)
+    rounds = math.ceil(shots / len(by_label))
+    selected_rounds = [[] for _ in range(rounds)]
+    for examples in by_label.values():
+        indices = rng.sample(range(len(examples)), rounds % len(examples))
+        while len(indices) < rounds:
+            indices += rng.sample(range(len(examples)), min(rounds - len(indices), len(examples)))
+        for index, example_index in enumerate(indices):
+            selected_rounds[index].append(examples[example_index])
+    for examples in selected_rounds:
+        rng.shuffle(examples)
+    return [item for group in selected_rounds for item in group][:shots]
+
+
+def helmet_source(family: str, seed: int):
+    if family == "trec_coarse":
+        source = load_dataset("CogComp/trec", trust_remote_code=True)
+        return source["train"], source["test"], "text", "coarse_label"
+    if family == "trec_fine":
+        source = load_dataset("CogComp/trec", trust_remote_code=True)
+        return source["train"], source["test"], "text", "fine_label"
+    if family == "banking77":
+        source = load_dataset("PolyAI/banking77", trust_remote_code=True)
+        return source["train"], source["test"], "text", "label"
+    if family == "clinic150":
+        source = load_dataset("clinc/clinc_oos", "plus")
+        return source["train"], source["validation"], "text", "intent"
+    if family == "nlu":
+        source = load_dataset("xingkunliuxtracta/nlu_evaluation_data", trust_remote_code=True)["train"]
+        split = source.train_test_split(test_size=0.1, seed=seed)
+        return split["train"], split["test"], "text", "label"
+    raise ValueError(family)
+
+
+def build_helmet(seed: int, limit: int):
+    rows = []
+    user_template = (
+        'Use the provided mapping from the text to label to assign a label to the text. '
+        'Only output "label: {{label}}" and nothing else. \n\n{context}\n\n{question}'
+    )
+    for family, (dataset_name, num_labels) in HELMET_SPECS.items():
+        shots = int(dataset_name.split("shot")[0].split("_")[-1])
+        train, test, text_field, label_field = helmet_source(family, seed)
+        samples = balanced(test, limit, label_field, seed)
+        for sample_index, sample in enumerate(samples):
+            local_seed = (int(hashlib.sha256(sample[text_field].encode()).hexdigest(), 16) + seed) % 2**31
+            demos = balanced(train, shots, label_field, local_seed)
+            mapping = list(range(num_labels))
+            random.Random(local_seed).shuffle(mapping)
+            context = "\n\n".join(
+                f"{demo[text_field]}\nlabel: {mapping[int(demo[label_field])]}" for demo in demos
+            )
+            rows.append({
+                "row_id": f"helmet_icl:8k:{family}:{sample_index}",
+                "benchmark": "helmet_icl",
+                "config": family,
+                "prompt": user_template.format(context=context, question=sample[text_field]) + "\nlabel:",
+                "answer": str(mapping[int(sample[label_field])]),
+                "max_new_tokens": 20,
+                "metadata": {"dataset": dataset_name, "shots": shots},
+            })
+    return rows
+
+
+def build_mrcr(tokenizer, selection_path: Path, limit: int):
+    selection = [json.loads(line) for line in selection_path.read_text().splitlines() if line.strip()]
+    chosen = sorted(
+        (row for row in selection if int(row["context_k"]) == 8),
+        key=lambda row: int(row["sample_index"]),
+    )[:limit]
+    if len(chosen) < limit:
+        raise ValueError(f"MRCR selection contains {len(chosen)} 8K rows, requested {limit}")
+    files = [
+        hf_hub_download("openai/mrcr", filename=f"2needle/2needle_{shard}.parquet", repo_type="dataset")
+        for shard in (0, 1)
+    ]
+    frame = pd.concat([pd.read_parquet(path) for path in files], ignore_index=True)
+    rows = []
+    for selected in chosen:
+        source = frame.iloc[int(selected["source_index"])]
+        messages = json.loads(source["prompt"])
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+        answer = str(source["answer"])
+        sample_index = int(selected["sample_index"])
+        rows.append({
+            "row_id": f"mrcr:8k:2needle:{sample_index}",
+            "benchmark": "mrcr",
+            "config": "2needle",
+            "prompt": prompt,
+            "answer": answer,
+            "prefix": str(source["random_string_to_prepend"]),
+            "max_new_tokens": min(768, len(tokenizer(answer, add_special_tokens=False).input_ids) + 64),
+            "metadata": {"source_index": int(selected["source_index"]), "official_tokens": int(selected["official_tokens"])},
+        })
+    return rows
+
+
+def normalize(text: str) -> str:
+    text = text.lower()
+    text = re.sub(r"\b(a|an|the)\b", " ", text)
+    text = "".join(ch for ch in text if ch not in string.punctuation)
+    return " ".join(text.split())
+
+
+def score(row, prediction: str) -> float:
+    if row["benchmark"] == "helmet_icl":
+        first = prediction.strip().splitlines()[0] if prediction.strip() else ""
+        parsed = re.sub(r"^label:", "", first, flags=re.I).strip()
+        return float(normalize(parsed) == normalize(row["answer"]))
+    prefix = row["prefix"]
+    if not prediction.startswith(prefix):
+        return 0.0
+    response = prediction.removeprefix(prefix).strip()
+    reference = row["answer"].removeprefix(prefix).strip()
+    return float(SequenceMatcher(None, response, reference).ratio())
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--benchmark", choices=("helmet_icl", "mrcr"), required=True)
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument("--tokenizer-path", required=True)
+    parser.add_argument(
+        "--attn-implementation",
+        default="flash_attention_2",
+        choices=("eager", "sdpa", "flash_attention_2"),
+        help="Transformers attention backend; AHA 4-D local masks require SDPA or eager.",
+    )
+    parser.add_argument("--selection", type=Path)
+    parser.add_argument("--rows", type=Path, help="Frozen aligned prompt rows (JSONL)")
+    parser.add_argument("--method", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--seed", type=int, default=20260710)
+    parser.add_argument("--limit", type=int, default=1, help="Samples per benchmark config")
+    parser.add_argument(
+        "--stop-new-line",
+        action="store_true",
+        help="Use HELMET's newline stop-token policy during generation.",
+    )
+    parser.add_argument(
+        "--duo-sink-size",
+        type=int,
+        help="Optional inference-only override for the AHA sink token count.",
+    )
+    parser.add_argument(
+        "--duo-recent-size",
+        type=int,
+        help="Optional inference-only override for the AHA recent-window token count.",
+    )
+    args = parser.parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True, use_fast=False)
+    if args.rows:
+        rows = [json.loads(line) for line in args.rows.read_text().splitlines() if line.strip()]
+    else:
+        rows = (
+            build_helmet(args.seed, args.limit)
+            if args.benchmark == "helmet_icl"
+            else build_mrcr(tokenizer, args.selection, args.limit)
+        )
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation=args.attn_implementation,
+    )
+    for field, value in (
+        ("duo_sink_size", args.duo_sink_size),
+        ("duo_recent_size", args.duo_recent_size),
+    ):
+        if value is not None:
+            if value < 0:
+                parser.error(f"--{field.replace('_', '-')} must be non-negative")
+            previous = getattr(model.config, field, None)
+            setattr(model.config, field, value)
+            print(f"[eval] {field} override: {previous!r} -> {value!r}", flush=True)
+    model = model.to("cuda").eval()
+    stop_token_ids = model.generation_config.eos_token_id
+    stop_token_ids = list(stop_token_ids) if isinstance(stop_token_ids, list) else [stop_token_ids]
+    if args.stop_new_line:
+        newline_tokens = ["\n", "Ċ", "ĊĊ", "<0x0A>"]
+        stop_token_ids += [tokenizer.convert_tokens_to_ids(token) for token in newline_tokens]
+    stop_token_ids = sorted(
+        {
+            token_id
+            for token_id in stop_token_ids
+            if token_id is not None and token_id != tokenizer.unk_token_id
+        }
+    )
+    print(f"[eval] stop_token_ids={stop_token_ids}", flush=True)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    outputs = []
+    for ordinal, row in enumerate(rows, 1):
+        encoded = tokenizer(row["prompt"], return_tensors="pt", add_special_tokens=False)
+        prompt_tokens = int(encoded.input_ids.shape[1])
+        print(f"[{args.benchmark}] {ordinal}/{len(rows)} {row['config']} tokens={prompt_tokens}", flush=True)
+        encoded = {key: value.to("cuda") for key, value in encoded.items()}
+        with torch.inference_mode():
+            generated = model.generate(
+                **encoded,
+                do_sample=False,
+                max_new_tokens=row["max_new_tokens"],
+                use_cache=True,
+                eos_token_id=stop_token_ids,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+        prediction = tokenizer.decode(generated[0, prompt_tokens:], skip_special_tokens=True).strip()
+        outputs.append({
+            **row,
+            "method": args.method,
+            "prompt_tokens": prompt_tokens,
+            "prompt_sha256": hashlib.sha256(row["prompt"].encode()).hexdigest(),
+            "prediction": prediction,
+            "score": score(row, prediction),
+        })
+    args.output.write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in outputs))
+    summary = args.output.with_suffix(".summary.csv")
+    with summary.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=["method", "benchmark", "config", "n", "score", "prompt_tokens"])
+        writer.writeheader()
+        for row in outputs:
+            writer.writerow({"method": args.method, "benchmark": row["benchmark"], "config": row["config"], "n": 1, "score": row["score"], "prompt_tokens": row["prompt_tokens"]})
+    print(summary)
+
+
+if __name__ == "__main__":
+    main()

recipe/scripts/eval_sweep.sh

@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+REPO="$(cd "$RECIPE/.." && pwd)"
+OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
+EVAL_ROOT="${EVAL_ROOT:-$OUTPUT_ROOT/eval}"
+VANILLA_MODEL="${VANILLA_DIR:-$REPO}"
+AHA_MODEL="${AHA_MODEL:-$OUTPUT_ROOT/aha/stage2/checkpoint-25}"
+L2A_MODEL="${L2A_MODEL:-$OUTPUT_ROOT/l2a_style/stage2/checkpoint-25}"
+GPU_LIST="${GPU_LIST:-0}"
+THRESHOLDS=(0.45 0.50 0.525 0.55 0.575 0.60 0.625 0.65)
+BENCHMARKS=(ruler_a ruler_b babilong helmet mrcr)
+IFS=',' read -r -a GPUS <<< "$GPU_LIST"
+declare -A SLOT_PIDS=()
+mkdir -p "$EVAL_ROOT"
+
+for model in "$VANILLA_MODEL" "$AHA_MODEL" "$L2A_MODEL"; do
+  [[ -f "$model/config.json" ]] || { echo "Missing model: $model" >&2; exit 2; }
+done
+
+jobs=()
+for benchmark in "${BENCHMARKS[@]}"; do
+  jobs+=("vanilla|$VANILLA_MODEL|0.5|$benchmark")
+done
+for threshold in "${THRESHOLDS[@]}"; do
+  slug="${threshold/./}"
+  for benchmark in "${BENCHMARKS[@]}"; do
+    jobs+=("token_kv_head_t${slug}|$AHA_MODEL|$threshold|$benchmark")
+    jobs+=("token_t${slug}|$L2A_MODEL|$threshold|$benchmark")
+  done
+done
+
+launch() {
+  local gpu="$1" method="$2" model="$3" threshold="$4" benchmark="$5"
+  EVAL_ROOT="$EVAL_ROOT" RULER_LIMIT=20 BABILONG_LIMIT=50 \
+    AHA_GATE_HARD_THRESHOLD="$threshold" \
+    "$SCRIPT_DIR/eval_cell.sh" "$method" "$model" "$gpu" "$benchmark" &
+  LAST_PID="$!"
+}
+
+for i in "${!jobs[@]}"; do
+  slot=$((i % ${#GPUS[@]}))
+  if [[ -n "${SLOT_PIDS[$slot]:-}" ]]; then
+    wait "${SLOT_PIDS[$slot]}"
+  fi
+  IFS='|' read -r method model threshold benchmark <<< "${jobs[$i]}"
+  launch "${GPUS[$slot]}" "$method" "$model" "$threshold" "$benchmark"
+  SLOT_PIDS[$slot]="$LAST_PID"
+done
+
+status=0
+for pid in "${SLOT_PIDS[@]}"; do
+  wait "$pid" || status=1
+done
+[[ "$status" == 0 ]] || exit "$status"
+
+python "$SCRIPT_DIR/summarize_qwen1p7b_router_granularity_20260714.py" \
+  --repo "$REPO" \
+  --eval-root "$EVAL_ROOT" \
+  --baseline-root "$EVAL_ROOT/vanilla" \
+  --input-dir "$RECIPE/data/eval_inputs" \
+  --output-dir "$OUTPUT_ROOT/summary"

recipe/scripts/evaluate.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
+IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
+GPU_COUNT="$(nvidia-smi -L | wc -l)"
+if [[ "$GPU_COUNT" -lt 1 ]]; then
+  echo "At least one NVIDIA GPU is required." >&2
+  exit 2
+fi
+GPU_LIST="$(seq -s, 0 $((GPU_COUNT - 1)))"
+
+docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
+docker run --rm --gpus all --ipc=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -e GPU_LIST="$GPU_LIST" \
+  -v "$REPO:/workspace" \
+  -w /workspace/recipe \
+  "$IMAGE" bash scripts/eval_sweep.sh

recipe/scripts/export_tuned_vanilla_aha_hotstart.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""Convert a tuned vanilla Qwen3 checkpoint into a quality-safe AHA hot start.
+
+All vanilla weights are preserved. Dynamic gate weights start at zero and a
+trainable bias initializes every gate to the requested full-attention
+probability, so hard routing is initially exactly vanilla full attention.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+
+import torch
+from torch import nn
+from transformers import AutoTokenizer
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from modeling_aha_qwen3 import AHAQwen3Config, AHAQwen3ForCausalLM
+
+
+def add_zero_bias(linear: nn.Linear) -> nn.Linear:
+    new = nn.Linear(
+        linear.in_features,
+        linear.out_features,
+        bias=True,
+        device=linear.weight.device,
+        dtype=linear.weight.dtype,
+    )
+    with torch.no_grad():
+        new.weight.copy_(linear.weight)
+        new.bias.zero_()
+    return new
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--vanilla-path", required=True)
+    parser.add_argument("--output-path", required=True)
+    parser.add_argument("--window-size", type=int, default=128)
+    parser.add_argument("--gate-init-full-prob", type=float, default=0.90)
+    parser.add_argument(
+        "--local-kind", choices=("sliding_window", "sink_recent"), default="sliding_window"
+    )
+    parser.add_argument(
+        "--router-granularity",
+        choices=("token", "token_kv_head"),
+        default="token_kv_head",
+        help="Native dynamic gate shape. token is one shared gate per layer/token.",
+    )
+    args = parser.parse_args()
+
+    if not 0.5 < args.gate_init_full_prob < 1.0:
+        raise ValueError("--gate-init-full-prob must be strictly between 0.5 and 1")
+    output = Path(args.output_path)
+    output.mkdir(parents=True, exist_ok=True)
+
+    AHAQwen3Config.register_for_auto_class()
+    AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+    model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
+        args.vanilla_path,
+        aha_window_size=args.window_size,
+        aha_local_kind=args.local_kind,
+        aha_router_granularity=args.router_granularity,
+        aha_mode="dynamic",
+        aha_gate_target=0.70,
+        aha_reg_weight=-1.0,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+    )
+
+    # Qwen3 normally has bias-free projections. A gate bias gives a stable,
+    # token-independent full-attention hot start while zero biases on q/k/v/o
+    # keep the original attention computation bit-for-bit unchanged.
+    if not model.config.attention_bias:
+        for layer in model.model.layers:
+            attn = layer.self_attn
+            attn.q_proj = add_zero_bias(attn.q_proj)
+            attn.k_proj = add_zero_bias(attn.k_proj)
+            attn.v_proj = add_zero_bias(attn.v_proj)
+            attn.o_proj = add_zero_bias(attn.o_proj)
+        model.config.attention_bias = True
+
+    q_rows = model.config.num_attention_heads * model.config.head_dim
+    gate_logit = math.log(args.gate_init_full_prob / (1.0 - args.gate_init_full_prob))
+    with torch.no_grad():
+        for layer in model.model.layers:
+            q_proj = layer.self_attn.q_proj
+            q_proj.weight[q_rows:].zero_()
+            q_proj.bias[q_rows:].fill_(gate_logit)
+
+    # Keep a distinct LM head so the custom checkpoint has an explicit,
+    # self-contained state dict under safetensors.
+    model.config.tie_word_embeddings = False
+    if model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr():
+        model.lm_head.weight = nn.Parameter(model.lm_head.weight.detach().clone())
+    model.config.aha_hotstart_source = str(Path(args.vanilla_path).resolve())
+    model.config.aha_gate_init_full_prob = args.gate_init_full_prob
+
+    tokenizer = AutoTokenizer.from_pretrained(args.vanilla_path, trust_remote_code=True)
+    model.save_pretrained(output, safe_serialization=True)
+    tokenizer.save_pretrained(output)
+    gate_rows = model.model.layers[0].self_attn.aha_router_outputs
+    gate_parameters = len(model.model.layers) * gate_rows * (
+        model.config.hidden_size + 1
+    )
+    (output / "aha_hotstart_manifest.json").write_text(
+        json.dumps(
+            {
+                "source_checkpoint": str(Path(args.vanilla_path).resolve()),
+                "router_granularity": args.router_granularity,
+                "native_gate_rows_per_layer": gate_rows,
+                "effective_gate_parameters": gate_parameters,
+                "gate_init_full_probability": args.gate_init_full_prob,
+                "gate_weight_init": "zeros",
+                "gate_bias_logit": gate_logit,
+                "local_attention": {
+                    "kind": args.local_kind,
+                    "sink_size": model.config.duo_sink_size,
+                    "recent_size": model.config.duo_recent_size,
+                },
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+    print(
+        f"saved={output} init_full_prob={args.gate_init_full_prob:.4f} "
+        f"gate_logit={gate_logit:.6f} hard_sparsity=0.0 "
+        f"router_granularity={args.router_granularity} gate_params={gate_parameters:,}"
+    )
+
+
+if __name__ == "__main__":
+    main()

recipe/scripts/reproduce.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
+IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
+GPU_COUNT="$(nvidia-smi -L | wc -l)"
+if [[ "$GPU_COUNT" -lt 1 ]]; then
+  echo "At least one NVIDIA GPU is required." >&2
+  exit 2
+fi
+GPU_AHA="${GPU_AHA:-0}"
+if [[ -z "${GPU_L2A+x}" ]]; then
+  GPU_L2A=0
+  if [[ "$GPU_COUNT" -ge 2 ]]; then
+    GPU_L2A=1
+  fi
+fi
+
+for gpu in "$GPU_AHA" "$GPU_L2A"; do
+  if [[ ! "$gpu" =~ ^[0-9]+$ ]] || (( gpu >= GPU_COUNT )); then
+    echo "Invalid GPU index $gpu; nvidia-smi reports $GPU_COUNT visible GPU(s)." >&2
+    exit 2
+  fi
+done
+
+echo "Training mode: one GPU per experimental variant (no DDP/FSDP/DeepSpeed)."
+echo "AHA variant GPU: $GPU_AHA"
+echo "L2A-style variant GPU: $GPU_L2A"
+if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
+  echo "The two variants will run sequentially on the same GPU."
+else
+  echo "The two variants will run concurrently on two independent GPUs."
+fi
+if [[ "$GPU_COUNT" -gt 2 ]]; then
+  echo "Visible GPUs: $GPU_COUNT; this recipe intentionally does not use all GPUs."
+fi
+
+docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
+docker run --rm --gpus all --ipc=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -e GPU_AHA="$GPU_AHA" -e GPU_L2A="$GPU_L2A" \
+  -v "$REPO:/workspace" \
+  -w /workspace/recipe \
+  "$IMAGE" bash scripts/train_both.sh

recipe/scripts/summarize_qwen1p7b_router_granularity_20260714.py

@@ -0,0 +1,403 @@
+#!/usr/bin/env python3
+"""Audit and summarize the matched shared-gate vs KV-head-gate sweep."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import glob
+import hashlib
+import json
+import subprocess
+from collections import defaultdict
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+
+
+SUITES = ("RULER-local-full13", "BabiLong", "HELMET-ICL", "MRCR")
+THRESHOLDS = (0.45, 0.50, 0.525, 0.55, 0.575, 0.60, 0.625, 0.65)
+ARMS = ("token", "token_kv_head")
+SLUGS = {
+    "RULER-local-full13": "ruler",
+    "BabiLong": "babilong",
+    "HELMET-ICL": "helmet",
+    "MRCR": "mrcr",
+}
+SCORE_DEFINITIONS = {
+    "RULER-local-full13": "unweighted macro of lm-eval 8192 string-match scores over the local 13-config set",
+    "BabiLong": "unweighted macro exact-match accuracy over qa1-qa5",
+    "HELMET-ICL": "unweighted macro label exact-match over five ICL configurations",
+    "MRCR": "unweighted macro of prefix-check plus SequenceMatcher scores over 2/4/8-needle configurations",
+}
+
+
+def threshold_slug(value: float) -> str:
+    return {
+        0.45: "045",
+        0.50: "050",
+        0.525: "0525",
+        0.55: "055",
+        0.575: "0575",
+        0.60: "060",
+        0.625: "0625",
+        0.65: "065",
+    }[value]
+
+
+def one(pattern: str) -> Path:
+    paths = [Path(path) for path in glob.glob(pattern, recursive=True)]
+    if len(paths) != 1:
+        raise RuntimeError(f"expected one path for {pattern}, found {paths}")
+    return paths[0]
+
+
+def read_jsonl(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
+
+
+def collect_scores(root: Path) -> dict[tuple[str, str], dict]:
+    rows: dict[tuple[str, str], dict] = {}
+    for split in ("a", "b"):
+        result = one(str(root / f"ruler8k_splits/{split}/lm_eval/**/results_*.json"))
+        payload = json.loads(result.read_text())
+        for config, metrics in payload["results"].items():
+            if config in payload["n-samples"]:
+                rows[("RULER-local-full13", config)] = {
+                    "score": float(metrics["8192,none"]),
+                    "n": int(payload["n-samples"][config]["effective"]),
+                }
+    result = one(str(root / "babilong8k_qa1_qa5_n50/**/results_*.json"))
+    payload = json.loads(result.read_text())
+    for config, metrics in payload["results"].items():
+        if config in payload["n-samples"]:
+            rows[("BabiLong", config)] = {
+                "score": float(metrics["acc,none"]),
+                "n": int(payload["n-samples"][config]["effective"]),
+            }
+    for suite, filename in (
+        ("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
+        ("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
+    ):
+        grouped: dict[str, list[dict]] = defaultdict(list)
+        for row in read_jsonl(root / filename):
+            grouped[row["config"]].append(row)
+        for config, values in grouped.items():
+            rows[(suite, config)] = {
+                "score": sum(float(row["score"]) for row in values) / len(values),
+                "n": len(values),
+            }
+    return rows
+
+
+def sample_signatures(root: Path, score_rows: dict[tuple[str, str], dict]) -> dict[str, str]:
+    signatures: dict[str, str] = {}
+    for suite, rel_dirs in (
+        ("RULER-local-full13", ("ruler8k_splits/a/lm_eval", "ruler8k_splits/b/lm_eval")),
+        ("BabiLong", ("babilong8k_qa1_qa5_n50",)),
+    ):
+        configs = [config for row_suite, config in score_rows if row_suite == suite]
+        for config in configs:
+            matches = []
+            for rel_dir in rel_dirs:
+                matches.extend(root.glob(f"{rel_dir}/**/samples_{config}_*.jsonl"))
+            if len(matches) != 1:
+                raise RuntimeError(f"expected one sample log for {suite}/{config}: {matches}")
+            values = [
+                (
+                    row.get("doc_id"),
+                    row.get("doc_hash"),
+                    row.get("prompt_hash"),
+                    row.get("target_hash"),
+                )
+                for row in read_jsonl(matches[0])
+            ]
+            signatures[f"{suite}/{config}"] = hashlib.sha256(
+                json.dumps(values, sort_keys=True).encode()
+            ).hexdigest()
+    for suite, filename in (
+        ("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
+        ("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
+    ):
+        grouped: dict[str, list[tuple]] = defaultdict(list)
+        for row in read_jsonl(root / filename):
+            grouped[row["config"]].append(
+                (row["row_id"], row.get("prompt_sha256"), row.get("target"))
+            )
+        for config, values in grouped.items():
+            signatures[f"{suite}/{config}"] = hashlib.sha256(
+                json.dumps(values, sort_keys=True).encode()
+            ).hexdigest()
+    return signatures
+
+
+def read_sparsity(root: Path, suite: str) -> dict:
+    if suite == "RULER-local-full13":
+        parts = [
+            json.loads((root / f"sparsity_ruler_{split}.json").read_text())
+            for split in ("a", "b")
+        ]
+        sparse = sum(int(part["sparse_decisions"]) for part in parts)
+        total = sum(int(part["total_decisions"]) for part in parts)
+        granularities = {part.get("router_granularity") for part in parts}
+        native = sum(int(part.get("native_router_decisions", 0)) for part in parts)
+        phases: dict[str, dict[str, int]] = defaultdict(lambda: {"sparse": 0, "total": 0})
+        for part in parts:
+            for phase, values in part.get("by_phase", {}).items():
+                phases[phase]["sparse"] += int(values["sparse_decisions"])
+                phases[phase]["total"] += int(values["total_decisions"])
+    else:
+        part = json.loads((root / f"sparsity_{SLUGS[suite]}.json").read_text())
+        sparse = int(part["sparse_decisions"])
+        total = int(part["total_decisions"])
+        granularities = {part.get("router_granularity")}
+        native = int(part.get("native_router_decisions", 0))
+        phases = {
+            phase: {
+                "sparse": int(values["sparse_decisions"]),
+                "total": int(values["total_decisions"]),
+            }
+            for phase, values in part.get("by_phase", {}).items()
+        }
+    return {
+        "sparse_decisions": sparse,
+        "total_decisions": total,
+        "sparsity": sparse / total,
+        "full_attention_usage": 1.0 - sparse / total,
+        "router_granularities": sorted(value for value in granularities if value),
+        "native_router_decisions": native or None,
+        "by_phase": phases,
+    }
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def command_output(command: list[str], cwd: Path) -> str:
+    try:
+        return subprocess.run(
+            command, cwd=cwd, check=True, text=True, capture_output=True
+        ).stdout.strip()
+    except (OSError, subprocess.CalledProcessError):
+        return "unavailable in downloaded HF snapshot"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--repo", type=Path, default=Path("/data/sjm/AHA/AHA-Qwen3")
+    )
+    parser.add_argument(
+        "--eval-root",
+        type=Path,
+        default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_router_granularity_20260714/eval"),
+    )
+    parser.add_argument(
+        "--baseline-root",
+        type=Path,
+        default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/vanilla"),
+    )
+    parser.add_argument("--output-dir", type=Path)
+    parser.add_argument("--input-dir", type=Path)
+    args = parser.parse_args()
+    output = args.output_dir or args.eval_root.parent / "summary"
+    output.mkdir(parents=True, exist_ok=True)
+
+    baseline = collect_scores(args.baseline_root)
+    baseline_signatures = sample_signatures(args.baseline_root, baseline)
+    curve_rows = []
+    per_config = []
+    points = []
+    alignment = {}
+    done_files = ("ruler_a.DONE", "ruler_b.DONE", "babilong.DONE", "helmet.DONE", "mrcr.DONE")
+
+    for arm in ARMS:
+        for threshold in THRESHOLDS:
+            method = f"{arm}_t{threshold_slug(threshold)}"
+            root = args.eval_root / method
+            missing = [name for name in done_files if not (root / name).exists()]
+            if missing:
+                raise RuntimeError(f"incomplete {method}: missing {missing}")
+            scores = collect_scores(root)
+            if set(scores) != set(baseline):
+                raise RuntimeError(f"configuration mismatch for {method}")
+            if any(scores[key]["n"] != baseline[key]["n"] for key in baseline):
+                raise RuntimeError(f"sample-count mismatch for {method}")
+            signatures = sample_signatures(root, scores)
+            mismatches = sorted(
+                key for key in baseline_signatures if signatures.get(key) != baseline_signatures[key]
+            )
+            alignment[method] = {"aligned": not mismatches, "mismatches": mismatches}
+            if mismatches:
+                raise RuntimeError(f"prompt/target hash mismatch for {method}: {mismatches}")
+
+            suite_rows = []
+            sparse_sum = total_sum = 0
+            for suite in SUITES:
+                keys = sorted(key for key in baseline if key[0] == suite)
+                baseline_score = sum(baseline[key]["score"] for key in keys) / len(keys)
+                score = sum(scores[key]["score"] for key in keys) / len(keys)
+                sparsity = read_sparsity(root, suite)
+                expected_granularity = arm
+                if sparsity["router_granularities"] != [expected_granularity]:
+                    raise RuntimeError(
+                        f"{method}/{suite} recorded {sparsity['router_granularities']}, "
+                        f"expected {[expected_granularity]}"
+                    )
+                sparse_sum += sparsity["sparse_decisions"]
+                total_sum += sparsity["total_decisions"]
+                row = {
+                    "arm": arm,
+                    "router_granularity": arm,
+                    "threshold": threshold,
+                    "benchmark": suite,
+                    "config_count": len(keys),
+                    "samples_per_config": ",".join(
+                        str(value) for value in sorted({baseline[key]["n"] for key in keys})
+                    ),
+                    "total_samples": sum(baseline[key]["n"] for key in keys),
+                    "tuned_vanilla_score": baseline_score,
+                    "quality_floor_95pct": 0.95 * baseline_score,
+                    "score": score,
+                    "retention": score / baseline_score if baseline_score else None,
+                    "quality_pass": score >= 0.95 * baseline_score,
+                    **{key: sparsity[key] for key in (
+                        "sparse_decisions", "total_decisions", "sparsity",
+                        "full_attention_usage", "native_router_decisions", "by_phase"
+                    )},
+                }
+                curve_rows.append(row)
+                suite_rows.append(row)
+                for key in keys:
+                    per_config.append({
+                        "arm": arm,
+                        "threshold": threshold,
+                        "benchmark": suite,
+                        "config": key[1],
+                        "n": baseline[key]["n"],
+                        "tuned_vanilla_score": baseline[key]["score"],
+                        "score": scores[key]["score"],
+                    })
+            points.append({
+                "arm": arm,
+                "threshold": threshold,
+                "all_suite_quality_pass": all(row["quality_pass"] for row in suite_rows),
+                "decision_weighted_sparsity": sparse_sum / total_sum,
+                "full_attention_usage": 1.0 - sparse_sum / total_sum,
+                "sparse_decisions": sparse_sum,
+                "total_decisions": total_sum,
+                "suites": suite_rows,
+            })
+
+    headline = {}
+    for arm in ARMS:
+        valid = [
+            point for point in points
+            if point["arm"] == arm and point["all_suite_quality_pass"]
+        ]
+        if valid:
+            best = max(valid, key=lambda point: point["decision_weighted_sparsity"])
+            headline[arm] = {"found": True, **best}
+        else:
+            headline[arm] = {
+                "found": False,
+                "statement": "No measured threshold preserved at least 95% of tuned vanilla on every suite.",
+            }
+
+    input_dir = args.input_dir or args.baseline_root.parent / "inputs"
+    frozen_inputs = [
+        input_dir / "helmet_icl_8k_n50_per_config.jsonl",
+        input_dir / "mrcr_8k_2_4_8needle_n10_per_config.jsonl",
+    ]
+    payload = {
+        "protocol": {
+            "model": "Qwen3-1.7B tuned vanilla",
+            "context_length": 8192,
+            "arms": {
+                "token": "L2A-style shared-gate: one native gate per token/layer",
+                "token_kv_head": "AHA: one native gate per token/KV-head/layer",
+            },
+            "local_attention": {"sink_tokens": 64, "recent_tokens": 256},
+            "inference": "strict AHA routing in both prefill and decode; no force-full heads or full-decode fallback",
+            "thresholds": list(THRESHOLDS),
+            "quality_rule": "every suite macro score >= 95% of aligned tuned-vanilla macro score",
+            "headline_rule": "highest decision-weighted measured sparsity among thresholds passing every suite",
+            "sparsity_definition": "hard local routes / token x KV-head x layer effective decisions",
+            "suite_scope": "RULER local full13 means this repository's fixed 13-config set, not complete upstream RULER",
+        },
+        "score_definitions": SCORE_DEFINITIONS,
+        "headline": headline,
+        "points": points,
+        "curves": curve_rows,
+        "per_config": per_config,
+        "alignment": {"all_aligned": all(value["aligned"] for value in alignment.values()), "methods": alignment},
+        "appendix_note": (
+            "Upstream L2A results are not merged into this table because architecture, "
+            "training data, objective, and evaluation protocols differ."
+        ),
+        "reproducibility": {
+            "git_commit": command_output(["git", "rev-parse", "HEAD"], args.repo),
+            "git_status_short": command_output(["git", "status", "--short"], args.repo),
+            "container": "dmtd-repro",
+            "container_image": command_output(
+                ["docker", "inspect", "-f", "{{.Config.Image}}", "dmtd-repro"], args.repo
+            ),
+            "baseline_root": str(args.baseline_root),
+            "eval_root": str(args.eval_root),
+            "frozen_input_sha256": {
+                str(path): sha256(path) for path in frozen_inputs if path.exists()
+            },
+        },
+    }
+    (output / "router_granularity_results.json").write_text(
+        json.dumps(payload, indent=2) + "\n"
+    )
+
+    flat_curve_rows = [
+        {key: value for key, value in row.items() if key != "by_phase"}
+        for row in curve_rows
+    ]
+    with (output / "quality_sparsity_curves.csv").open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(flat_curve_rows[0]))
+        writer.writeheader()
+        writer.writerows(flat_curve_rows)
+    with (output / "per_config_scores.csv").open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(per_config[0]))
+        writer.writeheader()
+        writer.writerows(per_config)
+
+    figure, axes = plt.subplots(2, 2, figsize=(12, 9), constrained_layout=True)
+    colors = {"token": "#d62728", "token_kv_head": "#1f77b4"}
+    labels = {"token": "L2A-style shared-gate", "token_kv_head": "AHA KV-head gate"}
+    for axis, suite in zip(axes.flat, SUITES):
+        for arm in ARMS:
+            values = sorted(
+                (row for row in curve_rows if row["arm"] == arm and row["benchmark"] == suite),
+                key=lambda row: row["sparsity"],
+            )
+            axis.plot(
+                [100 * row["sparsity"] for row in values],
+                [100 * row["retention"] for row in values],
+                marker="o", color=colors[arm], label=labels[arm],
+            )
+            for row in values:
+                axis.annotate(f"{row['threshold']:.3g}", (100 * row["sparsity"], 100 * row["retention"]), fontsize=7)
+        axis.axhline(95, color="black", linestyle="--", linewidth=1)
+        axis.set_title(suite)
+        axis.set_xlabel("Measured effective sparsity (%)")
+        axis.set_ylabel("Retention vs tuned vanilla (%)")
+        axis.grid(alpha=0.25)
+    axes.flat[0].legend()
+    figure.suptitle("Qwen3-1.7B matched router-granularity quality–sparsity curves")
+    figure.savefig(output / "quality_sparsity_curves.png", dpi=180)
+    plt.close(figure)
+    print(json.dumps(headline, indent=2))
+
+
+if __name__ == "__main__":
+    main()

recipe/scripts/train_aha.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec "$SCRIPT_DIR/train_one.sh" aha "${1:-0}"

recipe/scripts/train_both.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+GPU_AHA="${GPU_AHA:-0}"
+GPU_L2A="${GPU_L2A:-1}"
+cd "$RECIPE"
+
+echo "AHA = per-token, per-KV-head gate; assigned visible GPU $GPU_AHA."
+echo "L2A-style = per-token, head-shared gate; assigned visible GPU $GPU_L2A."
+echo "Each variant is a separate world-size-1 run with global batch size 1."
+
+python "$SCRIPT_DIR/validate_release.py"
+python -m unittest -v tests.test_router_granularity
+
+if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
+  "$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA"
+  "$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A"
+else
+  "$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA" & aha_pid=$!
+  "$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A" & l2a_pid=$!
+  status=0
+  wait "$aha_pid" || status=1
+  wait "$l2a_pid" || status=1
+  [[ "$status" == 0 ]] || exit "$status"
+fi
+
+echo "AHA selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/aha/stage2/checkpoint-25"
+echo "L2A-style selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/l2a_style/stage2/checkpoint-25"

recipe/scripts/train_l2a_style.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec "$SCRIPT_DIR/train_one.sh" l2a_style "${1:-0}"

recipe/scripts/train_one.sh

@@ -0,0 +1,127 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# -lt 1 || $# -gt 2 ]]; then
+  echo "Usage: $0 aha|l2a_style [GPU_ID]" >&2
+  exit 2
+fi
+
+ARM="$1"
+GPU="${2:-0}"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+REPO="$(cd "$RECIPE/.." && pwd)"
+VANILLA_DIR="${VANILLA_DIR:-$REPO}"
+DATA_DIR="${DATA_DIR:-$RECIPE/data/am_distilled_long_mix}"
+OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
+
+case "$ARM" in
+  aha) GRANULARITY=token_kv_head ;;
+  l2a_style) GRANULARITY=token ;;
+  *) echo "Unknown arm: $ARM (expected aha or l2a_style)" >&2; exit 2 ;;
+esac
+
+ARM_OUT="$OUTPUT_ROOT/$ARM"
+HOTSTART="$ARM_OUT/hotstart"
+STAGE1="$ARM_OUT/stage1"
+STAGE2="$ARM_OUT/stage2"
+LOG_DIR="$ARM_OUT/logs"
+mkdir -p "$LOG_DIR"
+
+export CUDA_VISIBLE_DEVICES="$GPU"
+export PYTHONUNBUFFERED=1
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
+
+if [[ ! -f "$VANILLA_DIR/model.safetensors" ]]; then
+  echo "Missing tuned-vanilla checkpoint: $VANILLA_DIR/model.safetensors" >&2
+  exit 2
+fi
+
+if [[ ! -f "$HOTSTART/config.json" ]]; then
+  python "$SCRIPT_DIR/export_tuned_vanilla_aha_hotstart.py" \
+    --vanilla-path "$VANILLA_DIR" \
+    --output-path "$HOTSTART" \
+    --window-size 128 \
+    --local-kind sink_recent \
+    --gate-init-full-prob 0.90 \
+    --router-granularity "$GRANULARITY" \
+    2>&1 | tee "$LOG_DIR/hotstart.log"
+fi
+
+if [[ ! -f "$STAGE1/checkpoint-300/config.json" ]]; then
+  AHA_TRAIN_GATE_HARD_THRESHOLD=0.50 \
+  python "$RECIPE/dynamic_duo_train.py" \
+    --aha_checkpoint "$HOTSTART" \
+    --model_path "$VANILLA_DIR" \
+    --output_dir "$STAGE1" \
+    --data_source am_distilled \
+    --am_dataset_path "$DATA_DIR" \
+    --am_dataset_split train \
+    --am_label_mode full \
+    --max_length 8192 \
+    --num_steps 300 \
+    --warmup_ratio 0.10 \
+    --lr 3e-5 \
+    --reg_weight 0.1 \
+    --ce_weight 0.0 \
+    --batch_size 1 \
+    --grad_accum 1 \
+    --save_steps 100 \
+    --log_steps 10 \
+    --seed 42 \
+    --dtype bfloat16 \
+    --attn_impl sdpa \
+    --aha_local_kind sink_recent \
+    --router_granularity "$GRANULARITY" \
+    2>&1 | tee "$LOG_DIR/stage1.log"
+fi
+
+if [[ ! -f "$STAGE2/checkpoint-75/config.json" ]]; then
+  MODEL_PATH="$VANILLA_DIR" \
+  AHA_CHECKPOINT_PATH="$STAGE1/checkpoint-300" \
+  DATASET_PATH="$DATA_DIR" \
+  DATASET_SPLIT=train \
+  OUTPUT_DIR="$STAGE2" \
+  AHA_MODE=dynamic \
+  AHA_ROUTER_GRANULARITY="$GRANULARITY" \
+  AHA_LOCAL_KIND=sink_recent \
+  AHA_CE_WEIGHT=1.0 \
+  AHA_DISTILL_WEIGHT=0.5 \
+  AHA_REG_WEIGHT=0.01 \
+  AHA_TRAIN_GATE_HARD_THRESHOLD=0.58 \
+  GROUPED_LR=1 \
+  GATE_ONLY=0 \
+  LEARNING_RATE=3e-6 \
+  GATE_LEARNING_RATE=3e-6 \
+  BACKBONE_LEARNING_RATE=3e-7 \
+  LR_SCHEDULER_TYPE=constant_with_warmup \
+  WARMUP_RATIO=0.10 \
+  WEIGHT_DECAY=0.0 \
+  MAX_SEQ_LENGTH=8192 \
+  PER_DEVICE_TRAIN_BATCH_SIZE=1 \
+  GRADIENT_ACCUMULATION_STEPS=1 \
+  MAX_STEPS=75 \
+  LOGGING_STEPS=5 \
+  SAVE_STEPS=25 \
+  SAVE_TOTAL_LIMIT=3 \
+  FREEZE_EMBEDDINGS_LM_HEAD=1 \
+  REPORT_TO=none \
+  SEED=47 \
+  python "$RECIPE/sft.py" \
+    2>&1 | tee "$LOG_DIR/stage2.log"
+fi
+
+python - "$STAGE2/checkpoint-25" "$GRANULARITY" <<'PY'
+import json
+import sys
+from pathlib import Path
+
+checkpoint = Path(sys.argv[1])
+expected = sys.argv[2]
+config = json.loads((checkpoint / "config.json").read_text())
+actual = config.get("aha_router_granularity", "token_kv_head")
+if actual != expected:
+    raise RuntimeError(f"checkpoint granularity {actual!r} != {expected!r}")
+print(f"selected_checkpoint={checkpoint} router_granularity={actual}")
+PY

recipe/scripts/validate_release.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""Fail fast if a downloaded release is incomplete or has drifted."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from collections import Counter
+from pathlib import Path
+
+from datasets import load_from_disk
+
+
+RECIPE = Path(__file__).resolve().parents[1]
+REPO = RECIPE.parent
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def read_jsonl(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
+
+
+def main() -> None:
+    manifest = json.loads((RECIPE / "manifest.json").read_text())
+    checked = []
+    for relative, expected in manifest["files"].items():
+        path = (RECIPE / relative).resolve()
+        if not path.exists():
+            raise FileNotFoundError(path)
+        if path.name == "model.safetensors" and os.environ.get("SKIP_LARGE_HASH") == "1":
+            continue
+        actual = sha256(path)
+        if actual != expected:
+            raise RuntimeError(f"SHA256 mismatch for {path}: {actual} != {expected}")
+        checked.append(str(path.relative_to(REPO)))
+
+    config = json.loads((REPO / "config.json").read_text())
+    if config.get("model_type") != "qwen3":
+        raise RuntimeError(f"unexpected tuned-vanilla model_type: {config.get('model_type')}")
+    dataset = load_from_disk(str(RECIPE / "data/am_distilled_long_mix"))
+    if len(dataset["train"]) != 1024:
+        raise RuntimeError(f"expected 1024 training rows, found {len(dataset['train'])}")
+
+    helmet = read_jsonl(RECIPE / "data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl")
+    mrcr = read_jsonl(RECIPE / "data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl")
+    helmet_counts = Counter(row["config"] for row in helmet)
+    mrcr_counts = Counter(row["config"] for row in mrcr)
+    if sorted(helmet_counts.values()) != [50] * 5:
+        raise RuntimeError(f"unexpected HELMET counts: {helmet_counts}")
+    if sorted(mrcr_counts.values()) != [10] * 3:
+        raise RuntimeError(f"unexpected MRCR counts: {mrcr_counts}")
+
+    print(json.dumps({
+        "status": "ok",
+        "checked_sha256": checked,
+        "model_type": config["model_type"],
+        "training_rows": len(dataset["train"]),
+        "helmet_rows": len(helmet),
+        "mrcr_rows": len(mrcr),
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()

recipe/sft.py

@@ -0,0 +1,642 @@
+"""
+AHA-Qwen3 SFT training script.
+
+Loads pretrained Qwen3-0.6B, converts to AHA-Qwen3 (adds gate to q_proj),
+and trains on am-distilled data. Only gate weights are randomly initialized;
+all other weights come from the pretrained model.
+"""
+
+import json
+import math
+import os
+import shutil
+import socket
+import subprocess
+from typing import Optional
+
+import datasets
+import torch
+import torch.distributed as dist
+from trl import SFTConfig, SFTTrainer
+from transformers import AutoTokenizer, TrainerCallback
+from transformers.trainer_utils import get_last_checkpoint
+
+from modeling_aha_qwen3 import (
+    AHA_ROUTER_GRANULARITY,
+    AHAQwen3Config,
+    AHAQwen3ForCausalLM,
+    aha_router_output_size,
+)
+from router_training_utils import RowWiseAdamW, configure_gate_only
+
+
+class AHASFTTrainer(SFTTrainer):
+    """Accumulate CE vs gate-aux vs distill breakdown and gate density; log with HF `loss` (total)."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # ``AHAQwen3ForCausalLM.forward`` accepts ``**kwargs`` for attention
+        # backends, but its CE loss is an ordinary local token mean and does
+        # not consume Trainer's ``num_items_in_batch``.  Transformers 4.57
+        # otherwise mistakes the variadic signature for a globally normalized
+        # loss, multiplies it by world size, and skips the normal gradient-
+        # accumulation division.  Mark the actual loss contract explicitly so
+        # lr=3e-5 has the same meaning at every world size.
+        self.model_accepts_loss_kwargs = False
+        self._aha_ce_sum = 0.0
+        self._aha_gate_aux_sum = 0.0
+        self._aha_distill_sum = 0.0
+        self._aha_gate_soft_sum = 0.0
+        self._aha_gate_hard_sum = 0.0
+        self._aha_metric_count = 0
+        # Cached for loss_total aggregation in log().
+        self._aha_distill_weight = float(
+            getattr(self.model.config, "aha_distill_weight", 0.0)
+        )
+        self._aha_ce_weight = float(
+            getattr(self.model.config, "aha_ce_weight", 1.0)
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        loss, outputs = super().compute_loss(
+            model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch
+        )
+        if getattr(outputs, "ce_loss", None) is not None:
+            self._aha_metric_count += 1
+            self._aha_ce_sum += float(outputs.ce_loss.detach().float().mean().cpu())
+            if getattr(outputs, "gate_aux_loss", None) is not None:
+                self._aha_gate_aux_sum += float(outputs.gate_aux_loss.detach().float().mean().cpu())
+            if getattr(outputs, "distill_loss", None) is not None:
+                self._aha_distill_sum += float(outputs.distill_loss.detach().float().mean().cpu())
+            if getattr(outputs, "gate_soft_mean", None) is not None:
+                self._aha_gate_soft_sum += float(outputs.gate_soft_mean.detach().float().mean().cpu())
+            if getattr(outputs, "gate_hard_mean", None) is not None:
+                self._aha_gate_hard_sum += float(outputs.gate_hard_mean.detach().float().mean().cpu())
+        if return_outputs:
+            return (loss, outputs)
+        return loss
+
+    def _distributed_mean_of_sums(self, sum_val: float, count: int) -> float:
+        if count <= 0:
+            return float("nan")
+        device = self.accelerator.device
+        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+            t = torch.tensor([sum_val, float(count)], device=device, dtype=torch.float64)
+            dist.all_reduce(t, op=dist.ReduceOp.SUM)
+            return (t[0] / t[1]).item()
+        return sum_val / float(count)
+
+    def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None:
+        if self._aha_metric_count > 0:
+            n = self._aha_metric_count
+            ce_m = self._distributed_mean_of_sums(self._aha_ce_sum, n)
+            aux_m = self._distributed_mean_of_sums(self._aha_gate_aux_sum, n)
+            distill_m = self._distributed_mean_of_sums(self._aha_distill_sum, n)
+            gs_m = self._distributed_mean_of_sums(self._aha_gate_soft_sum, n)
+            gh_m = self._distributed_mean_of_sums(self._aha_gate_hard_sum, n)
+            if not math.isnan(ce_m):
+                logs["ce_loss"] = round(ce_m, 4)
+            if not math.isnan(aux_m):
+                logs["gate_aux_loss"] = round(aux_m, 6)
+            if not math.isnan(distill_m) and distill_m != 0.0:
+                logs["distill_loss"] = round(distill_m, 6)
+            # loss_total = aha_ce_weight * ce + gate_aux
+            #              + aha_distill_weight * distill_loss
+            # (distill_m / ce_m are unweighted; apply weights here so the
+            # sum matches the scalar actually added to total loss.)
+            if not math.isnan(ce_m) and not math.isnan(aux_m):
+                lt = self._aha_ce_weight * ce_m + aux_m
+                if not math.isnan(distill_m):
+                    lt = lt + self._aha_distill_weight * distill_m
+                logs["loss_total"] = round(lt, 4)
+                if "loss" in logs and lt > 1e-8:
+                    # With ``model_accepts_loss_kwargs=False``, HF's logging
+                    # window and the raw model loss should agree up to normal
+                    # batch-to-batch weighting differences. A large integer
+                    # ratio is a regression in Trainer loss normalization.
+                    logs["hf_loss_ratio"] = round(logs["loss"] / lt, 2)
+            if not math.isnan(gs_m):
+                logs["gate_soft_mean"] = round(gs_m, 4)
+            if not math.isnan(gh_m):
+                logs["gate_hard_mean"] = round(gh_m, 4)
+            self._aha_ce_sum = 0.0
+            self._aha_gate_aux_sum = 0.0
+            self._aha_distill_sum = 0.0
+            self._aha_gate_soft_sum = 0.0
+            self._aha_gate_hard_sum = 0.0
+            self._aha_metric_count = 0
+        return super().log(logs, start_time)
+
+
+class AHAManifestCallback(TrainerCallback):
+    """Copy the immutable run manifest into each Trainer checkpoint."""
+
+    def __init__(self, manifest_path: str):
+        self.manifest_path = manifest_path
+
+    def on_save(self, args, state, control, **kwargs):
+        if state.is_world_process_zero and os.path.exists(self.manifest_path):
+            checkpoint_dir = os.path.join(
+                args.output_dir, f"checkpoint-{state.global_step}"
+            )
+            if os.path.isdir(checkpoint_dir):
+                shutil.copyfile(
+                    self.manifest_path,
+                    os.path.join(checkpoint_dir, "aha_training_manifest.json"),
+                )
+        return control
+
+
+def get_optional_int(name: str) -> Optional[int]:
+    value = os.environ.get(name)
+    return None if value in {None, ""} else int(value)
+
+
+def resolve_resume_checkpoint(output_dir: str) -> Optional[str]:
+    resume_from_checkpoint = os.environ.get("RESUME_FROM_CHECKPOINT")
+    if resume_from_checkpoint in {None, ""}:
+        return None
+    if resume_from_checkpoint == "latest":
+        if not os.path.isdir(output_dir):
+            return None
+        return get_last_checkpoint(output_dir)
+    return resume_from_checkpoint
+
+
+def parse_fsdp_options(name: str) -> list[str]:
+    value = os.environ.get(name, "").replace(",", " ").strip()
+    return [item for item in value.split() if item]
+
+
+def parse_bool(name: str, default: bool = False) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return value.lower() in {"1", "true", "yes", "on"}
+
+
+def parse_report_to() -> list[str]:
+    """Comma-separated integrations, e.g. REPORT_TO=wandb or wandb,tensorboard. Empty / none / off -> []."""
+    raw = os.environ.get("REPORT_TO", "").strip()
+    if not raw or raw.lower() in ("none", "off"):
+        return []
+    return [x.strip() for x in raw.split(",") if x.strip()]
+
+
+def build_grouped_lr_optimizer(model: torch.nn.Module) -> torch.optim.Optimizer:
+    """Build an AdamW optimizer with separate *effective* LR for gate rows.
+
+    Dynamic-mode gate logits live in the final router rows of
+    each attention ``q_proj`` instead of in a separate Parameter. PyTorch
+    optimizer groups cannot split one Parameter by row, so q_proj tensors are
+    placed in the gate-lr group and :class:`RowWiseAdamW` scales their realized
+    Q-row updates by ``backbone_lr / gate_lr``. Other trainable parameters use
+    the backbone-lr group directly.
+    """
+    if getattr(model.config, "aha_mode", "dynamic") != "dynamic":
+        raise ValueError("GROUPED_LR currently only supports AHA_MODE=dynamic")
+    if gate_learning_rate <= 0.0 or backbone_learning_rate < 0.0:
+        raise ValueError(
+            "GATE_LEARNING_RATE must be positive and BACKBONE_LEARNING_RATE must be non-negative"
+        )
+
+    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
+    q_row_scale = backbone_learning_rate / gate_learning_rate
+
+    gate_params = []
+    gate_param_ids = set()
+
+    row_scales = []
+    for layer in model.model.layers:
+        q_proj = layer.self_attn.q_proj
+        for p in (q_proj.weight, q_proj.bias):
+            if p is None or not p.requires_grad:
+                continue
+            gate_params.append(p)
+            gate_param_ids.add(id(p))
+            row_scales.append((p, q_rows, q_row_scale))
+
+    backbone_params = [
+        p for p in model.parameters()
+        if p.requires_grad and id(p) not in gate_param_ids
+    ]
+    if not gate_params:
+        raise ValueError("GROUPED_LR found no trainable q_proj gate parameters")
+
+    gate_n = sum(p.numel() for p in gate_params)
+    effective_gate_n = sum(
+        layer.self_attn.aha_router_outputs
+        * (layer.self_attn.q_proj.in_features + (layer.self_attn.q_proj.bias is not None))
+        for layer in model.model.layers
+    )
+    backbone_n = sum(p.numel() for p in backbone_params)
+    print(
+        "  GROUPED_LR optimizer: "
+        f"gate_lr={gate_learning_rate:.3e}, backbone_lr={backbone_learning_rate:.3e}, "
+        f"q_row_update_scale={q_row_scale:.3g}, "
+        f"q_proj_params={gate_n:,}, effective_gate_params={effective_gate_n:,}, "
+        f"other_trainable={backbone_n:,}"
+    )
+
+    param_groups = [{"params": gate_params, "lr": gate_learning_rate}]
+    if backbone_params:
+        param_groups.append({"params": backbone_params, "lr": backbone_learning_rate})
+    return RowWiseAdamW(
+        param_groups,
+        row_scales=row_scales,
+        weight_decay=weight_decay,
+        betas=(adam_beta1, adam_beta2),
+    )
+
+
+# === Configuration (override via environment variables) ===
+model_path = os.environ.get("MODEL_PATH", "/workspace/AHA/models/Qwen3-0.6B")
+aha_checkpoint_path = os.environ.get("AHA_CHECKPOINT_PATH", "")
+dataset_path = os.environ.get("DATASET_PATH", "/workspace/Direct-Multitoken-Decoding/am-distilled-8192")
+dataset_split = os.environ.get("DATASET_SPLIT", "train")
+output_dir = os.environ.get("OUTPUT_DIR", "/workspace/AHA/AHA-Qwen3/ckpts/aha_qwen3_w1024")
+aha_window_size = int(os.environ.get("AHA_WINDOW_SIZE", "1024"))
+aha_lambda = float(os.environ.get("AHA_LAMBDA", "3e-4"))
+aha_distill_weight = float(os.environ.get("AHA_DISTILL_WEIGHT", "0.0"))
+# Language-modeling CE weight. Default 1.0 = standard SFT. Set to 0.0 to
+# drop CE from the training objective and shape the gate purely via
+# ``aux + distill`` (useful with GATE_ONLY + frozen backbone).
+aha_ce_weight = float(os.environ.get("AHA_CE_WEIGHT", "1.0"))
+# Hinge target on mean(gate_soft). Aux loss is ``λ · max(0, ḡ - τ)``.
+# Default 1.0 keeps the legacy unconditional aux (τ=1 makes the clamp
+# vacuous). Set e.g. 0.15 to cap global-attention density at 15% of
+# tokens in steady state -- this is the safety floor that prevents
+# gate collapse when training without CE. See AHAQwen3Config.
+aha_gate_target = float(os.environ.get("AHA_GATE_TARGET", "1.0"))
+# Direct Duo-style sparsity regularizer. When set to a non-negative value,
+# the AHA loss uses ``AHA_REG_WEIGHT * mean(gate_soft)`` instead of the legacy
+# hinge controlled by AHA_LAMBDA/AHA_GATE_TARGET.
+aha_reg_weight = float(os.environ.get("AHA_REG_WEIGHT", "-1.0"))
+# AHA mode: "dynamic" (per-(token, kv_head) MLP gate) or "duo" (per-(layer,
+# kv_head) static scalar, DuoAttention-style). DUO mode ignores GATE_ONLY
+# and instead uses its own ``DUO_ALPHA_ONLY`` freeze path.
+aha_mode = os.environ.get("AHA_MODE", "dynamic")
+aha_router_granularity_request = os.environ.get(
+    "AHA_ROUTER_GRANULARITY", ""
+).strip()
+if aha_router_granularity_request and aha_router_granularity_request not in {
+    "token",
+    "token_kv_head",
+}:
+    raise ValueError(
+        "AHA_ROUTER_GRANULARITY must be 'token' or 'token_kv_head', got "
+        f"{aha_router_granularity_request!r}"
+    )
+# Override aha_local_kind on a loaded ckpt: "sink_recent" or
+# "sliding_window". Useful for ablating sink dependence in dynamic mode
+# from the same hot-start ckpt without retraining the source. Empty
+# string keeps the ckpt's existing config value.
+aha_local_kind_override = os.environ.get("AHA_LOCAL_KIND", "").strip()
+duo_sink_size = int(os.environ.get("DUO_SINK_SIZE", "64"))
+duo_recent_size = int(os.environ.get("DUO_RECENT_SIZE", "256"))
+duo_alpha_init = float(os.environ.get("DUO_ALPHA_INIT", "1.0"))
+duo_alpha_only = parse_bool("DUO_ALPHA_ONLY", default=(aha_mode == "duo"))
+gate_only = parse_bool("GATE_ONLY", default=False)
+grouped_lr = parse_bool("GROUPED_LR", default=False)
+max_seq_length = int(os.environ.get("MAX_SEQ_LENGTH", "8192"))
+min_train_tokens = get_optional_int("MIN_TRAIN_TOKENS")
+max_train_tokens = get_optional_int("MAX_TRAIN_TOKENS")
+max_train_samples = get_optional_int("MAX_TRAIN_SAMPLES")
+learning_rate = float(os.environ.get("LEARNING_RATE", "3e-5"))
+gate_learning_rate = float(os.environ.get("GATE_LEARNING_RATE", str(learning_rate)))
+backbone_learning_rate = float(os.environ.get("BACKBONE_LEARNING_RATE", str(learning_rate)))
+lr_scheduler_type = os.environ.get("LR_SCHEDULER_TYPE", "cosine")
+warmup_ratio = float(os.environ.get("WARMUP_RATIO", "0.03"))
+adam_beta1 = float(os.environ.get("ADAM_BETA1", "0.9"))
+adam_beta2 = float(os.environ.get("ADAM_BETA2", "0.999"))
+weight_decay = float(os.environ.get("WEIGHT_DECAY", "0.0"))
+max_grad_norm = float(os.environ.get("MAX_GRAD_NORM", "1.0"))
+per_device_batch_size = int(os.environ.get("PER_DEVICE_TRAIN_BATCH_SIZE", "2"))
+gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", "16"))
+num_train_epochs = float(os.environ.get("NUM_TRAIN_EPOCHS", "1"))
+max_steps = int(os.environ.get("MAX_STEPS", "-1"))  # -1 disables the cap; positive overrides num_train_epochs
+logging_steps = int(os.environ.get("LOGGING_STEPS", "10"))
+save_steps = int(os.environ.get("SAVE_STEPS", "500"))
+save_total_limit = int(os.environ.get("SAVE_TOTAL_LIMIT", "2"))
+save_strategy = os.environ.get("SAVE_STRATEGY", "steps")
+seed = int(os.environ.get("SEED", "42"))
+resume_from_checkpoint = resolve_resume_checkpoint(output_dir)
+report_to = parse_report_to()
+wandb_run_name = os.environ.get("WANDB_RUN_NAME") or os.environ.get("RUN_NAME") or os.path.basename(output_dir.rstrip("/"))
+fsdp_options = parse_fsdp_options("FSDP_OPTIONS")
+fsdp_transformer_layer_cls = os.environ.get("FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP", "")
+fsdp_activation_checkpointing = parse_bool("FSDP_ACTIVATION_CHECKPOINTING", default=True)
+freeze_embeddings_lm_head = parse_bool("FREEZE_EMBEDDINGS_LM_HEAD", default=True)
+
+
+def main():
+    AHAQwen3Config.register_for_auto_class()
+    AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f"Loading Qwen3 from {model_path} and converting to AHA-Qwen3...")
+    if aha_checkpoint_path:
+        print(f"  aha_checkpoint_path={aha_checkpoint_path}")
+    print(
+        f"  aha_mode={aha_mode}, "
+        f"aha_router_granularity={aha_router_granularity_request or 'checkpoint/default'}, "
+        f"aha_window_size={aha_window_size}, "
+        f"aha_lambda={aha_lambda}, aha_distill_weight={aha_distill_weight}, "
+        f"aha_ce_weight={aha_ce_weight}, aha_gate_target={aha_gate_target}, "
+        f"aha_reg_weight={aha_reg_weight}, "
+        f"gate_only={gate_only}, duo_alpha_only={duo_alpha_only}, grouped_lr={grouped_lr}"
+    )
+    if aha_mode == "duo":
+        print(
+            f"  duo_sink_size={duo_sink_size}, duo_recent_size={duo_recent_size}, "
+            f"duo_alpha_init={duo_alpha_init}"
+        )
+    print(f"  dataset={dataset_path}[{dataset_split}]")
+    print(f"  output_dir={output_dir}")
+    print(f"  max_seq_length={max_seq_length}")
+    print(
+        "  optimizer/schedule: "
+        f"lr={learning_rate:.3e}, scheduler={lr_scheduler_type}, warmup_ratio={warmup_ratio}, "
+        f"betas=({adam_beta1}, {adam_beta2}), weight_decay={weight_decay}, "
+        f"max_grad_norm={max_grad_norm}"
+    )
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    effective_batch_size = per_device_batch_size * gradient_accumulation_steps * world_size
+    print(
+        "  batch: "
+        f"per_device={per_device_batch_size}, grad_accum={gradient_accumulation_steps}, "
+        f"world_size={world_size}, effective_global={effective_batch_size}"
+    )
+    if min_train_tokens is not None or max_train_tokens is not None:
+        print(f"  token_filter=[{min_train_tokens}, {max_train_tokens}]")
+    if max_train_samples is not None:
+        print(f"  max_train_samples={max_train_samples}")
+    if resume_from_checkpoint is not None:
+        print(f"  resume_from_checkpoint={resume_from_checkpoint}")
+    if fsdp_options:
+        print(f"  fsdp_options={fsdp_options}")
+        if fsdp_transformer_layer_cls:
+            print(f"  fsdp_transformer_layer_cls_to_wrap={fsdp_transformer_layer_cls}")
+    print(f"  report_to={report_to if report_to else 'none'}")
+    if report_to and "wandb" in report_to:
+        print(
+            "  wandb: set WANDB_PROJECT (required for cloud UI); optional WANDB_ENTITY, WANDB_RUN_GROUP, WANDB_TAGS. "
+            "Run `wandb login` once. Metrics include train/ce_loss, train/gate_aux_loss, train/loss_total, train/loss (HF)."
+        )
+
+    if aha_checkpoint_path:
+        model = AHAQwen3ForCausalLM.from_pretrained_aha(
+            aha_checkpoint_path,
+            torch_dtype=torch.bfloat16,
+            attn_implementation="sdpa",
+        )
+        # Config carried inside the checkpoint may predate aha_distill_weight /
+        # aha_ce_weight / aha_gate_target; honor the env override either way.
+        model.config.aha_distill_weight = aha_distill_weight
+        model.config.aha_ce_weight = aha_ce_weight
+        model.config.aha_lambda = aha_lambda
+        model.config.aha_gate_target = aha_gate_target
+        model.config.aha_reg_weight = aha_reg_weight
+        loaded_granularity = getattr(
+            model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
+        )
+        if (
+            aha_router_granularity_request
+            and aha_router_granularity_request != loaded_granularity
+        ):
+            raise ValueError(
+                "AHA_ROUTER_GRANULARITY does not match checkpoint architecture: "
+                f"requested={aha_router_granularity_request!r}, "
+                f"checkpoint={loaded_granularity!r}"
+            )
+        if aha_local_kind_override:
+            if aha_local_kind_override not in {"sink_recent", "sliding_window"}:
+                raise ValueError(
+                    f"AHA_LOCAL_KIND must be 'sink_recent' or 'sliding_window', got {aha_local_kind_override!r}"
+                )
+            prev = getattr(model.config, "aha_local_kind", None)
+            model.config.aha_local_kind = aha_local_kind_override
+            print(f"  AHA_LOCAL_KIND override: {prev!r} -> {aha_local_kind_override!r}")
+    else:
+        model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
+            model_path,
+            aha_window_size=aha_window_size,
+            aha_lambda=aha_lambda,
+            aha_distill_weight=aha_distill_weight,
+            aha_ce_weight=aha_ce_weight,
+            aha_gate_target=aha_gate_target,
+            aha_reg_weight=aha_reg_weight,
+            aha_mode=aha_mode,
+            aha_router_granularity=(
+                aha_router_granularity_request or AHA_ROUTER_GRANULARITY
+            ),
+            duo_sink_size=duo_sink_size,
+            duo_recent_size=duo_recent_size,
+            duo_alpha_init=duo_alpha_init,
+            torch_dtype=torch.bfloat16,
+            attn_implementation="sdpa",
+        )
+
+    # Historical Qwen experiments froze embeddings and lm_head. The rebuttal
+    # scaling protocol disables this to match the OLMo-2 end-to-end recipe,
+    # where the complete pretrained model and routers are fine-tuned jointly.
+    if freeze_embeddings_lm_head:
+        for param in model.model.embed_tokens.parameters():
+            param.requires_grad = False
+        for param in model.lm_head.parameters():
+            param.requires_grad = False
+    print(f"  freeze_embeddings_lm_head={freeze_embeddings_lm_head}")
+
+    if duo_alpha_only:
+        # DuoAttention-style: only the per-(layer, kv_head) scalar
+        # ``full_attention_heads`` is trainable; everything else is
+        # frozen. For Qwen3-0.6B this yields 224 trainable scalars.
+        if getattr(model.config, "aha_mode", "dynamic") != "duo":
+            raise ValueError(
+                "DUO_ALPHA_ONLY=1 requires AHA_MODE=duo; otherwise there are no alpha params."
+            )
+        for param in model.parameters():
+            param.requires_grad = False
+        unfrozen = 0
+        for layer in model.model.layers:
+            p = layer.self_attn.full_attention_heads
+            p.requires_grad = True
+            unfrozen += p.numel()
+        print(f"  DUO_ALPHA_ONLY: trainable scalars {unfrozen:,}")
+    elif gate_only:
+        # Gate-only training: freeze everything, then unfreeze just the
+        # gate rows of q_proj via gradient masks. q_proj.weight / bias
+        # layout is [num_heads*head_dim Q rows, native router rows].
+        setup = configure_gate_only(model)
+        print(
+            f"  GATE_ONLY: granularity={model.config.aha_router_granularity}, "
+            f"gate_rows/layer={setup.gate_rows}, "
+            f"effective gate parameters {setup.effective_parameter_count:,} "
+            "(Q rows grad-masked to 0)"
+        )
+
+    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    n_total = sum(p.numel() for p in model.parameters())
+    print(f"  Trainable: {n_trainable:,} / {n_total:,} ({n_trainable/n_total:.1%})")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+    print(f"Loading dataset from {dataset_path}...")
+    dataset_dict = datasets.load_from_disk(dataset_path)
+    train_dataset = dataset_dict[dataset_split]
+    print(f"  loaded_rows={len(train_dataset):,}")
+
+    if min_train_tokens is not None:
+        train_dataset = train_dataset.filter(lambda example: example["num_tokens"] >= min_train_tokens)
+        print(f"  rows_after_min_train_tokens={len(train_dataset):,}")
+    if max_train_tokens is not None:
+        train_dataset = train_dataset.filter(lambda example: example["num_tokens"] <= max_train_tokens)
+        print(f"  rows_after_max_train_tokens={len(train_dataset):,}")
+    if max_train_samples is not None:
+        keep = min(max_train_samples, len(train_dataset))
+        train_dataset = train_dataset.shuffle(seed=seed).select(range(keep))
+        print(f"  rows_after_max_train_samples={len(train_dataset):,}")
+    if len(train_dataset) == 0:
+        raise ValueError("Training dataset is empty after filtering.")
+
+    if int(os.environ.get("RANK", "0")) == 0:
+        try:
+            git_commit = subprocess.check_output(
+                ["git", "rev-parse", "HEAD"], text=True
+            ).strip()
+            git_status = subprocess.check_output(
+                ["git", "status", "--short"], text=True
+            ).splitlines()
+        except (OSError, subprocess.CalledProcessError):
+            git_commit, git_status = None, []
+        manifest = {
+            "schema": "aha-qwen3-training-manifest-v1",
+            "git_commit": git_commit,
+            "git_status": git_status,
+            "host": socket.gethostname(),
+            "model_path": str(model_path),
+            "source_aha_checkpoint": str(aha_checkpoint_path) or None,
+            "router_granularity": getattr(
+                model.config, "aha_router_granularity", AHA_ROUTER_GRANULARITY
+            ),
+            "native_gate_rows_per_layer": aha_router_output_size(model.config),
+            "effective_gate_parameters": sum(
+                layer.self_attn.aha_router_outputs
+                * (
+                    layer.self_attn.q_proj.in_features
+                    + (layer.self_attn.q_proj.bias is not None)
+                )
+                for layer in model.model.layers
+            ),
+            "effective_sparsity_denominator": "token x KV-head x layer",
+            "dataset": {
+                "path": str(dataset_path),
+                "split": dataset_split,
+                "rows_after_filtering": len(train_dataset),
+                "max_sequence_length": max_seq_length,
+            },
+            "training": {
+                "gate_only": gate_only,
+                "grouped_lr": grouped_lr,
+                "learning_rate": learning_rate,
+                "gate_learning_rate": gate_learning_rate,
+                "backbone_learning_rate": backbone_learning_rate,
+                "regularizer_weight": aha_reg_weight,
+                "ce_weight": aha_ce_weight,
+                "attention_distill_weight": aha_distill_weight,
+                "train_gate_threshold": float(
+                    os.environ.get("AHA_TRAIN_GATE_HARD_THRESHOLD", "0.5")
+                ),
+                "max_steps": max_steps,
+                "seed": seed,
+                "per_device_batch_size": per_device_batch_size,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "world_size": world_size,
+            },
+        }
+        manifest_path = os.path.join(output_dir, "aha_training_manifest.json")
+        with open(manifest_path, "w") as f:
+            json.dump(manifest, f, indent=2)
+    else:
+        manifest_path = os.path.join(output_dir, "aha_training_manifest.json")
+
+    fsdp_config = None
+    gradient_checkpointing = True
+    gradient_checkpointing_kwargs = {"use_reentrant": False}
+    if fsdp_options:
+        fsdp_config = {}
+        if fsdp_transformer_layer_cls:
+            fsdp_config["transformer_layer_cls_to_wrap"] = [fsdp_transformer_layer_cls]
+        if fsdp_activation_checkpointing:
+            # Transformers warns that FSDP should use fsdp_config.activation_checkpointing
+            # instead of Trainer-level gradient_checkpointing to avoid redundant all-gathers.
+            fsdp_config["activation_checkpointing"] = True
+            gradient_checkpointing = False
+            gradient_checkpointing_kwargs = None
+
+    sft_config = SFTConfig(
+        output_dir=output_dir,
+        per_device_train_batch_size=per_device_batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        learning_rate=learning_rate,
+        num_train_epochs=num_train_epochs,
+        max_steps=max_steps,
+        max_length=max_seq_length,
+        lr_scheduler_type=lr_scheduler_type,
+        warmup_ratio=warmup_ratio,
+        adam_beta1=adam_beta1,
+        adam_beta2=adam_beta2,
+        weight_decay=weight_decay,
+        max_grad_norm=max_grad_norm,
+        logging_steps=logging_steps,
+        save_steps=save_steps,
+        save_total_limit=save_total_limit,
+        save_strategy=save_strategy,
+        bf16=True,
+        tf32=True,
+        gradient_checkpointing=gradient_checkpointing,
+        gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
+        dataloader_drop_last=True,
+        remove_unused_columns=True,
+        report_to=report_to if report_to else "none",
+        run_name=wandb_run_name,
+        seed=seed,
+        # The custom model returns a standard local mean CE and does not use
+        # ``num_items_in_batch``. Keep legacy OLMo-style per-microbatch means
+        # instead of Transformers' global token-count scaling path.
+        average_tokens_across_devices=False,
+        fsdp=fsdp_options,
+        fsdp_config=fsdp_config,
+    )
+
+    trainer = AHASFTTrainer(
+        model=model,
+        args=sft_config,
+        train_dataset=train_dataset,
+        processing_class=tokenizer,
+        optimizers=(build_grouped_lr_optimizer(model), None) if grouped_lr else (None, None),
+        callbacks=[AHAManifestCallback(manifest_path)],
+    )
+
+    print(
+        "Starting training… Extra log fields: ce_loss, gate_aux_loss, loss_total (=ce+gate_aux), "
+        "gate_soft_mean, gate_hard_mean, hf_loss_ratio (= key `loss` / loss_total). "
+        "ce_loss / loss_total are raw model means (language quality + AHA aux). "
+        "Key `loss` is HuggingFace Trainer’s gradient-accumulation-normalized aggregate; "
+        "hf_loss_ratio should stay near 1.0."
+    )
+    print("Starting training...")
+    train_result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+    trainer.save_model(output_dir)
+    trainer.save_state()
+
+    metrics = train_result.metrics
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+    print(f"Training completed. Metrics: {metrics}")
+
+
+if __name__ == "__main__":
+    main()

recipe/source_training/all_results.json

@@ -0,0 +1,9 @@
+{
+    "ce_loss": 0.8887,
+    "total_flos": 6.189458802853544e+18,
+    "train_loss": 0.9857371548739723,
+    "train_ppl": 2.431889,
+    "train_runtime": 23521.9665,
+    "train_samples_per_second": 37.607,
+    "train_steps_per_second": 0.294
+}

recipe/source_training/train_results.json

@@ -0,0 +1,9 @@
+{
+    "ce_loss": 0.8887,
+    "total_flos": 6.189458802853544e+18,
+    "train_loss": 0.9857371548739723,
+    "train_ppl": 2.431889,
+    "train_runtime": 23521.9665,
+    "train_samples_per_second": 37.607,
+    "train_steps_per_second": 0.294
+}

recipe/tests/__init__.py

@@ -0,0 +1 @@
+"""Release-package tests."""

recipe/tests/test_router_granularity.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+import torch
+from transformers import Qwen3Config, Qwen3ForCausalLM
+
+import modeling_aha_qwen3 as aha_module
+from modeling_aha_qwen3 import (
+    AHAQwen3Config,
+    AHAQwen3ForCausalLM,
+    aha_router_output_size,
+)
+from router_training_utils import RowWiseAdamW, configure_gate_only
+
+
+def base_config() -> Qwen3Config:
+    config = Qwen3Config(
+        vocab_size=97,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        max_position_embeddings=64,
+        attention_dropout=0.0,
+        attention_bias=True,
+        tie_word_embeddings=False,
+    )
+    config._attn_implementation = "eager"
+    return config
+
+
+def aha_config(granularity: str, *, force_gate_value=None) -> AHAQwen3Config:
+    payload = base_config().to_dict()
+    payload.update(
+        aha_window_size=2,
+        aha_local_kind="sliding_window",
+        aha_mode="dynamic",
+        aha_router_granularity=granularity,
+        aha_force_gate_value=force_gate_value,
+        aha_reg_weight=1.0,
+        aha_ce_weight=0.0,
+        model_type="aha_qwen3",
+    )
+    config = AHAQwen3Config(**payload)
+    config._attn_implementation = "eager"
+    return config
+
+
+def copy_base_weights(base: Qwen3ForCausalLM, target: AHAQwen3ForCausalLM) -> None:
+    source = base.state_dict()
+    destination = target.state_dict()
+    q_rows = target.config.num_attention_heads * target.config.head_dim
+    with torch.no_grad():
+        for name, value in source.items():
+            if name not in destination:
+                continue
+            if destination[name].shape == value.shape:
+                destination[name].copy_(value)
+            elif name.endswith("self_attn.q_proj.weight"):
+                destination[name][:q_rows].copy_(value)
+                destination[name][q_rows:].zero_()
+            elif name.endswith("self_attn.q_proj.bias"):
+                destination[name][:q_rows].copy_(value)
+                destination[name][q_rows:].zero_()
+            else:
+                raise AssertionError(f"unexpected shape mismatch for {name}")
+    target.load_state_dict(destination)
+
+
+def aligned_models():
+    torch.manual_seed(7)
+    base = Qwen3ForCausalLM(base_config()).eval()
+    models = {}
+    for granularity in ("token", "token_kv_head"):
+        model = AHAQwen3ForCausalLM(
+            aha_config(granularity, force_gate_value=1.0)
+        ).eval()
+        copy_base_weights(base, model)
+        models[granularity] = model
+    return base, models
+
+
+class RouterGranularityTest(unittest.TestCase):
+    def test_auto_class_checkpoint_embeds_modeling_source(self):
+        AHAQwen3Config.register_for_auto_class()
+        AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+        with tempfile.TemporaryDirectory() as tmp:
+            AHAQwen3ForCausalLM(aha_config("token")).save_pretrained(
+                tmp, safe_serialization=True
+            )
+            self.assertTrue((Path(tmp) / "modeling_aha_qwen3.py").exists())
+
+    def test_projection_and_effective_gate_shapes(self):
+        _, models = aligned_models()
+        input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        expected_q_rows = 4 * 8
+        for granularity, native_rows in (("token", 1), ("token_kv_head", 2)):
+            with self.subTest(granularity=granularity):
+                model = models[granularity]
+                attention = model.model.layers[0].self_attn
+                self.assertEqual(attention.q_proj.out_features, expected_q_rows + native_rows)
+                self.assertEqual(attention.aha_router_outputs, native_rows)
+                output = model.model(input_ids=input_ids, use_cache=False)
+                self.assertEqual(len(output.all_gate_soft), 2)
+                self.assertEqual(output.all_gate_soft[0].shape, (1, 5, 2))
+                self.assertEqual(output.all_gate_hard[0].shape, (1, 5, 2))
+
+    def test_force_open_matches_full_attention_logits(self):
+        base, models = aligned_models()
+        input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
+        with torch.no_grad():
+            expected = base(input_ids=input_ids, use_cache=False).logits
+            for granularity, model in models.items():
+                with self.subTest(granularity=granularity):
+                    actual = model(input_ids=input_ids, use_cache=False).logits
+                    torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)
+
+    def test_force_closed_uses_identical_local_branch(self):
+        _, models = aligned_models()
+        input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
+        for model in models.values():
+            model.config.aha_force_gate_value = 0.0
+        with torch.no_grad():
+            token_logits = models["token"](input_ids=input_ids, use_cache=False).logits
+            head_logits = models["token_kv_head"](
+                input_ids=input_ids, use_cache=False
+            ).logits
+        torch.testing.assert_close(token_logits, head_logits, atol=1e-6, rtol=1e-5)
+
+    def test_regularizer_uses_effective_kv_head_denominator(self):
+        _, models = aligned_models()
+        probability = 0.73
+        bias = math.log(probability / (1.0 - probability))
+        input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
+        outputs = {}
+        for granularity, model in models.items():
+            model.train()
+            model.config.aha_force_gate_value = None
+            q_rows = model.config.num_attention_heads * model.config.head_dim
+            with torch.no_grad():
+                for layer in model.model.layers:
+                    layer.self_attn.q_proj.weight[q_rows:].zero_()
+                    layer.self_attn.q_proj.bias[q_rows:].fill_(bias)
+            outputs[granularity] = model(
+                input_ids=input_ids, labels=input_ids.clone(), use_cache=False
+            )
+        torch.testing.assert_close(
+            outputs["token"].gate_soft_mean,
+            outputs["token_kv_head"].gate_soft_mean,
+        )
+        torch.testing.assert_close(
+            outputs["token"].gate_aux_loss,
+            outputs["token_kv_head"].gate_aux_loss,
+        )
+        self.assertAlmostEqual(outputs["token"].gate_soft_mean.item(), probability, places=6)
+
+    def test_gate_only_updates_only_appended_rows(self):
+        for granularity, expected_rows in (("token", 1), ("token_kv_head", 2)):
+            with self.subTest(granularity=granularity):
+                model = AHAQwen3ForCausalLM(aha_config(granularity))
+                setup = configure_gate_only(model)
+                self.assertEqual(setup.gate_rows, expected_rows)
+                q_proj = model.model.layers[0].self_attn.q_proj
+                before = q_proj.weight.detach().clone()
+                q_proj.weight.sum().backward()
+                self.assertEqual(
+                    torch.count_nonzero(q_proj.weight.grad[: setup.q_rows]).item(), 0
+                )
+                self.assertGreater(
+                    torch.count_nonzero(q_proj.weight.grad[setup.q_rows :]).item(), 0
+                )
+                torch.optim.SGD(setup.parameters, lr=0.1).step()
+                torch.testing.assert_close(
+                    q_proj.weight[: setup.q_rows], before[: setup.q_rows]
+                )
+                self.assertFalse(
+                    torch.equal(q_proj.weight[setup.q_rows :], before[setup.q_rows :])
+                )
+
+    def test_rowwise_adamw_applies_real_ten_x_lr_ratio(self):
+        parameter = torch.nn.Parameter(torch.zeros(3, 1))
+        optimizer = RowWiseAdamW(
+            [{"params": [parameter], "lr": 0.1}],
+            row_scales=[(parameter, 2, 0.1)],
+            weight_decay=0.0,
+            betas=(0.9, 0.999),
+        )
+        parameter.grad = torch.ones_like(parameter)
+        optimizer.step()
+        backbone_update = parameter[:2].abs().mean().item()
+        gate_update = parameter[2:].abs().mean().item()
+        self.assertAlmostEqual(gate_update / backbone_update, 10.0, places=5)
+
+    def test_one_step_smoke_is_finite_and_reloadable(self):
+        input_ids = torch.tensor([[1, 4, 2, 8, 3, 9]])
+        for granularity in ("token", "token_kv_head"):
+            with self.subTest(granularity=granularity), tempfile.TemporaryDirectory() as tmp:
+                model = AHAQwen3ForCausalLM(aha_config(granularity)).train()
+                setup = configure_gate_only(model)
+                optimizer = torch.optim.AdamW(setup.parameters, lr=3e-5)
+                output = model(
+                    input_ids=input_ids, labels=input_ids.clone(), use_cache=False
+                )
+                self.assertTrue(torch.isfinite(output.loss).item())
+                output.loss.backward()
+                self.assertTrue(
+                    all(
+                        parameter.grad is None
+                        or torch.isfinite(parameter.grad).all().item()
+                        for parameter in setup.parameters
+                    )
+                )
+                optimizer.step()
+                model.save_pretrained(tmp, safe_serialization=True)
+                reloaded = AHAQwen3ForCausalLM.from_pretrained_aha(
+                    tmp, torch_dtype=torch.float32, attn_implementation="eager"
+                )
+                self.assertEqual(
+                    reloaded.config.aha_router_granularity, granularity
+                )
+
+    def test_legacy_default_and_save_load_roundtrip(self):
+        legacy = AHAQwen3Config(**base_config().to_dict())
+        self.assertEqual(legacy.aha_router_granularity, "token_kv_head")
+        self.assertEqual(aha_router_output_size(legacy), legacy.num_key_value_heads)
+
+        with tempfile.TemporaryDirectory() as tmp:
+            model = AHAQwen3ForCausalLM(aha_config("token"))
+            model.save_pretrained(tmp, safe_serialization=True)
+            loaded = AHAQwen3ForCausalLM.from_pretrained_aha(
+                tmp, torch_dtype=torch.float32, attn_implementation="eager"
+            )
+            self.assertEqual(loaded.config.aha_router_granularity, "token")
+            self.assertEqual(loaded.model.layers[0].self_attn.aha_router_outputs, 1)
+
+    def test_sparsity_tracker_records_native_and_effective_counts(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            output = os.path.join(tmp, "sparsity.json")
+            with mock.patch.dict(os.environ, {"AHA_SPARSITY_STATS_PATH": output}):
+                tracker = aha_module._AHAInferenceSparsityTracker()
+                gate_hard = torch.tensor([[[1.0, 1.0], [0.0, 0.0]]])
+                gate_soft = torch.tensor([[[0.8, 0.8], [0.2, 0.2]]])
+                tracker.update(
+                    gate_hard,
+                    gate_soft,
+                    layer_idx=0,
+                    phase="prefill",
+                    router_granularity="token",
+                    native_router_width=1,
+                )
+                tracker.update(
+                    gate_hard[:, :1],
+                    gate_soft[:, :1],
+                    layer_idx=0,
+                    phase="decode",
+                    router_granularity="token",
+                    native_router_width=1,
+                )
+                tracker.write_stats()
+                payload = json.loads(Path(output).read_text())
+            self.assertEqual(payload["router_granularity"], "token")
+            self.assertEqual(payload["native_router_decisions"], 3)
+            self.assertEqual(payload["effective_router_decisions"], 6)
+            self.assertAlmostEqual(payload["sparsity"], 1 / 3)
+            self.assertEqual(payload["by_phase"]["decode"]["total_decisions"], 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}