Distributed training splits data, matrix operations or network layers across accelerators. With mixed-precision Adam, a 175-billion-parameter model needs about 2.8 TB for its weights, gradients and optimizer state alone, more than one accelerator can hold.
How to study this pagePredict · Check · Generalize · Continue
PredictSay what you expect.
Pick a concrete input and say what the equation should give before you read the math.
CheckCheck against the math and the papers.
Work through the math, then read the cited papers for the original argument.
GeneralizeName what stays the same.
Identify the part of the idea that survives a change of notation or model, so you can recognize it in another paper or system.
ContinueFollow the connections.
Use the prerequisites and related concepts at the end of the page to fill the next gap.
Why It Matters for Modern Models
Distributed training splits data, matrix operations or network layers across accelerators. With mixed-precision Adam, a 175-billion-parameter model needs about 2.8 TB for its weights, gradients and optimizer state alone, more than one accelerator can hold.
Combining data, tensor and pipeline parallelism, Narayanan et al. (2021) ran training iterations of a 1-trillion-parameter model at 502 petaFLOP/s on 3,072 GPUs.
Communication patterns explain why some architectures and parallel layouts scale better than others.
Questions to Think About
ZeRO stages 1 and 2 add no communication over plain data parallelism, while stage 3 adds about 50%; why does partitioning the parameters themselves cost extra communication (Rajbhandari et al., 2020)?
Why does the one-forward-one-backward (1F1B) pipeline schedule reduce activation memory without shrinking the bubble, and how does the interleaved schedule shrink the bubble by giving each device several smaller stages (Narayanan et al., 2021)?
Core Math
Key Equation
Memory per GPU (ZeRO-3)≈Nd16Ψbytes
Training a large model is split across many accelerators in three main ways, which are often combined ("3D parallelism").
Data parallelism. Every device holds a full copy of the model and processes a different part of the batch; gradients are averaged with an all-reduce before each update:
∇L=N1i=1∑N∇Li
Tensor parallelism. A single matrix multiplication is split across devices, for example by columns (Shoeybi et al., 2019):
Y=XW=X[W1∣W2]=[XW1∣XW2]
Pipeline parallelism. Consecutive groups of layers live on different devices, and each batch is split into m micro-batches that flow through the p stages. Stages sit idle at the start and end of each batch; this "pipeline bubble" costs a fraction (p−1)/m of the ideal compute time (Huang et al., 2019; Narayanan et al., 2021).
ZeRO. Mixed-precision training with Adam stores about 16 bytes per parameter: 2 for 16-bit weights, 2 for 16-bit gradients and 12 for 32-bit optimizer state (master weights, momentum and variance). Plain data parallelism keeps all of it on every device. ZeRO partitions the optimizer state (stage 1), then also the gradients (stage 2), then also the parameters (stage 3) across the Nd data-parallel devices (Rajbhandari et al., 2020). With stage 3, a model with Ψ parameters needs about
Memory per GPU≈Nd16Ψbytes
for weights, gradients and optimizer state, plus the activations.
Key Papers
Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism