AR progress
Done
- Milestone 1: data pipeline (tokenizer, prepare_data, dataset) — 10M FineWeb tokens at /fsx/ajay-stu/ar-data/bin/fineweb.bin
- Milestone 2: model architecture — 4.022B params, replicates Qwen3-4B
- Milestone 3: single-GPU training loop — loss drops 12 -> 7.2 on 0.5B test. WORKS.
- Fixed: weight init bug (logits were exploding to std~32; added small-normal + residual scaling)
Next: Milestone 4 — scale to 8 GPUs (FSDP)
Open issues
- Storage: /fsx is 100% full. Need /fsxb write access from admin for real run.
- TODO: intra-document masking (option 2) before real training
- TODO: full 4B needs FSDP + bf16 + grad checkpointing (single GPU OOMs)
Run that works
python -m ar.train.train --bin /fsx/ajay-stu/ar-data/bin/fineweb.bin --model_size small --max_steps 200 --batch_size 2 --seq_len 1024 --warmup_steps 20
Session 2 update
- Milestone 4 DONE: FSDP2 multi-GPU works. Full 4B trains across 8 H100s (was OOM on 1 GPU).
- train_fsdp.py is the multi-GPU script. Launch: torchrun --nproc_per_node=8 -m ar.train.train_fsdp ...
- Verified: 2-GPU (0.5B) and 8-GPU (full 4B) both run, loss at correct 11.9 baseline.
BLOCKER for real training: storage + data
- Only have 10M tokens (40MB, english only) — too small, would overfit.
- Need: (1) /fsxb access or other big storage, (2) prep real multi-source dataset (200B+ tokens for strong base)
- Storage frees up in mornings when others aren't using cluster.
Next session (morning)
- Sort storage (check /fsxb, /fsx free space)
- Prep real dataset: all 3 sources, target 200B+ tokens (uint32 = ~800GB, so storage matters)
- THEN launch real training run (this is the overnight-worthy job)
Minor TODOs
- torch.manual_seed() in setup to silence DTensor RNG warning
- intra-document masking (option 2) before real run
- Milestones remaining: M5 multi-node, M6 efficiency, M7 sampling/eval