Why One GPU Is Never Enough
Lesson 1 of 5 in Distributed Training.
Start with the intuition everyone brings from inference: a model with N parameters stored at 2 bytes each needs 2·N bytes of memory. A 7-billion-parameter model is 14 GB — fits on one good accelerator with room to spare. So why does training the same model demand a cluster?
Because training carries luggage that inference never packs. To take one gradient step you need, live in memory at the same time:
- The Weights — the numbers being learned.
- The gradients — one number per parameter, saying which direction to nudge it.
- The optimizer states. Adam, the de-facto training optimizer, keeps two running statistics per parameter (a momentum average and a variance average), and Mixed precision training adds a full-precision master copy of the weights on top. The bookkeeping ends up taking roughly three times the memory of the weights and gradients combined.
- The activations — every intermediate result of the forward pass, kept around so the backward pass can compute gradients. These scale with batch size and sequence length, on top of everything above.
The first three are what the ZeRO paper calls model states, and they are the immovable furniture: they exist for every parameter, at every step, regardless of batch size.
Bar chart showing five components of per-parameter training memory: fp16 weights at 2 bytes, fp16 gradients at 2 bytes, fp32 master weights at 4 bytes, fp32 Adam momentum at 4 bytes, and fp32 Adam variance at 4 bytes, totaling 16 bytes per parameter. The three optimizer-state bars together dwarf the weights and gradients.
Run the multiplier and the single-device dream dies fast: at 16 bytes per parameter, a 70-billion-parameter model carries about 1.1 TB of model states — before a single activation is stored. No accelerator holds that.
And memory is only the first wall. The second is time. The compute-and-budgets module gave you C ≈ 6·N·D FLOPs for a training run; divide a modern pre-training budget by one accelerator’s throughput and the wall-clock answer comes out in years, not weeks. A run that must finish this quarter needs thousands of devices working simultaneously — which means the real question of this module is not whether to distribute but along which axes.
The 16 bytes, line by line — and what the accounting leaves out
The ZeRO paper (Rajbhandari et al. 2019, §3.1) counts model-state memory for a model with Ψ parameters trained with Adam in mixed precision: 2Ψ bytes for fp16 weights, 2Ψ bytes for fp16 gradients, and K·Ψ bytes of optimizer states with K = 12 — an fp32 master copy of the weights (4Ψ) plus fp32 momentum (4Ψ) and variance (4Ψ). Total: (2 + 2 + 12)·Ψ = 16Ψ bytes. That is where “Adam’s bookkeeping is three times the weights and gradients combined” comes from: 12 of the 16 bytes are optimizer states.
Why the fp32 master copy at all? Mixed-precision training (Micikevicius et al. 2017) runs the expensive math in 16-bit for speed, but per-step weight updates are often too small to survive 16-bit rounding — so the optimizer accumulates them into a full-precision copy. The precision-and-stability module picks this thread up, including how BF16 (bfloat16) shifts the details without changing the shape of the accounting.
Two honest caveats. First, the exact K depends on optimizer and precision recipe — 16 bytes is the canonical mixed-precision-Adam figure, not a law of nature. Second, model states are not the whole bill: activations, temporary buffers, and memory fragmentation come on top (ZeRO calls these residual states and attacks them separately). Activation memory scales with batch size × sequence length × depth, and is why activation checkpointing — recompute instead of store — appears in the pipeline lesson.
Interactive checkpoint quiz (2 questions) — open this page in a browser to take it.