Unit 3.3 · about 4 h
306 / 324AdamW, schedules and the global token batch
Learning rate, batch size and schedule interact, and runs compare by tokens, not steps. Until you can state any run's token batch, tuning is guessing.
- derive
- implement
You can now
- Derive one
AdamWstep from the two moments, the bias correction and the decoupled decay. - Implement
AdamWand the warmup-cosine schedule so both pass the CS336 public tests. - Derive the global token batch and step count from sequences, length, accumulation and parallelism.
Covers
- AdamW moments and decoupled weight decay
- Warmup, cosine decay and schedule state
- Gradient accumulation and global token batch
Study, in this order
4 steps
- CS336 Assignment 1 handout: `AdamW` and its accountingStanford CS336 Assignment 1 - training and decoding contract
Section 4.3,
AdamW(pp. 31 to 32): Algorithm 1 with moments, bias-corrected step and decoupled decay, typical betas(0.9, 0.999)versus(0.9, 0.95), then theadamw_accountingdeliverables for memory, FLOPs and MFU time - CS336 Assignment 1 handout: cosine schedule with warmupStanford CS336 Assignment 1 - training and decoding contract
Section 4.4, Learning rate scheduling (p. 33): linear warmup, cosine annealing and post-annealing phases as a pure function of step \(t\)
- LLMs from Scratch Appendix D: warmup, cosine and the modified loopLLMs from Scratch companion - training and generation
Sections D.1 (warmup over 0.1 percent to 20 percent of steps), D.2 (cosine decay) and D.4 (the modified training function combining both with gradient clipping)
- Modern LLM Notebook 10-training-loss: mini `Trainer` and batch arithmeticModern LLM Notebook - configuration, loss and generation
Section 6 (mini
Trainerskeleton with optimizer and scheduler step order) and Homework 2 (effective batch size from per-device batch, GPU count and accumulation steps)
Focus on
- Weight decay sits outside the gradient in
AdamW: parameters shrink by \(\alpha\) times \(\lambda\) directly, not through the moments. - The schedule is a pure function of \(t\), so resuming needs \(t\) itself rather than a schedule object.
- Global tokens per optimizer step equal sequences times length times accumulation times data-parallel width.
- Appendix D sizes warmup as a fraction of total steps, typically 0.1 to 20 percent; many runs fix a step count instead. Either way, the warmup length belongs in the config.
Reproduce closed-book
AdamW update from Algorithm 1, the learning rate at a step inside each of the three schedule phases, and the step count for a 327.68M-token budget.Traps
- Adding weight decay to the gradients and calling it
AdamW. That isAdamwith L2, and the decay gets divided by the adaptive denominator. - Stepping the scheduler per microbatch under accumulation. The schedule then decays faster than intended because it sees more steps.
- Comparing two runs by step count at different batch sizes. Fix tokens first: batch times steps times length is the common currency.
Practice
Problem (adamw) in Section 4.3; expose your optimizer through get_adamw_cls in tests/adapters.py and pass test_adamw in tests/test_optimizer.py
The test trains a small linear model for 1000 steps and matches PyTorch AdamW or the reference snapshot. Implement it in your own repository.
CS336 Assignment 1: `learning_rate_schedule`
Problem (learning_rate_schedule) in Section 4.4; expose your function through get_lr_cosine_schedule in tests/adapters.py and pass test_get_lr_cosine_schedule in tests/test_optimizer.py
Modern LLM Notebook 10-training-loss Homework 2: effective batch size
Homework 2: compute effective_batch_size from per-device batch, GPU count and accumulation steps; the cell asserts 64
Derive the token batch for your own planned run the same way before trusting any step-based comparison.
Self-test
-
How does
AdamWapply weight decay differently fromAdamwith L2 regularization?AdamWsubtracts \(\alpha\) times \(\lambda\) times \(\theta\) from the parameters directly. With L2 inAdam, the decay enters the gradient and is rescaled by the adaptive denominator, so large-gradient parameters get less decay. -
Per-device batch 2, 4 GPUs, 8 accumulation steps. What is the effective batch size?
64 samples. \(2 \times 4 \times 8\) accumulation groups form one optimizer step, matching the notebook Homework 2 assert. -
Batch size 1024, context 256, budget 327.68M tokens. How many optimizer steps?
1250. Each step consumes \(1024 \times 256=262144\) tokens, and \(327680000/262144=1250\). -
Training resumes at step \(t\). What schedule state must the checkpoint hold?
Only \(t\). The schedule is a closed-form function of \(t\), so restoring the exact step restores the exact rate. The optimizer moments are separate state and must be restored alongside.
Completion criterion
Stop when
Your AdamW and schedule pass both optimizer tests including phase boundaries, and you can state any run's token batch and step count from its config.
Leads to
- 3.4 Keeping training numerically healthy Protect the tuned loop with clipping, scaling and event logging.
- 3.5 Checkpoints that really resume Persist optimizer moments and the schedule position with the model.
Comes after
- 3.2 Feeding the model: batches, initialization and first stepsHand the batches to
AdamWand count the global token batch.
Resources for this unit
- LLMs from Scratch companion - training and generation Use the ch05 training and generation code for the small pretraining loop. Add full checkpoint state, explicit loss masks and resume-equivalence tests that the base example does not supply.
- Stanford CS336 Assignment 1 - training and decoding contract Use cross_entropy, adamw, adamw_accounting, learning_rate_schedule, gradient_clipping, data_loading, checkpointing, training_together and decoding specifications plus public tests. Debug on CPU and implement independently in your own repository.
- Modern LLM Notebook - configuration, loss and generation Use English notebooks 09-model-config, 10-training-loss and 20-generation to visualize configuration ownership, shifted loss and decoding policies. Change edge cases and cross-check probability calculations against your own implementation.
Also in this section
- PyTorch saving, loading and reproducibility references Use the state-dict and reproducibility notes to enumerate all state needed for the supported resume guarantee. Test model, optimizer, schedule, scaler, data order and RNG state across an interruption.
- Hugging Face LLM Course - library translation Only after the blank-file loop passes, map tokenization, model configuration, data collation, generation configuration and stopping behavior into Transformers APIs. Keep the from-scratch tests as the authority.
- PyTorch Automatic Mixed Precision recipe Use the
autocast,GradScaler, unscale-before-clip, scaler checkpointing and inf-NaN troubleshooting sections to add mixed precision to the blank-file loop. Keep the CS336 clipping contract as the authority for the clipping rule.