pure-aha-longctx

Pure AHA long-context training handoff

This package runs pure AHA training for the BooksumMultipasskey retrieve setup. It does not use DuoAttention 50% fixed heads.

Main entry point:

./run_training.sh

Use MODE=smoke ./run_training.sh for a 1-GPU smoke test, or MODE=preflight ./run_training.sh for a 20-step 8-GPU preflight before the full run.

This repo contains dataset/artifact files stored with Git LFS. Before cloning on a fresh machine, make sure git-lfs is installed and run:

git lfs install

Environment

Run inside any GPU environment with:

torch
transformers
datasets
flash-attn
numpy

Docker is fine, but not required. If using Docker, enter the container first, then run the same commands below.

The default model is downloaded from Hugging Face:

gradientai/Llama-3-8B-Instruct-Gradient-1048k

To use a local model, set MODEL=/path/to/model.

1-GPU smoke test

CUDA_VISIBLE_DEVICES=0 MODE=smoke ./run_training.sh

Success signs:

attention_mode=pure_aha static_sparsity=0.0000
training finished

Recommended 8-GPU 32K preflight before full training

Do this before spending a full 4000-step run. This uses the same 8-GPU, 32K configuration as the final run, but stops after 20 steps.

MODE=preflight ./run_training.sh

Preflight must show:

attention_mode=pure_aha
static_sparsity=0.0
config lr is 3e-4; logged lr starts lower because of warmup
metrics.jsonl reaches step 20
aha_gate_latest.pt exists
no OOM / NaN / NCCL failure

Only after this passes, run the full training below.

Full 8-GPU training

./run_training.sh

Default full training keeps the previous B7 Duo50+AHA data/context/step protocol, except for the intended method change to pure AHA and the corrected bias-free gate:

context length = 32000
num passkeys = 10
steps = 4000
lr = 3e-4
reg_weight = 0.05
sink_size = 128
recent_size = 256
gate_init_open_prob = unused; gate is bias-free and initializes at sigmoid(0)=0.5
gate_threshold = 0.5
max_grad_norm = 0.0

Outputs are written under:

results/aha_oracle/

Each run writes:

commit.txt
git_status.txt
config.json
metrics.jsonl
aha_gate_latest.pt
aha_gate_step*.pt