Home

← All posts

Model development

Post 9 of 26

Data parallelism: one update computed on many GPUs

11 Nov 2026 13 min read

Posts in this series · 26

The first post defined one training update precisely: which tokens are targets, how their losses are averaged, how the gradient becomes a step. Everything in this track keeps that update the same while spreading its work across GPUs. The previous post showed why we need to: a large run would take years on one GPU.

Data parallelism is the first and most common answer. Every GPU holds a full copy of the model. Each processes a different slice of the batch, computes gradients for its slice, and then all GPUs average their gradients so every copy applies the same update. If done correctly, the result is the update one enormous GPU would compute on the whole batch.

What must stay identical

The guarantee rests on three conditions:

  1. Same starting weights. PyTorch’s DistributedDataParallel (DDP) copies rank 0’s parameters and buffers to every process when it is constructed, and afterwards never broadcasts parameters again: it relies on every rank applying the same update (PyTorch 2.6 source).
  2. Disjoint data. Each rank must see a different slice of the global batch, usually arranged by a sampler that shards the dataset by rank. Two ranks reading the same examples silently doubles their weight.
  3. The same objective. The averaged gradient must be the gradient of the loss you intended. This is the condition people get wrong.

Batch sizes multiply. Each GPU processes a microbatch at a time, may accumulate several microbatches before synchronizing, and there are as many copies as GPUs:

global batch=microbatch×accumulation steps×data-parallel ranks.\text{global batch} = \text{microbatch} \times \text{accumulation steps} \times \text{data-parallel ranks}.

With 4 sequences per microbatch, 8 accumulation steps and 64 ranks, the global batch is 2,048 sequences, about 8.4 million tokens at length 4,096. That is the batch the optimizer sees, and the one learning rates and schedules must be tuned for.

The normalization trap

DDP averages gradients across ranks. If each rank computes the usual mean loss over its own targets, the result is the average of per-rank means, which is not the mean over all tokens when ranks hold different numbers of targets.

Use the numbers from the first post. Rank 0 holds three target tokens with losses 1, 1 and 1; rank 1 holds one target with loss 3.

Rank 1’s single token gets three times the weight of each of rank 0’s. This is the same mistake as mean-then-divide gradient accumulation, now across GPUs.

The fix mirrors the accumulation fix. Count targets on every rank, add the counts across ranks with an all-reduce, and divide each rank’s summed loss by the global count. Because DDP then divides the summed gradients by the number of ranks, multiply by the number of ranks to undo it:

n_local = (labels != -100).sum()
n_global = n_local.clone()
dist.all_reduce(n_global)                     # total target tokens across ranks
loss = F.cross_entropy(logits.view(-1, V), labels.view(-1),
                       ignore_index=-100, reduction="sum") * world_size / n_global
loss.backward()                               # DDP averages: gradients = global token mean

Running two DDP processes on CPU with the example above confirms it: with per-rank means, the averaged gradient equals the gradient of the mean of rank means; with the global count, it matches a single process computing the token mean over the whole batch to within 3×10−83 \times 10^{-8}. Hugging Face’s Trainer has an option for this, average_tokens_across_devices, which gathers the token count across ranks and scales the loss by the number of processes, as described in the first post.

Packing sequences to a fixed length makes target counts nearly equal and hides the problem; variable-length data, fine-tuning with masked prompts and the last ragged batch of an epoch expose it.

Try it: diagnose a mismatch. A colleague’s two-GPU run gives a different loss curve from a one-GPU run with the same global batch, data order and seed. Training packs documents into fixed-length rows, but masks padding at the end of each epoch’s last batch. Which part of the curve differs, and why?

Answer. Mostly nothing, then a blip at the end of each epoch. With packed rows every rank has nearly the same number of targets, so per-rank means and the token mean almost coincide. In the last batch one rank may hold a short, padded slice with few targets; averaging per-rank means gives those few tokens extra weight, as in the 1-versus-3 example. Normalizing by the global count removes the blip.

What communication costs

Averaging gradients is an all-reduce: every rank contributes a gradient tensor and every rank ends with the sum (or average). The usual ring algorithm passes chunks around a ring of GPUs; each rank sends and receives about

2 n−1n×gradient bytes2\,\frac{n-1}{n} \times \text{gradient bytes}

for nn ranks, almost independent of nn (Thakur et al., 2005). For a 7-billion-parameter model with BF16 gradients (14 GB) on 8 GPUs, each GPU sends about 24.5 GB per step. At an illustrative 100 GB/s per GPU that is about a quarter of a second, every step.

Two techniques hide most of it:

Whether communication can be hidden depends on the ratio of compute to communication: a large model with large microbatches has plenty of backward-pass time to overlap with; a small microbatch on a slow network does not.

Where data parallelism stops

Data parallelism scales throughput but not memory: every GPU still holds the full model, its gradients and its optimizer state, the 112 GB of a 7-billion-parameter model in the previous post. It also grows the global batch with the number of GPUs, and beyond some size larger batches stop reducing the number of steps needed, so extra GPUs stop paying off. The next post removes the memory limit by sharding the state across the data-parallel ranks.

What to remember

Sources

Further reading