Splitting the work

Parts 1 and 2 were about one device. This part adds more of them in the simplest way available.

The identity

The gradient of a sum is the sum of gradients. So the batch can be split, each device can compute on its slice, and the results added.

$$g = \frac{1}{N}\sum_{i=1}^{N} g_i$$

Every rank starts identical, computes its own gradient, all-reduces, and steps. Because they all apply the same averaged gradient to the same weights, they stay identical — and nothing else needs synchronising. One tensor, once per step.

This is gradient accumulation with the micro-batches running concurrently instead of sequentially, which is a nice way to see that they are the same idea against different resources: accumulation spends time to save memory, data parallelism spends devices.

Why the collective does not get worse

The thing I did not understand for an embarrassingly long time is why this scales. Surely averaging across a thousand ranks costs a thousand times as much as across two?

It does not, and the ring algorithm is why. Arrange the ranks in a circle, split the gradient into $N$ chunks. Each rank sends one chunk to its neighbour and adds what it receives; after $N-1$ steps every rank holds the complete sum for exactly one chunk. Then circulate those around, another $N-1$ steps.

$$2\,\frac{N-1}{N}\cdot S \;\longrightarrow\; 2S$$

Twice the gradient size per rank, approaching a constant. The latency term grows with $N$ — $2(N-1)$ hops — which is why small tensors are the problem and why frameworks bucket hundreds of parameter tensors into a few large messages.

That flat bandwidth cost is the reason data parallelism is the outermost layer of every distributed configuration. It is the only strategy whose communication does not degrade with scale.

The waste

Every rank holds a full copy of the parameters, gradients and optimiser state. That is part 1’s 16 bytes per parameter, replicated $N$ times, to store one model’s worth of information.

For 1024 ranks that is 1023 redundant copies of the optimiser state — which part 1 established is two thirds of the total.

ZeRO removes the redundancy in three stages. Shard the optimiser state; then the gradients; then the parameters themselves, gathering each layer’s weights just before they are needed and releasing them after. Stage 3 is $16P/N$ per device rather than $16P$.

$$\text{stage 1--2: unchanged bytes} \qquad \text{stage 3: } +\ \text{all-gather of parameters, per layer}$$

Stages 1 and 2 are close to free, because of the identity above — the all-reduce was already going to happen and is simply decomposed. Stage 3 genuinely adds traffic, since parameters must be gathered on the way forward and again on the way back, and whether that is affordable is a question about the interconnect.

The thing that has to stay identical

Everything above assumes the ranks remain in lockstep, and the ways that quietly fails are worth listing because none of them raises an error.

A different data order. DistributedSampler needs set_epoch called every epoch, or every rank uses the same permutation for the whole run. The model sees one fixed ordering forever. No error, slightly worse model.

A rank that skips a collective. Any branch that differs between ranks — an empty micro-batch, a conditional early exit — deadlocks the job at the next barrier. The symptom is a hang with no output, which is a worse failure than a crash because it looks like slowness.

Clipping in the wrong place. Gradient clipping before the all-reduce clips a partial gradient against a global threshold, so each rank scales differently and the averaged direction is wrong. Trains fine, produces a worse model.

Where it stops

Two limits. Data parallelism requires the model to fit on one device, which is what ZeRO stage 3 and the model-splitting strategies address. And the effective batch grows with the rank count, so past the critical batch size the extra devices are buying gradient variance reduction that the optimiser cannot use.

The last part is about what a real configuration looks like when four of these run at once, and about the failure that costs half a cluster’s throughput without producing a single error message.