Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* 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/.pytest_cache/.gitignore

@@ -0,0 +1,2 @@
+# Created by pytest automatically.
+*

recipe/.pytest_cache/CACHEDIR.TAG

@@ -0,0 +1,4 @@
+Signature: 8a477f597d28d172789f06886806bc55
+# This file is a cache directory tag created by pytest.
+# For information about cache directory tags, see:
+#	https://bford.info/cachedir/spec.html

recipe/.pytest_cache/README.md

@@ -0,0 +1,8 @@
+# pytest cache directory #
+
+This directory contains data from the pytest's cache plugin,
+which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
+
+**Do not** commit this to version control.
+
+See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

recipe/.pytest_cache/v/cache/nodeids

@@ -0,0 +1,12 @@
+[
+  "tests/test_router_granularity.py::RouterGranularityTest::test_auto_class_checkpoint_embeds_modeling_source",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_force_closed_uses_identical_local_branch",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_force_open_matches_full_attention_logits",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_gate_only_updates_only_appended_rows",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_legacy_default_and_save_load_roundtrip",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_one_step_smoke_is_finite_and_reloadable",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_projection_and_effective_gate_shapes",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_regularizer_uses_effective_kv_head_denominator",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_rowwise_adamw_applies_real_ten_x_lr_ratio",
+  "tests/test_router_granularity.py::RouterGranularityTest::test_sparsity_tracker_records_native_and_effective_counts"
+]

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/modeling_aha_qwen3.py

@@ -0,0 +1,1226 @@
+"""
+AHA-Qwen3: All-or-Here Attention applied to Qwen3.
+
+Minimal modification to standard Qwen3 attention: adds a binary gate that
+dynamically toggles between full attention and local attention. The persisted
+``aha_router_granularity`` config selects either one gate per (token, KV-head)
+or one gate per token shared by every head in the layer.
+
+Based on: transformers.models.qwen3.modeling_qwen3 (Qwen3Attention)
+Changes:
+  1. q_proj outputs either one extra gate value (``token``) or
+     ``num_key_value_heads`` values (``token_kv_head``). The shared-token
+     value is broadcast to all KV/query heads before attention mixing.
+  2. Gate: sigmoid → hard threshold → STE for training
+  3. Double-mix: global attention + local window attention, blended by
+     gate (broadcast across the GQA group)
+  4. Auxiliary loss: λ * mean(gate_soft) to encourage sparsity
+"""
+
+import atexit
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file
+
+from transformers import Qwen3Config, Qwen3ForCausalLM
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Attention,
+    Qwen3DecoderLayer,
+    Qwen3Model,
+    Qwen3RMSNorm,
+    apply_rotary_pos_emb,
+)
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.processing_utils import Unpack
+from transformers.utils.generic import TransformersKwargs
+
+# AHA constants
+AHA_WINDOW_SIZE = 1024
+AHA_LAMBDA = 3e-4
+AHA_DISTILL_WEIGHT = 0.0
+AHA_CE_WEIGHT = 1.0
+# Gate-density ceiling (hinge target). The aux loss only fires when
+# ``mean(gate_soft) > AHA_GATE_TARGET``, so sparsity cannot drop
+# below ``1 - AHA_GATE_TARGET`` in steady state regardless of how
+# large ``aha_lambda`` is. Setting this to 1.0 recovers the legacy
+# behaviour (unconditional downward pressure).
+AHA_GATE_TARGET = 1.0
+
+# DuoAttention baseline knobs. When ``aha_mode == "duo"`` the per-(token,
+# kv_head) gate is replaced by a per-(layer, kv_head) learnable scalar
+# (``duo_alpha``, 224 scalars total for Qwen3-0.6B), and the "local"
+# branch becomes a streaming mask with ``duo_sink_size`` attention sinks
+# plus a recent window of ``duo_recent_size``. This is a strict reduction
+# of the dynamic AHA gate and reproduces the DuoAttention paper's head
+# classifier (Xiao et al. 2024) on the Qwen3 backbone.
+AHA_MODE = "dynamic"
+DUO_SINK_SIZE = 64
+DUO_RECENT_SIZE = 256
+DUO_ALPHA_INIT = 1.0
+# Local-branch mask kind for the dual-branch attention (global vs local).
+#  - "sliding_window": last ``aha_window_size`` tokens, no attention sinks.
+#    Historical default for ``aha_mode="dynamic"`` (Bv3 was trained with this).
+#  - "sink_recent":    first ``duo_sink_size`` tokens + last ``duo_recent_size``
+#    tokens, matching the DuoAttention paper. Historical hardcoded default for
+#    ``aha_mode="duo"`` (still forced for duo regardless of this field).
+# § 9.4.13.3 found that hot-starting a dynamic-mode ckpt from a duo ckpt with
+# sliding_window local loses ~20pp LongBench retention vs the duo source,
+# almost entirely from passage_retrieval_en (sinks are critical for long
+# passage retrieval). Setting this to "sink_recent" for hot-started dynamic
+# ckpts closes that architectural gap.
+AHA_LOCAL_KIND = "sliding_window"
+AHA_ROUTER_GRANULARITY = "token_kv_head"
+
+
+def aha_router_output_size(config: Qwen3Config) -> int:
+    """Return the number of native dynamic-router logits per token."""
+
+    granularity = getattr(config, "aha_router_granularity", AHA_ROUTER_GRANULARITY)
+    if granularity == "token":
+        return 1
+    if granularity == "token_kv_head":
+        return int(config.num_key_value_heads)
+    raise ValueError(f"Unsupported aha_router_granularity={granularity!r}")
+
+
+class _AHAInferenceSparsityTracker:
+    """Accumulate exact inference hard-route counts without changing outputs."""
+
+    def __init__(self) -> None:
+        # Keep counters as device scalars until process exit. Calling .item()
+        # once per layer and decode token serialized CUDA and made the
+        # diagnostic tracker materially slow generation.
+        self.sparse_decisions: dict[int, torch.Tensor] = {}
+        self.total_decisions: dict[int, int] = {}
+        self.phase_sparse_decisions: dict[str, torch.Tensor] = {}
+        self.phase_total_decisions: dict[str, int] = {}
+        raw_thresholds = os.environ.get("AHA_COUNTERFACTUAL_THRESHOLDS", "")
+        self.counterfactual_thresholds = tuple(
+            sorted({float(value) for value in raw_thresholds.split(",") if value.strip()})
+        )
+        if any(value < 0.0 or value > 1.0 for value in self.counterfactual_thresholds):
+            raise ValueError("AHA_COUNTERFACTUAL_THRESHOLDS values must be in [0, 1]")
+        self.counterfactual_sparse: dict[float, torch.Tensor] = {}
+        self.router_granularity: Optional[str] = None
+        self.native_router_decisions = 0
+        atexit.register(self.write_stats)
+
+    def update(
+        self,
+        gate_hard: torch.Tensor,
+        gate_soft: torch.Tensor,
+        layer_idx: int,
+        phase: str,
+        router_granularity: str,
+        native_router_width: int,
+    ) -> None:
+        if self.router_granularity is None:
+            self.router_granularity = router_granularity
+        elif self.router_granularity != router_granularity:
+            raise RuntimeError(
+                "AHA sparsity tracker received mixed router granularities: "
+                f"{self.router_granularity!r} and {router_granularity!r}"
+            )
+        sparse = torch.count_nonzero(gate_hard == 0).to(torch.int64)
+        total = int(gate_hard.numel())
+        effective_width = int(gate_hard.shape[-1])
+        self.native_router_decisions += total // effective_width * native_router_width
+        if layer_idx not in self.sparse_decisions:
+            self.sparse_decisions[layer_idx] = sparse
+        else:
+            self.sparse_decisions[layer_idx].add_(sparse)
+        self.total_decisions[layer_idx] = self.total_decisions.get(layer_idx, 0) + total
+        if phase not in self.phase_sparse_decisions:
+            self.phase_sparse_decisions[phase] = sparse.clone()
+        else:
+            self.phase_sparse_decisions[phase].add_(sparse)
+        self.phase_total_decisions[phase] = self.phase_total_decisions.get(phase, 0) + total
+        if self.counterfactual_thresholds:
+            flat = gate_soft.detach().reshape(-1, 1)
+            thresholds = flat.new_tensor(self.counterfactual_thresholds).reshape(1, -1)
+            counts = (flat <= thresholds).sum(dim=0, dtype=torch.int64)
+            for threshold, count in zip(self.counterfactual_thresholds, counts.unbind()):
+                if threshold not in self.counterfactual_sparse:
+                    self.counterfactual_sparse[threshold] = count
+                else:
+                    self.counterfactual_sparse[threshold].add_(count)
+
+    def write_stats(self) -> None:
+        output = os.environ.get("AHA_SPARSITY_STATS_PATH")
+        if not output:
+            return
+        sparse_by_layer = {
+            layer_idx: int(value.item())
+            for layer_idx, value in self.sparse_decisions.items()
+        }
+        sparse = sum(sparse_by_layer.values())
+        total = sum(self.total_decisions.values())
+        per_layer = {}
+        for layer_idx in sorted(self.total_decisions):
+            layer_sparse = sparse_by_layer.get(layer_idx, 0)
+            layer_total = self.total_decisions[layer_idx]
+            per_layer[str(layer_idx)] = {
+                "sparse_decisions": layer_sparse,
+                "total_decisions": layer_total,
+                "sparsity": layer_sparse / layer_total if layer_total else None,
+            }
+        payload = {
+            "definition": (
+                "hard AHA gate zeros / all token x KV-head x layer decisions, "
+                "token-weighted over prefill and decode"
+            ),
+            "router_granularity": self.router_granularity,
+            "native_router_decisions": self.native_router_decisions,
+            "effective_router_decisions": total,
+            "native_router_note": (
+                "Native decisions count learned gate logits before broadcast; effective "
+                "decisions always use token x KV-head x layer for matched sparsity."
+            ),
+            "sparse_decisions": sparse,
+            "total_decisions": total,
+            "sparsity": sparse / total if total else None,
+            "full_attention_usage": 1.0 - sparse / total if total else None,
+            "per_layer": per_layer,
+            "by_phase": {
+                phase: {
+                    "sparse_decisions": int(self.phase_sparse_decisions[phase].item()),
+                    "total_decisions": phase_total,
+                    "sparsity": int(self.phase_sparse_decisions[phase].item()) / phase_total
+                    if phase_total else None,
+                    "full_attention_usage": 1.0
+                    - int(self.phase_sparse_decisions[phase].item()) / phase_total
+                    if phase_total else None,
+                }
+                for phase, phase_total in sorted(self.phase_total_decisions.items())
+            },
+        }
+        if self.counterfactual_thresholds:
+            payload["counterfactual_sparsity_by_threshold"] = {
+                f"{threshold:.6g}": {
+                    "sparse_decisions": int(self.counterfactual_sparse[threshold].item()),
+                    "total_decisions": total,
+                    "sparsity": int(self.counterfactual_sparse[threshold].item()) / total
+                    if total else None,
+                }
+                for threshold in self.counterfactual_thresholds
+            }
+            payload["counterfactual_note"] = (
+                "Routing counts only; benchmark quality must be measured in a separate "
+                "run with AHA_GATE_HARD_THRESHOLD set to the selected threshold."
+            )
+        path = Path(output)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(payload, indent=2) + "\n")
+
+
+_AHA_INFERENCE_SPARSITY_TRACKER: Optional[_AHAInferenceSparsityTracker] = None
+
+
+def _track_aha_inference_sparsity(
+    gate_hard: torch.Tensor,
+    gate_soft: torch.Tensor,
+    layer_idx: int,
+    phase: str,
+    router_granularity: str,
+    native_router_width: int,
+) -> None:
+    global _AHA_INFERENCE_SPARSITY_TRACKER
+    if not os.environ.get("AHA_SPARSITY_STATS_PATH"):
+        return
+    if _AHA_INFERENCE_SPARSITY_TRACKER is None:
+        _AHA_INFERENCE_SPARSITY_TRACKER = _AHAInferenceSparsityTracker()
+    _AHA_INFERENCE_SPARSITY_TRACKER.update(
+        gate_hard.detach(), gate_soft.detach(), layer_idx, phase,
+        router_granularity, native_router_width,
+    )
+
+
+def _get_optional_nonnegative_int_env(name: str) -> Optional[int]:
+    raw = os.environ.get(name, "").strip()
+    if not raw:
+        return None
+    value = int(raw)
+    if value < 0:
+        raise ValueError(f"{name} must be non-negative, got {value}")
+    return value
+
+
+def _apply_low_alpha_force_full_heads_from_env(model: nn.Module) -> None:
+    """Auto-protect the weakest static-full heads in duo_dynamic mode.
+
+    ``AHA_FORCE_LOW_ALPHA_FULL_HEADS=K`` selects the K smallest
+    ``full_attention_heads`` values among heads that static Duo still kept
+    full (alpha > 0.5). These are the fragile boundary full heads: raising the
+    dynamic threshold can close them, but static Duo's alpha says they should
+    not be fully discarded. The selected heads are OR-ed with any manually
+    supplied ``AHA_FORCE_FULL_HEADS`` mask.
+    """
+    k = _get_optional_nonnegative_int_env("AHA_FORCE_LOW_ALPHA_FULL_HEADS")
+    if not k:
+        return
+    if getattr(model.config, "aha_mode", "dynamic") != "duo_dynamic":
+        return
+
+    candidates: list[tuple[float, int, int]] = []
+    with torch.no_grad():
+        for layer_idx, layer in enumerate(model.model.layers):
+            attn = layer.self_attn
+            alpha = getattr(attn, "full_attention_heads", None)
+            if alpha is None:
+                continue
+            for head_idx, value in enumerate(alpha.detach().float().cpu().tolist()):
+                if value > 0.5:
+                    candidates.append((float(value), layer_idx, head_idx))
+    selected = sorted(candidates)[: min(k, len(candidates))]
+    if not selected:
+        return
+
+    selected_by_layer: dict[int, list[int]] = {}
+    for _, layer_idx, head_idx in selected:
+        selected_by_layer.setdefault(layer_idx, []).append(head_idx)
+    with torch.no_grad():
+        for layer_idx, heads in selected_by_layer.items():
+            mask = model.model.layers[layer_idx].self_attn._aha_force_full_heads_mask
+            for head_idx in heads:
+                mask[head_idx] = True
+    model.config.aha_force_low_alpha_full_heads = int(k)
+    formatted = ",".join(f"{layer}:{head}" for _, layer, head in selected)
+    print(
+        f"[AHA] AHA_FORCE_LOW_ALPHA_FULL_HEADS={k} selected {len(selected)} "
+        f"duo_dynamic heads: {formatted}",
+        flush=True,
+    )
+
+
+class AHAQwen3Config(Qwen3Config):
+    """Qwen3Config with AHA-specific parameters."""
+    model_type = "aha_qwen3"
+
+    def __init__(
+        self,
+        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=-1.0,
+        aha_force_gate_value=None,
+        aha_mode=AHA_MODE,
+        duo_sink_size=DUO_SINK_SIZE,
+        duo_recent_size=DUO_RECENT_SIZE,
+        duo_alpha_init=DUO_ALPHA_INIT,
+        aha_local_kind=AHA_LOCAL_KIND,
+        aha_router_granularity=AHA_ROUTER_GRANULARITY,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.aha_window_size = aha_window_size
+        self.aha_lambda = aha_lambda
+        # Per-layer attention distillation weight. When > 0 and model is
+        # training, each layer contributes ``mean(((1-g)·(global-local))**2)``
+        # to the total loss. This is a dense per-(token, head) signal that
+        # pushes the gate toward 1 where SWA is insufficient to approximate
+        # full attention. Cost: ~0 extra forward compute (global/local are
+        # already computed for the mix).
+        self.aha_distill_weight = aha_distill_weight
+        # Language-modeling CE weight for the training objective. Default 1.0
+        # keeps the standard SFT behaviour. Set to 0.0 to drop CE entirely
+        # and train the gate purely from ``aux + distill`` (useful when the
+        # backbone is frozen and we only want to shape the gate). Evaluation
+        # always reports unweighted ``ce_loss`` regardless of this setting.
+        self.aha_ce_weight = aha_ce_weight
+        # Direct Duo-style sparsity regularizer. When >= 0, this replaces the
+        # legacy hinge term and uses ``aha_reg_weight * mean(gate_soft)``.
+        # Default -1 keeps older checkpoints on their original hinge objective.
+        self.aha_reg_weight = float(aha_reg_weight)
+        # Optional override used by Duo-style dynamic distillation. When set
+        # to 1.0, dynamic mode behaves as a full-attention teacher for one
+        # forward; None keeps the learned gate.
+        self.aha_force_gate_value = aha_force_gate_value
+        # Target ceiling on mean(gate_soft). Aux loss is hinge-shaped:
+        # ``λ · max(0, mean(gate_soft) - τ)``. Setting τ < 1.0 guarantees
+        # the gate cannot drop below ``1 - τ`` local-attention fraction
+        # in steady state -- the failure mode we observed with the legacy
+        # unconditional aux, where ``λ = 1.0`` drove mean(gate_hard) to
+        # 0.005 (99.5% local) and killed long-context accuracy
+        # (gsm8k 74% -> 19%). Default 1.0 is a no-op for backward
+        # compatibility with checkpoints trained on the old loss.
+        self.aha_gate_target = aha_gate_target
+        # DuoAttention-style baseline. ``dynamic`` keeps the original per-
+        # (token, kv_head) MLP gate. ``duo`` replaces it with a per-(layer,
+        # kv_head) learnable scalar and swaps the local branch's mask for a
+        # streaming (sink + recent) pattern. ``duo_dynamic`` keeps the Duo
+        # static mask for KV-capacity planning, but applies a dynamic AHA gate
+        # only on the static-full heads: effective_full = duo_full AND dyn_gate.
+        # Sink / recent sizes are fixed at train time; set ``duo_sink_size=0``
+        # for the sink-ablation run.
+        if aha_mode not in ("dynamic", "duo", "duo_dynamic"):
+            raise ValueError(
+                f"aha_mode must be 'dynamic', 'duo', or 'duo_dynamic', got {aha_mode!r}"
+            )
+        self.aha_mode = aha_mode
+        if aha_router_granularity not in ("token", "token_kv_head"):
+            raise ValueError(
+                "aha_router_granularity must be 'token' or 'token_kv_head', "
+                f"got {aha_router_granularity!r}"
+            )
+        # Persist this even for static Duo checkpoints. Dynamic AHA uses it to
+        # choose the native q_proj gate-row count; old configs omit the field
+        # and therefore retain the historical token_kv_head behavior.
+        self.aha_router_granularity = aha_router_granularity
+        self.duo_sink_size = int(duo_sink_size)
+        self.duo_recent_size = int(duo_recent_size)
+        self.duo_alpha_init = float(duo_alpha_init)
+        if aha_local_kind not in ("sliding_window", "sink_recent"):
+            raise ValueError(
+                f"aha_local_kind must be 'sliding_window' or 'sink_recent', "
+                f"got {aha_local_kind!r}"
+            )
+        # When aha_mode="duo" the local branch is always sink+recent regardless
+        # of this field (forced to the historical DuoAttention behaviour). The
+        # field only controls the dynamic mode's local mask.
+        self.aha_local_kind = aha_local_kind
+
+
+def _build_streaming_causal_mask(
+    seq_len_q: int,
+    seq_len_kv: int,
+    sink_size: int,
+    recent_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    kv_offset: int = 0,
+) -> torch.Tensor:
+    """Construct a streaming attention mask (sink + recent, causal).
+
+    Returns an additive mask of shape ``[1, 1, Q, K]`` where
+    disallowed positions are ``finfo(dtype).min`` and allowed positions
+    are ``0``. A query at position ``q`` (absolute in the sequence,
+    with ``kv_offset`` tokens already in the KV cache) attends to a key
+    at position ``k`` iff ``k <= q_abs`` AND (``k < sink_size`` OR
+    ``q_abs - k < recent_size``).
+    """
+    q_abs = torch.arange(seq_len_q, device=device).unsqueeze(-1) + kv_offset  # [Q, 1]
+    k_abs = torch.arange(seq_len_kv, device=device).unsqueeze(0)               # [1, K]
+    causal = k_abs <= q_abs
+    sink = k_abs < sink_size if sink_size > 0 else torch.zeros_like(causal)
+    recent = (q_abs - k_abs) < recent_size
+    allowed = causal & (sink | recent)
+    mask = torch.zeros(seq_len_q, seq_len_kv, dtype=dtype, device=device)
+    mask.masked_fill_(~allowed, torch.finfo(dtype).min)
+    return mask.unsqueeze(0).unsqueeze(0)
+
+
+class AHAQwen3Attention(nn.Module):
+    """Qwen3 attention with AHA gate for per-head local/global routing."""
+
+    def __init__(self, config: AHAQwen3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim ** -0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.aha_window_size = config.aha_window_size
+
+        # q_proj: extra router outputs for either one shared per-token gate or
+        # one per-(token, KV-head) gate.
+        # DUO mode uses a per-(layer, kv_head) static scalar instead, so
+        # q_proj keeps its vanilla Qwen3 shape and is binary-compatible
+        # with the base model checkpoint.
+        self._has_dyn_gate_logits = getattr(config, "aha_mode", "dynamic") in (
+            "dynamic",
+            "duo_dynamic",
+        )
+        self.aha_router_outputs = aha_router_output_size(config)
+        q_out_dim = config.num_attention_heads * self.head_dim
+        if self._has_dyn_gate_logits:
+            q_out_dim += self.aha_router_outputs
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            q_out_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        # Qwen3 applies q_norm/k_norm per head_dim (not full projection)
+        self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+        # DuoAttention baseline: per-(layer, kv_head) learnable scalar.
+        # Populated in ``duo`` and ``duo_dynamic``. In ``duo_dynamic`` this
+        # scalar is the frozen Duo capacity mask; the dynamic gate is only
+        # allowed to close heads where this scalar says "full". Registered
+        # as ``full_attention_heads`` to match the name used in the
+        # official DuoAttention repo (``duo_attn.patch.llama``), so
+        # per-head statistics (e.g. retrieval heads vs streaming heads)
+        # can be read back with the same tooling.
+        if getattr(config, "aha_mode", "dynamic") in ("duo", "duo_dynamic"):
+            init = float(getattr(config, "duo_alpha_init", 1.0))
+            self.full_attention_heads = nn.Parameter(
+                torch.full((config.num_key_value_heads,), init, dtype=torch.float32)
+            )
+        else:
+            self.full_attention_heads = None
+        force_full_mask = torch.zeros(config.num_key_value_heads, dtype=torch.bool)
+        for raw in os.environ.get("AHA_FORCE_FULL_HEADS", "").split(","):
+            raw = raw.strip()
+            if not raw:
+                continue
+            layer_s, head_s = raw.split(":", 1)
+            if int(layer_s) == int(layer_idx):
+                force_full_mask[int(head_s)] = True
+        self.register_buffer(
+            "_aha_force_full_heads_mask",
+            force_full_mask,
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[dict],  # dict with "global" and "local" keys
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        # === Q projection + (optional) gate extraction ===
+        q_proj_out = self.q_proj(hidden_states)
+        if self._has_dyn_gate_logits:
+            query_states, gate = torch.split(
+                q_proj_out,
+                [
+                    self.config.num_attention_heads * self.head_dim,
+                    self.aha_router_outputs,
+                ],
+                dim=-1,
+            )
+        else:
+            query_states = q_proj_out
+            # DUO mode: use a zero placeholder so downstream dtype/device
+            # checks stay unchanged. The actual gate comes from the
+            # per-kv-head scalar ``full_attention_heads`` below.
+            gate = torch.zeros(
+                *input_shape, self.config.num_key_value_heads,
+                dtype=query_states.dtype, device=query_states.device,
+            )
+
+        # Qwen3-style: norm then reshape then transpose
+        query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        # === Gate computation ===
+        # gate shape: (B, Q, H_kv) — per-(token, kv_head) in dynamic mode,
+        # or per-kv_head broadcast to (B, Q, H_kv) in DUO mode.
+        # At inference, threshold is read from env var AHA_GATE_HARD_THRESHOLD
+        # so users can trade sparsity for quality without re-training.
+        # Default 0.5 keeps training-time behaviour. Lowering to e.g. 0.35
+        # lets more (token, kv_head) pairs route through global attention,
+        # which recovers long-CoT quality at the cost of sparsity.
+        aha_mode = getattr(self.config, "aha_mode", "dynamic")
+        if aha_mode == "duo":
+            # DuoAttention: alpha is per-(layer, kv_head) learnable scalar.
+            # gate_soft = clamp(alpha, 0, 1); gate_hard = (alpha > 0.5).
+            # Both broadcast to (B, Q, H_kv) for a uniform downstream API
+            # (so the distill / aux / blend paths do not need to branch).
+            qhead_alpha = getattr(self, "_aha_qhead_full_attention_heads", None)
+            if qhead_alpha is not None:
+                # Eval-only oracle diagnostic: route at Q-head granularity
+                # instead of KV-head/GQA-group granularity. This does not
+                # represent a KV-capacity-saving deployment point; it isolates
+                # whether GQA grouping itself is the quality bottleneck.
+                alpha = qhead_alpha.to(gate.dtype).clamp(0.0, 1.0)
+                gate_soft = alpha.view(1, 1, -1).expand(*input_shape, -1)
+                import os as _os
+                _thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
+                gate_hard = (alpha > _thr).to(gate.dtype).view(1, 1, -1).expand(*input_shape, -1)
+                gate_ste = gate_hard
+
+            else:
+                alpha = self.full_attention_heads.to(gate.dtype).clamp(0.0, 1.0)
+                gate_soft = alpha.view(1, 1, -1).expand(*input_shape, -1)
+                if self.training:
+                    gate_hard = (alpha > 0.5).to(gate.dtype).view(1, 1, -1).expand(*input_shape, -1)
+                    # Straight-through estimator through the clamped alpha so
+                    # L1 / distill gradients flow into the scalar.
+                    gate_ste = gate_soft
+                else:
+                    import os as _os
+                    _thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
+                    gate_hard = (alpha > _thr).to(gate.dtype).view(1, 1, -1).expand(*input_shape, -1)
+                    gate_ste = gate_hard
+        elif aha_mode == "duo_dynamic":
+            # Duo + AHA-on-full-heads:
+            #   duo_full_mask = 1 for static Duo full/retrieval heads
+            #   dyn_gate      = per-token AHA decision inside those heads
+            #   effective_full = duo_full_mask AND dyn_gate
+            #
+            # Duo streaming heads are locked to local/streaming forever. This
+            # preserves Duo's KV-capacity saving: only static-full heads need a
+            # full KV cache. ``aha_force_gate_value=1`` therefore reproduces
+            # the static Duo deployment point exactly, not all-full attention.
+            alpha = self.full_attention_heads.to(gate.dtype).clamp(0.0, 1.0)
+            duo_full = (alpha > 0.5).to(gate.dtype).view(1, 1, -1)
+            dyn_soft = torch.sigmoid(gate)
+            force_gate_value = getattr(self.config, "aha_force_gate_value", None)
+            if force_gate_value is not None:
+                dyn_soft = torch.full_like(dyn_soft, float(force_gate_value))
+                gate_soft = dyn_soft * duo_full
+                gate_hard = gate_soft
+                gate_ste = gate_soft
+            elif self.training:
+                train_thr = float(os.environ.get("AHA_TRAIN_GATE_HARD_THRESHOLD", "0.5"))
+                gate_soft = dyn_soft * duo_full
+                dyn_hard = (dyn_soft > train_thr).to(gate.dtype)
+                gate_hard = dyn_hard * duo_full
+                gate_ste = gate_hard + (gate_soft - gate_soft.detach())
+            else:
+                import os as _os
+                _thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
+                gate_soft = dyn_soft * duo_full
+                dyn_hard = (dyn_soft > _thr).to(gate.dtype)
+                gate_hard = dyn_hard * duo_full
+                gate_ste = gate_hard
+            force_full = self._aha_force_full_heads_mask.to(gate.dtype).view(1, 1, -1)
+            if bool(force_full.any().item()):
+                force_full = force_full * duo_full
+                gate_hard = torch.maximum(gate_hard, force_full)
+                gate_ste = torch.maximum(gate_ste, force_full)
+        else:
+            gate_soft = torch.sigmoid(gate)
+            force_gate_value = getattr(self.config, "aha_force_gate_value", None)
+            if force_gate_value is not None:
+                force = float(force_gate_value)
+                gate_soft = torch.full_like(gate_soft, force)
+                gate_hard = torch.full_like(gate_soft, force)
+                gate_ste = gate_soft
+            elif self.training:
+                train_thr = float(os.environ.get("AHA_TRAIN_GATE_HARD_THRESHOLD", "0.5"))
+                gate_hard = (gate_soft > train_thr).to(gate.dtype)
+                gate_ste = gate_hard + (gate_soft - gate_soft.detach())
+            else:
+                import os as _os
+                _thr = float(_os.environ.get("AHA_GATE_HARD_THRESHOLD", "0.5"))
+                gate_hard = (gate_soft > _thr).to(gate.dtype)
+                gate_ste = gate_hard
+
+        # Normalize the downstream contract to effective
+        # token x KV-head decisions. In shared-token mode there is only one
+        # learned logit per token, but regularization, sparsity accounting, and
+        # attention mixing all see the same KV-head denominator as the
+        # head-granular arm.
+        if gate_soft.shape[-1] == 1 and self.config.num_key_value_heads != 1:
+            effective_shape = (*gate_soft.shape[:-1], self.config.num_key_value_heads)
+            gate_soft = gate_soft.expand(effective_shape)
+            gate_hard = gate_hard.expand(effective_shape)
+            gate_ste = gate_ste.expand(effective_shape)
+
+        if (
+            aha_mode in ("duo", "duo_dynamic")
+            and not self.training
+            and os.environ.get("AHA_DUO_PREFILL_FULL", "0") == "1"
+            and query_states.shape[-2] == key_states.shape[-2]
+        ):
+            # Diagnostic parity with the official DuoAttention HF eval path:
+            # initial prefill is full attention; streaming heads are only
+            # localized during subsequent decode steps.
+            gate_hard = torch.ones_like(gate_hard)
+            gate_ste = torch.ones_like(gate_ste)
+
+        if not self.training:
+            phase = (
+                "decode"
+                if hidden_states.shape[1] == 1 and key_states.shape[-2] > 1
+                else "prefill"
+            )
+            _track_aha_inference_sparsity(
+                gate_hard, gate_soft, self.layer_idx, phase,
+                getattr(
+                    self.config,
+                    "aha_router_granularity",
+                    AHA_ROUTER_GRANULARITY,
+                ),
+                self.aha_router_outputs if self._has_dyn_gate_logits
+                else self.config.num_key_value_heads,
+            )
+
+        # === Attention: double_mix (global + local, blended by gate) ===
+        from transformers.models.qwen3.modeling_qwen3 import eager_attention_forward
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        # Global attention (full context)
+        global_attn_output, _ = attention_interface(
+            self, query_states, key_states, value_states,
+            attention_mask["global"],
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        if getattr(self.config, "_aha_teacher_full_fastpath", False):
+            attn_output = global_attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output = self.o_proj(attn_output)
+            return attn_output, gate_soft, gate_hard, None
+
+        # Local / streaming branch. The local mask was pre-built by the
+        # model's ``forward`` either as sliding-window (legacy dynamic mode
+        # default) or sink+recent (duo mode, and dynamic mode when
+        # ``aha_local_kind="sink_recent"``). Only pass ``sliding_window`` to
+        # the SDPA backend when the mask is the sliding-window kind, otherwise
+        # the backend will double-mask and wipe out the sink positions.
+        local_attn_kwargs = dict(
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+        )
+        local_kind = "sink_recent" if aha_mode in ("duo", "duo_dynamic") else getattr(
+            self.config, "aha_local_kind", "sliding_window"
+        )
+        if local_kind == "sliding_window":
+            local_attn_kwargs["sliding_window"] = self.aha_window_size
+        local_attn_output, _ = attention_interface(
+            self, query_states, key_states, value_states,
+            attention_mask["local"],
+            **local_attn_kwargs,
+            **kwargs,
+        )
+
+        # Blend: normal gate_ste is (B, Q, H_kv); broadcast to (B, Q, H, 1)
+        # by repeating each kv_head decision. The optional qhead oracle above
+        # already emits (B, Q, H), so it bypasses this broadcast.
+        if gate_ste.shape[-1] == self.config.num_attention_heads:
+            g_bqh = gate_ste
+        else:
+            g_bqh = gate_ste.repeat_interleave(self.num_key_value_groups, dim=-1)
+        attn_output = (
+            global_attn_output * g_bqh.unsqueeze(-1)
+            + local_attn_output * (1 - g_bqh.unsqueeze(-1))
+        )
+
+        # Per-layer attention distillation signal. Returns a scalar so the
+        # outer model can aggregate across layers. When gate → 1 the term
+        # vanishes; when gate → 0 on positions where global/local disagree
+        # strongly, the gradient pushes gate back toward 1.
+        #
+        # ``diff`` is ``detach``-ed so distill only trains the gate; it does
+        # not reshape attention outputs through the backbone. This also
+        # prevents the squared-difference graph from retaining activations
+        # for all 28 layers' global/local branches, which was the source of
+        # OOM when gradient checkpointing is disabled on the custom forward.
+        distill_per_layer: Optional[torch.Tensor] = None
+        if self.training and getattr(self.config, "aha_distill_weight", 0.0) > 0.0:
+            with torch.no_grad():
+                diff_sq = (global_attn_output - local_attn_output).float() ** 2
+            one_minus_g_sq = ((1.0 - g_bqh.float()) ** 2).unsqueeze(-1)
+            distill_per_layer = (one_minus_g_sq * diff_sq).mean()
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, gate_soft, gate_hard, distill_per_layer
+
+
+class AHAQwen3DecoderLayer(GradientCheckpointingLayer):
+    """Decoder layer with gradient-checkpointing support.
+
+    Inheriting from ``GradientCheckpointingLayer`` lets HuggingFace Trainer's
+    ``gradient_checkpointing_enable`` take effect on this layer without any
+    explicit ``_gradient_checkpointing_func`` call in the outer forward.
+    """
+
+    def __init__(self, config: AHAQwen3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = AHAQwen3Attention(config=config, layer_idx=layer_idx)
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[dict] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, gate_soft, gate_hard, distill_per_layer = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states, gate_soft, gate_hard, distill_per_layer
+
+
+@dataclass
+class AHACausalLMOutputWithPast(CausalLMOutputWithPast):
+    """Causal LM output with AHA loss breakdown and gate density stats (training / eval with labels)."""
+
+    ce_loss: Optional[torch.FloatTensor] = None
+    gate_aux_loss: Optional[torch.FloatTensor] = None
+    distill_loss: Optional[torch.FloatTensor] = None
+    gate_soft_mean: Optional[torch.FloatTensor] = None
+    gate_hard_mean: Optional[torch.FloatTensor] = None
+
+
+class AHAModelOutputWithPast(BaseModelOutputWithPast):
+    def __init__(self, last_hidden_state, past_key_values=None, hidden_states=None, attentions=None,
+                 all_gate_soft=None, all_gate_hard=None, all_distill_per_layer=None):
+        super().__init__(last_hidden_state=last_hidden_state, past_key_values=past_key_values,
+                         hidden_states=hidden_states, attentions=attentions)
+        self.all_gate_soft = all_gate_soft
+        self.all_gate_hard = all_gate_hard
+        # Tuple of per-layer distill scalars, or empty tuple when distillation is off.
+        self.all_distill_per_layer = all_distill_per_layer if all_distill_per_layer is not None else ()
+
+
+class AHAQwen3Model(Qwen3Model):
+    config_class = AHAQwen3Config
+
+    def __init__(self, config: AHAQwen3Config):
+        from transformers import PreTrainedModel
+        PreTrainedModel.__init__(self, config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [AHAQwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding
+        self.rotary_emb = Qwen3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> AHAModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # Build two masks: global (full causal) and local (sliding window
+        # for dynamic mode, streaming sink+recent for DUO mode).
+        orig_sliding_window = getattr(self.config, "sliding_window", None)
+        self.config.sliding_window = self.config.aha_window_size
+        mask_kwargs = {
+            "config": self.config,
+            "input_embeds": inputs_embeds,
+            "attention_mask": attention_mask,
+            "cache_position": cache_position,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        global_mask = create_causal_mask(**mask_kwargs)
+        aha_mode = getattr(self.config, "aha_mode", "dynamic")
+        # duo / duo_dynamic force sink_recent for backward compat. dynamic mode
+        # reads the local kind from config (default sliding_window for legacy
+        # Bv3 ckpts; hot-started ckpts from duo set sink_recent — see
+        # §9.4.13.3).
+        local_kind = "sink_recent" if aha_mode in ("duo", "duo_dynamic") else getattr(
+            self.config, "aha_local_kind", "sliding_window"
+        )
+        if local_kind == "sink_recent":
+            seq_len_q = inputs_embeds.shape[1]
+            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+            seq_len_kv = seq_len_q + past_seen
+            local_mask = _build_streaming_causal_mask(
+                seq_len_q=seq_len_q,
+                seq_len_kv=seq_len_kv,
+                sink_size=int(getattr(self.config, "duo_sink_size", 64)),
+                recent_size=int(getattr(self.config, "duo_recent_size", 256)),
+                device=inputs_embeds.device,
+                dtype=inputs_embeds.dtype,
+                kv_offset=past_seen,
+            )
+        else:
+            local_mask = create_sliding_window_causal_mask(**mask_kwargs)
+        causal_mask_dict = {"global": global_mask, "local": local_mask}
+        self.config.sliding_window = orig_sliding_window
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        all_gate_soft = ()
+        all_gate_hard = ()
+        all_distill_per_layer: tuple[torch.Tensor, ...] = ()
+        for decoder_layer in self.layers:
+            hidden_states, gate_soft, gate_hard, distill_per_layer = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_dict,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+            all_gate_soft += (gate_soft,)
+            all_gate_hard += (gate_hard,)
+            if distill_per_layer is not None:
+                all_distill_per_layer += (distill_per_layer,)
+
+        hidden_states = self.norm(hidden_states)
+        return AHAModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            all_gate_soft=all_gate_soft,
+            all_gate_hard=all_gate_hard,
+            all_distill_per_layer=all_distill_per_layer,
+        )
+
+
+class AHAQwen3ForCausalLM(Qwen3ForCausalLM):
+    config_class = AHAQwen3Config
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: AHAQwen3Config):
+        # Skip Qwen3ForCausalLM.__init__ which creates Qwen3Model
+        # Go to PreTrainedModel.__init__ directly
+        from transformers import PreTrainedModel
+        PreTrainedModel.__init__(self, config)
+        self.model = AHAQwen3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs,
+    ) -> AHACausalLMOutputWithPast:
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        ce_loss = None
+        gate_aux_loss = None
+        distill_loss = None
+        gate_soft_mean = None
+        gate_hard_mean = None
+
+        if labels is not None:
+            from torch.nn import CrossEntropyLoss
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = CrossEntropyLoss()
+            ce_loss = loss_fct(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1))
+
+        all_gate_soft = outputs.all_gate_soft
+        all_gate_hard = outputs.all_gate_hard
+        if labels is not None and len(all_gate_soft) > 0 and len(all_gate_hard) > 0:
+            gate_soft_stack = torch.stack(
+                [g.to(hidden_states.device) for g in all_gate_soft], dim=-1
+            ).float()
+            gate_hard_stack = torch.stack(
+                [g.to(hidden_states.device).float() for g in all_gate_hard], dim=-1
+            )
+            shift_gate_soft = gate_soft_stack[:, :-1, :].contiguous()
+            shift_gate_hard = gate_hard_stack[:, :-1, :].contiguous()
+            # Match the CE/PPL population exactly. The old unmasked mean
+            # counted padded sequence positions as real gate decisions, which
+            # made the training sparsity curve disagree with inference-time
+            # sparsity on variable-length batches.
+            valid_gate = shift_labels.ne(-100)
+            if attention_mask is not None:
+                valid_gate = valid_gate & attention_mask[:, 1:].to(torch.bool)
+            if bool(valid_gate.any()):
+                valid_gate = valid_gate[..., None, None].to(shift_gate_soft.dtype)
+                normalizer = (
+                    valid_gate.sum()
+                    * shift_gate_soft.shape[-2]
+                    * shift_gate_soft.shape[-1]
+                )
+                gate_soft_mean = (shift_gate_soft * valid_gate).sum() / normalizer
+                gate_hard_mean = (shift_gate_hard * valid_gate).sum() / normalizer
+            else:
+                gate_soft_mean = shift_gate_soft.mean()
+                gate_hard_mean = shift_gate_hard.mean()
+
+        # Per-layer distillation loss: mean over layers of
+        # ``mean(((1-g)·(global-local))**2)``. Populated only when training
+        # with ``config.aha_distill_weight > 0`` (see AHAQwen3Attention).
+        all_distill = outputs.all_distill_per_layer
+        if self.training and len(all_distill) > 0:
+            distill_loss = torch.stack(list(all_distill)).mean()
+
+        # === Loss assembly ===
+        # Eval path: standard LM loss for compatibility with HF metrics.
+        # Train path: weighted sum of (ce, gate_aux, distill) so callers can
+        # zero any term by setting its weight to 0 (e.g. ``aha_ce_weight=0``
+        # to train the gate purely from ``aux + distill``).
+        if not self.training:
+            loss = ce_loss
+        else:
+            ce_weight = float(getattr(self.config, "aha_ce_weight", 1.0))
+            distill_weight = float(getattr(self.config, "aha_distill_weight", 0.0))
+            terms: list[torch.Tensor] = []
+            if ce_loss is not None and ce_weight > 0.0:
+                terms.append(ce_weight * ce_loss)
+            if gate_soft_mean is not None:
+                reg_weight = float(getattr(self.config, "aha_reg_weight", -1.0))
+                if reg_weight >= 0.0:
+                    # Duo-style direct sparsity term:
+                    # alpha/gate near 1 means full/global attention, so penalize
+                    # mean gate usage directly.
+                    gate_aux_loss = reg_weight * gate_soft_mean
+                else:
+                    # Legacy hinge aux: penalise only the excess above the target
+                    # ceiling. Below ``τ`` the gradient is zero, so CE and
+                    # distill freely shape where the gate opens without an
+                    # unbounded downward push from aux. See AHAQwen3Config
+                    # for the rationale and the collapse that motivated it.
+                    gate_target = float(getattr(self.config, "aha_gate_target", 1.0))
+                    excess = torch.clamp(gate_soft_mean - gate_target, min=0.0)
+                    gate_aux_loss = self.config.aha_lambda * excess
+                terms.append(gate_aux_loss)
+            if distill_loss is not None and distill_weight > 0.0:
+                terms.append(distill_weight * distill_loss)
+            if terms:
+                loss = terms[0]
+                for t in terms[1:]:
+                    loss = loss + t
+
+        return AHACausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            ce_loss=ce_loss,
+            gate_aux_loss=gate_aux_loss,
+            distill_loss=distill_loss,
+            gate_soft_mean=gate_soft_mean,
+            gate_hard_mean=gate_hard_mean,
+        )
+
+    @classmethod
+    def from_pretrained_qwen3(cls, qwen3_path: str, aha_window_size: int = AHA_WINDOW_SIZE,
+                               aha_lambda: float = AHA_LAMBDA,
+                               aha_distill_weight: float = AHA_DISTILL_WEIGHT,
+                               aha_ce_weight: float = AHA_CE_WEIGHT,
+                               aha_gate_target: float = AHA_GATE_TARGET,
+                               aha_reg_weight: float = -1.0,
+                               aha_force_gate_value=None,
+                               aha_mode: str = AHA_MODE,
+                               duo_sink_size: int = DUO_SINK_SIZE,
+                               duo_recent_size: int = DUO_RECENT_SIZE,
+                               duo_alpha_init: float = DUO_ALPHA_INIT,
+                               aha_local_kind: str = AHA_LOCAL_KIND,
+                               aha_router_granularity: str = AHA_ROUTER_GRANULARITY,
+                               **kwargs):
+        """Load a pretrained Qwen3 model and convert to AHA-Qwen3.
+
+        Initializes the gate weights in q_proj randomly (the extra num_heads outputs).
+        All other weights are copied from the pretrained model.
+        """
+        # Load original Qwen3
+        original = Qwen3ForCausalLM.from_pretrained(qwen3_path, **kwargs)
+        original_config = original.config
+
+        # Create AHA config
+        config_dict = original_config.to_dict()
+        config_dict["aha_window_size"] = aha_window_size
+        config_dict["aha_lambda"] = aha_lambda
+        config_dict["aha_distill_weight"] = aha_distill_weight
+        config_dict["aha_ce_weight"] = aha_ce_weight
+        config_dict["aha_gate_target"] = aha_gate_target
+        config_dict["aha_reg_weight"] = aha_reg_weight
+        config_dict["aha_force_gate_value"] = aha_force_gate_value
+        config_dict["aha_mode"] = aha_mode
+        config_dict["duo_sink_size"] = duo_sink_size
+        config_dict["duo_recent_size"] = duo_recent_size
+        config_dict["duo_alpha_init"] = duo_alpha_init
+        config_dict["aha_local_kind"] = aha_local_kind
+        config_dict["aha_router_granularity"] = aha_router_granularity
+        config_dict["model_type"] = "aha_qwen3"
+        aha_config = AHAQwen3Config(**config_dict)
+        aha_config._attn_implementation = kwargs.get("attn_implementation", "sdpa")
+
+        # Create AHA model
+        aha_model = cls(aha_config)
+
+        # Copy weights
+        # embed_tokens, lm_head, norm
+        aha_model.model.embed_tokens.load_state_dict(original.model.embed_tokens.state_dict())
+        aha_model.lm_head.load_state_dict(original.lm_head.state_dict())
+        aha_model.model.norm.load_state_dict(original.model.norm.state_dict())
+        aha_model.model.rotary_emb.load_state_dict(original.model.rotary_emb.state_dict())
+
+        # Copy per-layer weights
+        num_heads = original_config.num_attention_heads
+        head_dim = aha_config.head_dim
+
+        for i in range(original_config.num_hidden_layers):
+            orig_layer = original.model.layers[i]
+            aha_layer = aha_model.model.layers[i]
+
+            # MLP + norms: direct copy
+            aha_layer.mlp.load_state_dict(orig_layer.mlp.state_dict())
+            aha_layer.input_layernorm.load_state_dict(orig_layer.input_layernorm.state_dict())
+            aha_layer.post_attention_layernorm.load_state_dict(orig_layer.post_attention_layernorm.state_dict())
+
+            # Attention: k_proj, v_proj, o_proj, q_norm, k_norm — direct copy
+            aha_attn = aha_layer.self_attn
+            orig_attn = orig_layer.self_attn
+            aha_attn.k_proj.load_state_dict(orig_attn.k_proj.state_dict())
+            aha_attn.v_proj.load_state_dict(orig_attn.v_proj.state_dict())
+            aha_attn.o_proj.load_state_dict(orig_attn.o_proj.state_dict())
+            aha_attn.q_norm.load_state_dict(orig_attn.q_norm.state_dict())
+            aha_attn.k_norm.load_state_dict(orig_attn.k_norm.state_dict())
+
+            # q_proj:
+            # - dynamic / duo_dynamic mode: copy original Q rows, init extra
+            #   gate rows random (~0.5).
+            # - duo mode: q_proj is identical shape to base Qwen3 (no gate
+            #   rows), so direct load_state_dict works.
+            orig_q_weight = orig_attn.q_proj.weight.data  # (num_heads*head_dim, hidden_size)
+            aha_q_weight = aha_attn.q_proj.weight.data
+            if aha_attn._has_dyn_gate_logits:
+                aha_q_weight[:num_heads * head_dim, :] = orig_q_weight
+                nn.init.normal_(aha_q_weight[num_heads * head_dim:, :], mean=0.0, std=0.01)
+                if orig_attn.q_proj.bias is not None and aha_attn.q_proj.bias is not None:
+                    aha_attn.q_proj.bias.data[:num_heads * head_dim] = orig_attn.q_proj.bias.data
+                    nn.init.zeros_(aha_attn.q_proj.bias.data[num_heads * head_dim:])
+            else:
+                aha_attn.q_proj.load_state_dict(orig_attn.q_proj.state_dict())
+
+        del original
+        return aha_model
+
+    @classmethod
+    def from_pretrained_aha(cls, checkpoint_path: str, **kwargs):
+        """Load an AHA-Qwen3 checkpoint saved by `trainer.save_model`.
+
+        Supports both single-file (``model.safetensors``) and multi-shard
+        (``model.safetensors.index.json`` + ``model-*-of-*.safetensors``)
+        layouts. The latter is automatically used by HF when a checkpoint
+        exceeds the single-file size threshold (e.g. Qwen3-1.7B+).
+        """
+        config = AHAQwen3Config.from_pretrained(checkpoint_path)
+        config._attn_implementation = kwargs.get(
+            "attn_implementation",
+            getattr(config, "_attn_implementation", "sdpa"),
+        )
+
+        model = cls(config)
+
+        single_path = os.path.join(checkpoint_path, "model.safetensors")
+        index_path = os.path.join(checkpoint_path, "model.safetensors.index.json")
+        if os.path.exists(single_path):
+            state_dict = load_file(single_path)
+        elif os.path.exists(index_path):
+            with open(index_path) as f:
+                shard_index = json.load(f)
+            shard_files = sorted(set(shard_index["weight_map"].values()))
+            state_dict = {}
+            for shard in shard_files:
+                state_dict.update(load_file(os.path.join(checkpoint_path, shard)))
+        else:
+            raise FileNotFoundError(
+                f"No safetensors file or index found in {checkpoint_path}. "
+                f"Expected either model.safetensors or model.safetensors.index.json."
+            )
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+        allowed_missing = {"lm_head.weight"} if getattr(config, "tie_word_embeddings", False) else set()
+        extra_missing = set(missing_keys) - allowed_missing
+        if extra_missing or unexpected_keys:
+            raise RuntimeError(
+                f"Unexpected checkpoint mismatch. missing={sorted(extra_missing)}, "
+                f"unexpected={sorted(unexpected_keys)}"
+            )
+        if "lm_head.weight" in missing_keys:
+            model.tie_weights()
+
+        _apply_low_alpha_force_full_heads_from_env(model)
+
+        torch_dtype = kwargs.get("torch_dtype")
+        if torch_dtype is not None:
+            model = model.to(dtype=torch_dtype)
+        return model

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(encoding="utf-8").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(encoding="utf-8").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="", encoding="utf-8") 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(encoding="utf-8").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(encoding="utf-8"))
+        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(encoding="utf-8"))
+    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(encoding="utf-8"))
+            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(encoding="utf-8"))
+        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", encoding="utf-8"
+    )
+
+    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="", encoding="utf-8") 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="", encoding="utf-8") 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(encoding="utf-8").splitlines() if line.strip()]
+
+
+def main() -> None:
+    manifest = json.loads((RECIPE / "manifest.json").read_text(encoding="utf-8"))
+    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(encoding="utf-8"))
+    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/smoke_test.py

@@ -0,0 +1,122 @@
+import os
+import sys
+import time
+import torch
+
+# Add recipe path to sys.path
+_RECIPE_ROOT = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _RECIPE_ROOT)
+
+from modeling_aha_qwen3 import AHAQwen3ForCausalLM, AHAQwen3Config
+from router_training_utils import RowWiseAdamW
+
+def main():
+    AHAQwen3Config.register_for_auto_class()
+    AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+
+    # The repo root is the parent directory of recipe
+    model_path = os.path.dirname(_RECIPE_ROOT)
+
+    print("Loading model in BF16...", flush=True)
+    model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
+        model_path,
+        aha_window_size=128,
+        aha_lambda=3e-4,
+        aha_distill_weight=0.0,
+        aha_ce_weight=1.0,
+        aha_gate_target=1.0,
+        aha_reg_weight=0.01,
+        aha_mode="dynamic",
+        aha_router_granularity="token",
+        duo_sink_size=64,
+        duo_recent_size=256,
+        duo_alpha_init=1.0,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+    )
+
+    # Freeze embeddings and LM head as in stage 2 SFT training
+    for param in model.model.embed_tokens.parameters():
+        param.requires_grad = False
+    for param in model.lm_head.parameters():
+        param.requires_grad = False
+
+    print("Moving model to CUDA...", flush=True)
+    model = model.to("cuda")
+
+    # Configure RowWiseAdamW optimizer
+    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 = 3e-7 / 3e-6  # backbone_lr / gate_lr
+
+    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
+    ]
+
+    param_groups = [{"params": gate_params, "lr": 3e-6}]
+    if backbone_params:
+        param_groups.append({"params": backbone_params, "lr": 3e-7})
+
+    optimizer = RowWiseAdamW(
+        param_groups,
+        row_scales=row_scales,
+        weight_decay=0.0,
+    )
+
+    # Allocate a batch of seq_len=8192
+    seq_len = 8192
+    print(f"Allocating dummy batch: batch_size=1, seq_len={seq_len}", flush=True)
+    input_ids = torch.randint(0, model.config.vocab_size, (1, seq_len), device="cuda")
+    labels = input_ids.clone()
+
+    # Enable gradient checkpointing
+    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+
+    # Warmup step (GPU caching, model trace creation, etc.)
+    print("Warmup step...", flush=True)
+    outputs = model(input_ids=input_ids, labels=labels)
+    loss = outputs.loss
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+
+    # Reset peak memory stats and time the next step
+    print("Starting measured smoke test step...", flush=True)
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    start_time = time.time()
+
+    outputs = model(input_ids=input_ids, labels=labels)
+    loss = outputs.loss
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+
+    torch.cuda.synchronize()
+    step_time = time.time() - start_time
+
+    peak_mem = torch.cuda.max_memory_allocated() / (1024 ** 3)
+    reserved_mem = torch.cuda.memory_reserved() / (1024 ** 3)
+
+    print("=== SMOKE TEST RESULTS ===", flush=True)
+    print(f"Peak VRAM: {peak_mem:.4f} GB", flush=True)
+    print(f"Reserved VRAM: {reserved_mem:.4f} GB", flush=True)
+    print(f"Step time: {step_time:.4f} seconds", flush=True)
+    print("==========================", flush=True)
+
+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(encoding="utf-8"))
+            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
+  }
+}