Data Parallelism and All-Reduce

Map, then reduce

To a data engineer this is familiar. Each GPU (a mapper) takes its shard of the batch, runs the forward and backward passes, and produces a full set of gradients. Then all GPUs reduce: they average those gradients, and every GPU applies the same averaged update. Because the average of per-shard gradients equals the gradient of the whole batch, the result is exactly what one enormous GPU would have computed.

Tiny example. Four GPUs each process 4 sequences of 4,096 tokens: 16 sequences and 65,536 tokens per step. Each GPU's gradient for some weight is 0.8, 1.2, 1.0 and 1.0; the all-reduce gives every GPU the average, 1.0, and every copy takes the identical step.

$$\text{global batch} = \text{micro-batch} \times \text{accumulation steps} \times \text{data-parallel GPUs}$$

All-reduce

An all-reduce is a collective operation: every participant contributes a vector and every participant receives the sum (or average). Efficient implementations pass pieces around a ring or tree so that each GPU sends and receives roughly twice the size of the gradients however many GPUs there are, which is why the pattern scales to thousands.

Data parallelSharded (ZeRO-3 / FSDP)Tensor parallelPipeline parallel

Each GPU holds

A full copy of the model, and a different slice of the batch.

What they exchange

Once per step, every GPU averages its gradients with all the others (an all-reduce), so all copies stay identical.

Used

Always, as the outermost layer of parallelism: it is how a run uses thousands of GPUs.

The catch

The whole model, its gradients and optimizer states must fit on every GPU.

Real runs combine them. Llama 3 405B used tensor parallelism across the 8 NVLink-connected GPUs of each server, 16 pipeline stages, and sharded data parallelism (FSDP) across the rest, plus a fourth kind, context parallelism, for very long sequences: up to 16,384 GPUs.

The catches

Memory. Every GPU holds a full copy of the weights, gradients and optimizer states: 16 bytes per parameter. For anything beyond a few billion parameters that doesn't fit, which is what sharded training fixes.

Batch size. Adding GPUs adds to the global batch. Goyal and colleagues showed that very large batches can match small-batch accuracy if the learning rate is scaled with the batch size and warmed up gradually Established, but there are limits beyond which bigger batches stop helping.

Lockstep. Every GPU waits for the slowest at each all-reduce, and in synchronous training a single GPU failure may require restarting the entire job Established, a theme of training stability.