Module 5 of 7

Training at scale

What changes when one training run spans thousands of GPUs: feeding data, numerical stability, data, tensor, and pipeline parallelism, sharded optimizer state, and surviving failures.

6 lessons47 min reading

A model that fits on one GPU trains with a loop and an optimizer. A model that does not, trained on trillions of tokens, needs a data pipeline that keeps every GPU busy, number formats that save memory without losing gradients, and a way to split the model, its gradients, and its optimizer state across thousands of devices. Then it needs to survive the hardware failing every few hours.

This module covers each piece with the arithmetic that sizes it: loader throughput, loss-scale behavior, all-reduce time, bytes per GPU under each ZeRO stage, tensor and pipeline communication and bubbles, and the checkpoint interval that maximizes useful work. Afterwards you should be able to pick a parallelism layout for a given model and cluster and explain where its time and memory go.

Lessons

  1. 5.aData pipelinesHow fast must the input pipeline be, and how do you shuffle and resume a dataset that is too large to hold?8 min
  2. 5.bMixed precision and stabilityHow do you train in 16-bit numbers without losing the gradients, and what makes large runs diverge?8 min
  3. 5.cData parallelism and all-reduceHow long does gradient synchronization take, and how do you hide it behind computation?8 min
  4. 5.dZeRO and FSDPHow much memory does each GPU need under each ZeRO stage, and what does sharding cost in communication?7 min
  5. 5.eTensor and pipeline parallelismWhen the model will not fit even sharded, how do you split it across GPUs, and what does each split cost?8 min
  6. 5.fCheckpoints and failuresHow often will a large training job fail, and how often should it checkpoint to lose the least work?8 min