---
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.