Unit 3.4 · about 3 h
307 / 324Keeping training numerically healthy
One bad batch can end a run, and mixed precision makes bad batches likelier. Clipping, scaling and logging turn a crash into a diagnosed event with a known remedy.
- explain
- measure
You can now
- Explain the clipping rule, where it sits in the step order, and how it changes with a gradient scaler.
- Explain
autocastdtype choice and howGradScalerprotects small gradients from underflow. - Log gradient norms, loss scale and skipped steps, and read overflow signals from the log.
Covers
- Gradient clipping and numerical-event logging
- Mixed-precision vocabulary, scaling and overflow signals
Study, in this order
3 steps
- CS336 Assignment 1 handout: gradient clippingStanford CS336 Assignment 1 - training and decoding contract
Section 4.5, Gradient clipping (p. 34): the global norm rule with max \(M\) and
eps1e-6, applied in place after the backward pass - PyTorch AMP recipe: `autocast`, `GradScaler` and clipping orderPyTorch Automatic Mixed Precision recipe
Sections Adding
torch.autocast(forward underautocast, never backward), AddingGradScaler(scale, backward, conditional step, update) and Inspecting/modifying gradients (unscale_before clipping), plus Troubleshooting forinfandNaNloss - LLMs from Scratch Appendix D: worked clipping exampleLLMs from Scratch companion - training and generation
Section D.3 (L2 norm worked example on the gradient matrix \(G\), before-and-after clipping comparison) and Section D.4 (clipping inside the modified training function)
Focus on
- Clipping rescales the whole gradient vector when its global norm passes \(M\). Direction is preserved; only the step length is capped.
- Under a scaler, gradients are scaled by \(s\), so
unscale_must run before the norm check or the threshold tests the wrong numbers. - Autocast keeps compute-heavy ops in
float16but reductions and the loss infloat32, because they need the dynamic range. - A skipped optimizer step is protection working, not a bug. The log should show the skip, the shrunken scale and the recovery.
Reproduce closed-book
unscale_, the skip-on-inf rule and the per-step values to log.Traps
- Clipping the scaled gradients. The norm test then uses a threshold meant for unscaled values and either never fires or always fires.
- Running
backwardunderautocast. Backward ops must run in the dtype chosen for theforwardops, outside the context. - Retrying a diverged run with the same seed and settings. The overflow recurs; lower the scale policy, the clip threshold or the learning rate first.
Practice
CS336 Assignment 1: `gradient_clipping`
Problem (gradient_clipping) in Section 4.5; connect your function to run_gradient_clipping in tests/adapters.py and pass test_gradient_clipping in tests/test_nn_utils.py
The test compares against clip_grad_norm_ with a frozen parameter and max_norm 1e-2. Implement it in your own repository.
PyTorch AMP recipe: add `autocast` and `GradScaler` to the default-precision loop
Sections Adding torch.autocast through All together: run the default-precision loop, convert it, and keep the output-dtype and loss-dtype asserts green
Run on any CUDA machine you can access; on CPU, read the scale and skip logic and wire the same order into your own loop without claiming speedups.
Self-test
-
The global gradient norm is 5 and \(M\) is 1. What does clipping do?
It multiplies every gradient by about \(\frac{1}{5}\), the factor \(M/(\text{norm}+\text{eps})\). The resulting norm sits just under 1 and the update direction is unchanged. -
Why must
scaler.unscale_run before gradient clipping?Backward under a scaler leaves gradients multiplied by the scale \(s\). Clipping tests the norm against \(M\), so on scaled values the comparison is wrong by a factor of \(s\) and corrupts the step. -
The scaler finds
infgradients and skipsoptimizer.step. What must still run, and what happens to the batch?scaler.update()still runs so the scale shrinks for the next step. The batch contributes no update; its gradients are discarded. -
Under
autocast, linear layers run infloat16but the loss stays infloat32. Why?Linears are compute-heavy and safe in half precision, while the loss is a reduction that needsfloat32dynamic range. The recipe asserts output isfloat16and loss isfloat32.
Completion criterion
Stop when
Your clipping matches the reference in test_gradient_clipping, and your AMP loop holds the recipe dtype asserts while you explain each skip from your own log.
Leads to
- 3.5 Checkpoints that really resume Add the scaler state to the checkpoint so protection survives a restart.
- 4.4 Uncertainty, robustness and honest claims Treat seed and run variability as uncertainty instead of noise to ignore.
Comes after
- 2.6 The modern decoder block: residuals, RMSNorm and SwiGLUTrack activation and gradient behavior during training.
- 3.2 Feeding the model: batches, initialization and first stepsWatch the first steps for exploding norms and overflow.
- 3.3 AdamW, schedules and the global token batchProtect the tuned loop with clipping, scaling and event logging.
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.
- 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.
Also in this section
- 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.
- 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.