Vol. XIV · Parallelism Plate 4 of 11data parallelism, replicas, gradient all-reduce, overlap, global batch size, scaling efficiencyCE

Data Parallelism.

Copy the model to every device, split the batch, average the gradients — the easiest way to use more chips.

In data parallelism each of n devices holds a full replica of the model and its optimiser state, takes a different slice of the global batch, computes gradients, and then all-reduces them so every replica applies the same averaged update. Nothing about the model changes; the global batch is n times the per-device batch. Because gradients for later layers are ready first during the backward pass, their all-reduce can start while earlier layers are still being computed, hiding most of the communication.

Two things eventually stop it. First, the all-reduce time — about 2 × 2N bytes ÷ link bandwidth, regardless of n — is fixed while each device’s compute per step shrinks as n grows at a fixed global batch; once communication exceeds compute, it can no longer be hidden. Second, the global batch cannot grow without limit: beyond the critical batch size (plate XI.2) extra examples buy little. And plain data parallelism needs the whole model and its 16 bytes per parameter on every device. The workshop finds where the first limit bites.

Below: compute time and gradient all-reduce time per step against the number of data-parallel replicas, for the chosen model and global batch, over the chosen network; and the resulting scaling efficiency.

What you will learn
Workshop
Per-step compute at 45% of H100 peak for the chosen global batch, divided over n replicas, and the ring all-reduce of BF16 gradients (2N bytes) over a 400 Gb/s network (50 GB/s) or NVLink (450 GB/s). Efficiency assumes communication overlaps perfectly with compute until it is longer.

Linear until it is not

For Llama 3 8B with a global batch of 4 million tokens over a 400 Gb/s network, each step’s compute falls as replicas are added while the gradient all-reduce stays near 0.56 seconds; they cross only at about 1,024 replicas, by which point each replica processes 4,096 tokens a step — half of one 8 192-token sequence. Over NVLink, at 1 024 replicas efficiency is still 100%. With a 1-million-token batch the crossing comes sooner, at 256 replicas. For a model this size the batch runs out before the network does: data parallelism scales well as long as each replica has enough work per step to hide the all-reduce.

But the model must fit. Llama 3 8B’s model states are 128.5 GB per replica — more than one H100 holds — so even an 8-billion-parameter model cannot use plain data parallelism on H100s without sharding its states (plate XIV.5). And the global batch cannot keep growing: past the critical batch more replicas buy little (plate XI.2). Data parallelism is therefore the outermost layer of every large layout — Llama 3 405B used 64 to 128 replicas — wrapped around the model-splitting schemes of the next plates.

Gradient compression and lower-precision all-reduce can cut communication further, but they change the arithmetic of the update and are used cautiously; overlapping and bucketing are the standard tools.

Terms

Data parallelism
Replicating the model on every device, splitting the batch, and averaging gradients.
Replica
One full copy of the model and its optimiser state.
Global batch
The tokens summed over all replicas in one step.
Overlap
Communicating gradients for finished layers while the backward pass continues.
Data parallelism scales almost linearly while each replica’s compute hides the gradient all-reduce — and only if the whole model fits on each replica.

Sources

  1. Li et al., “PyTorch Distributed: Experiences on Accelerating Data Parallel Training”, VLDB 2020 (gradient bucketing and overlap).
  2. McCandlish et al., “An Empirical Model of Large-Batch Training”, arXiv:1812.06162, 2018.