Learning path
Next moduleModule 28: Prompting and context engineering
Module 27
Training at scale
Train models too big or too slow for one GPU: mixed precision and clipping, gradient checkpointing, collective communication, data/tensor/pipeline parallelism, and ZeRO/FSDP sharding.
- 1. Gradient clipping and mixed precisionClip gradients by their global norm, train in bf16 or fp16 with loss scaling, and diagnose NaNs, skipped steps, and loss spikes.40 min
- 2. Gradient checkpointingTrade compute for memory by storing only some activations and recomputing the rest during backward; size activation memory and pick a checkpointing policy.35 min
- 3. Collective operations from scratchBuild ring reduce-scatter, all-gather, and all-reduce, prove the 2(N-1)/N bandwidth cost, and pick ring or tree with a latency-bandwidth model.40 min
- 4. Scaling and distributed trainingSize training memory per GPU, and choose between data, tensor, pipeline, and sharded data parallelism (and their 3D combination) for a model and cluster.45 min
- 5. ZeRO and FSDPShard optimizer state, gradients, and weights across data-parallel GPUs; compute per-GPU memory and communication for DDP and ZeRO stages 1 to 3, and pick a stage.40 min
A training run is limited by memory, then by communication. This module builds each technique that pushes those limits from scratch, and keeps one set of numbers throughout: a synthetic 7B-style model needs 16 bytes per parameter with mixed-precision AdamW, 108.63 GB in total, and every lesson shows a different way to make that fit.
| Lever | What it saves | What it costs |
|---|---|---|
| Mixed precision (bf16 / fp16) | half the weight and activation bytes, faster matmuls | range or precision limits; fp16 needs loss scaling |
| Gradient clipping | runs that would diverge on one bad batch | nothing, unless it clips every step |
| Gradient checkpointing | most activation memory | about 33% more compute per step |
| Data parallelism | wall-clock time | one gradient all-reduce per step; no memory saving |
| Tensor / pipeline parallelism | weight memory per GPU | per-layer communication / pipeline bubble |
| ZeRO / FSDP | duplicate optimizer state, gradients, weights | up to 1.5 times DDP's communication (stage 3) |
| # | Lesson | You will be able to |
|---|---|---|
| 1 | Gradient clipping and mixed precision | clip by global norm, choose bf16 or fp16, and run and debug a loss scaler |
| 2 | Gradient checkpointing | size activation memory and trade it for recompute |
| 3 | Collective operations from scratch | build ring all-reduce and cost any collective |
| 4 | Scaling and distributed training | size training memory and combine data, tensor, and pipeline parallelism |
| 5 | ZeRO and FSDP | compute per-GPU memory and communication for DDP and ZeRO stages 1 to 3 |
Interactive lab: Data parallelism and ZeRO memory