Large-scale distributed training (DDP, FSDP, ZeRO & Megatron-LM)

Hook

One GPU can't hold a modern model, let alone train it. So the model and its data get split across many GPUs — but "split" is not one operation. Split the batch, split the weight matrices, or split the layers, and you get three different training runs with three very different amounts of chatter between machines.

Intuition

Each step, this strategy communicates the full gradient, all-reduced across every GPU, once per training step.

data parallel

48

tensor parallel

24

pipeline parallel

8

selected: data parallel = 48 units, on a 2-layer, 4×4 model across 4 GPUs

The same tiny 2-layer model, the same 4 GPUs. Switch strategies and watch what actually gets sent over the network change — not just how much of it, but what kind of thing it even is.

Formalize

All three strategies eventually bottom out in a ring all-reduce: $P$ GPUs arranged in a ring, each sending and receiving $(P-1)/P$ of a payload during a reduce-scatter pass, then again during an all-gather pass, to end with every GPU holding the same combined result:

V_{\text{ring}}(s, P) = \frac{2(P-1)}{P}\,s

$s$ — the size of the payload being combined across GPUs.
$P$ — the number of GPUs.
$V_{\text{ring}}$ — total communication volume for one ring all-reduce of that payload.

The three strategies differ only in what payload goes through this operation, and how often:

V_{\text{data}} = V_{\text{ring}}(\text{params}, P) \qquad V_{\text{tensor}} = 2L \cdot V_{\text{ring}}(H, P) \qquad V_{\text{pipeline}} = 2H(L-1)

$\text{params}$ — the model's total parameter count (every GPU holds a full copy).
$H$ — the hidden dimension of one layer's activation.
$L$ — the number of layers.

Data parallel: sync the whole gradient
Every GPU holds the entire model and computes on its own slice of the batch, so what has to sync is the full gradient — one ring all-reduce over all params, every single step.
Tensor parallel: sync one activation, at every layer
Weights are sharded across GPUs (no gradient sync needed at all), but each layer boundary needs its partial activations combined — a much smaller payload ( $H$ , not params), paid $2L$ times per step (forward and backward, at every layer).
Pipeline parallel: no collective, ever
Layers are assigned whole to different GPUs. Only the boundary between adjacent stages needs an activation handed across — point-to-point, not a ring all-reduce — so it never pays the $2(P-1)/P$ collective overhead at all.

Play

data parallel

48.0

tensor parallel

24.0

pipeline parallel

8.0

pipeline parallel moves 8 units no matter how many GPUs join — it never runs a collective, only point-to-point hops between adjacent stages

Grow the cluster and data-parallel volume grows with it, since $2(P-1)/P \to 2$ as $P$ climbs — more GPUs means more ring hops. Tensor-parallel volume grows the same way, just against a much smaller payload. Pipeline-parallel volume doesn't move at all: it was never a function of $P$ in the first place.

Worked example

The model: 2 layers, 4×4 each, 32 parameters total, 4 GPUs
$\text{params} = 4 \times 4 \times 2 = 32$ . Ring all-reduce factor at $P=4$ : $2(4-1)/4 = 1.5$ .
Data parallel
$V_{\text{data}} = 1.5 \times 32 = 48$ — the entire gradient, moved once.
Tensor parallel
Per boundary: $1.5 \times 4 = 6$ . Two boundaries per layer (forward + backward) × 2 layers: $6 \times 2 \times 2 = 24$ .
Pipeline parallel
Only $L - 1 = 1$ inter-layer boundary, forward and backward: $4 \times 1 \times 2 = 8$ — the least of the three, because it never pays the ring all-reduce factor at all.

Checkpoint

Same 2-layer, 4×4 model, now spread across 8 GPUs. What is the data-parallel communication volume?

Pick a value to try it

Summary

V_{\text{data}} = \frac{2(P-1)}{P}\cdot\text{params}

None of the three strategies is "the" way to distribute training — they're different trades. Data parallelism is simple and scales the batch, but every GPU needs the whole model to fit in memory and the gradient sync grows with parameter count. Tensor parallelism shrinks the per-GPU memory footprint but pays a communication tax at every single layer. Pipeline parallelism communicates the least, but idles GPUs during the fill-and-drain of its pipeline unless the batch is sliced into enough microbatches to keep it full. Real large-model training (Megatron-LM, DeepSpeed, GPT-NeoX) combines all three at once, tuned to the cluster's actual network topology. Having built and trained something, the next chapter turns to serving it: finding the nearest match to a query among millions of stored vectors, without comparing against every one.