YAML Metadata Warning:empty or missing yaml metadata in repo card
Check out the documentation for more information.
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