Training at scale
P10.llm-systems-frontier.02 · Audience: it-ml · Prerequisites: Inference & serving
The systems that make large-model training feasible, one idea at a time: parallelism and ZeRO, mixed-precision training, scaling laws, and quantization — each with a small worked calculation. Interview depth for the LLM-systems track.
ⓘ Concept: One mental model to keep
The two walls. A large model does not fit in one GPU, and one GPU would be too slow. Everything here splits the model and batch across devices and uses fewer bits so more fits and it runs faster.
Why it matters — Every technique here is an answer to one of two walls — the model does not fit, or one GPU is too slow — so you either split work across devices or spend fewer bits per number.
Data parallel splits the batch (full replica per GPU); tensor parallel splits matmuls within a layer (chatty — keep on NVLink); pipeline parallel splits layers across GPUs; sequence parallel splits the sequence dimension for long contexts. Real runs combine them (3-D parallelism). ZeRO shards the redundant optimizer state (the memory hog) across the data-parallel group instead of replicating it.
| Strategy | What is split | When it wins |
|---|---|---|
| Data parallel (DP) | the batch (full replica per GPU) | model fits on one GPU; scale throughput |
| Tensor parallel (TP) | matmuls within a layer | a single layer is too big; keep on NVLink |
| Pipeline parallel (PP) | layers across GPUs | many layers; cross-node scaling |
| Sequence / context | the sequence dimension | very long contexts blow up activation memory |
Worked example — where the training memory actually goes
Adam keeps momentum + variance + an fp32 master copy — about 12 bytes/param of optimizer state. For a 7B model, that alone is:
optimizer state = 7e9 params * 12 bytes = 84 GB
(vs fp16 weights = 7e9 * 2 = 14 GB)
-> the optimizer state, not the weights, blows the memory budget
-> ZeRO/FSDP shard those 84 GB across the data-parallel GPUsThe 84 GB of optimizer state dwarfs the 14 GB of weights — which is exactly what ZeRO (stages 1→3 shard optimizer states → gradients → parameters) and FSDP attack. Interview point: optimizer state, not weights, is the memory hog.
⚡ Interview Ref — the quick-scan Reference face
The systems that make large-model training feasible: parallelism and ZeRO, mixed-precision training, scaling laws, and quantization. The two walls: a large model does not fit in one GPU, and one GPU is too slow — split the model and batch across devices, and use fewer bits.
1 — Parallelism strategies & ZeRO (q16)
DP splits the batch, TP splits matmuls within a layer (chatty — keep on NVLink), PP splits layers across GPUs (needs micro-batches to fill the pipeline bubble), sequence parallel splits the sequence dimension for long contexts; real runs combine them (3-D / 4-D). ZeRO shards the redundant optimizer state (Adam, ~12 bytes/param) across the DP group — stage 1 optimizer states, stage 2 also gradients, stage 3 also parameters. FSDP is the same idea. Optimizer state, not weights, is the memory hog.
2 — Mixed-precision training (q22)
Keep most math in 16-bit, protect sensitive parts in fp32. fp16 (5 exponent bits) has a narrow range — small gradients underflow, needs loss scaling (multiply loss by S before backward, unscale before the step; dynamic scaling adapts S). bf16 keeps fp32's 8-bit exponent (wide range, no loss scaling) at the cost of mantissa bits — now the default on A100/H100/TPU. Keep an fp32 master copy of the weights.
3 — Scaling laws: Kaplan vs Chinchilla (q28)
Loss falls as a power law in N, D, C. Kaplan (2020): spend extra compute on bigger models → large, under-trained models (GPT-3). Chinchilla (2022): scale N and D together, ~20 tokens per parameter. A 70B Chinchilla-optimal model wants ~1.4T tokens. Deployment twist: over-train smaller models past Chinchilla for cheap inference (the LLaMA strategy).
4 — Quantization (q20)
Store/run in low precision (int8/int4) to cut memory and bandwidth. Weight-only int4 (GPTQ/AWQ) shrinks the model and bytes read per decode step — ideal for memory-bound decode. Weight + activation (W8A8) is harder because of activation outliers. GPTQ minimises output error via Hessian info; AWQ protects salient channels; both are post-training (no retraining).
Interview one-liners
- Optimizer state (~12 bytes/param), not weights, is the memory hog — ZeRO/FSDP shard it; combine DP/TP/PP/sequence parallelism.
- bf16 > fp16 for stability (fp32 exponent, no loss scaling); keep an fp32 master copy.
- Chinchilla ≈ 20 tokens/param; over-train smaller models for cheap inference.
- int4 weight-only (GPTQ/AWQ) is the memory-bound-decode win; activation outliers make activation quantization hard.
📚 Go Further
Curated resources on large-scale training and quantization.
| Type | Resource |
|---|---|
| Reference | Rajbhandari et al. (2020), ZeRO: Memory Optimizations Toward Training Trillion Parameter Models |
| Reference | Micikevicius et al. (2018), Mixed Precision Training |
| Reference | Hoffmann et al. (2022), Training Compute-Optimal Large Language Models — Chinchilla |
| Reference | Kaplan et al. (2020), Scaling Laws for Neural Language Models |
| Reference | Frantar et al. (2022), GPTQ; Lin et al. (2023), AWQ |
| In-app | P10.llm-systems-frontier.01 — Inference & serving, for why decode is memory-bound |
Ask the mentor about this module
Ask a question about this content. The mentor explains and grounds its answer in what you are studying; asking is recorded as a learning signal, not a grade.
Where next?
Later in LLM Systems & Frontier
Go up a level