| # 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) |
| 1. Sort storage (check /fsxb, /fsx free space) |
| 2. Prep real dataset: all 3 sources, target 200B+ tokens (uint32 = ~800GB, so storage matters) |
| 3. 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 |
| |
