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
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:
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. 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 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.
Before training, the launcher automatically:
- verifies SHA-256 for the starting weights, training Arrow file, frozen evaluation inputs, and core training sources;
- checks that the bundled dataset has exactly 1,024 rows;
- checks the 250 HELMET and 30 MRCR frozen rows;
- 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 on codex/l2a-router-granularity. File hashes are recorded in 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.