DDP Math — Distributed Data Parallel

The setup: replicate the model, split the batch

DDP starts from a deliberately simple picture. On each of N GPUs you place an identical replica of the model — same parameters θ, byte-for-byte. A global minibatch of B examples is partitioned into N disjoint microbatches of size b = B / N, and replica r gets microbatch r. Every replica runs its own forward and backward pass on its own data, in parallel, with no communication at all.

The result of that local backward pass is a local gradient g_r = ∇θ L_r, where L_r is the mean loss over replica r’s b examples. The N replicas now hold N different gradients, because they saw different data. If each applied its own, the copies would diverge into N different models. The single job of DDP’s communication step is to reconcile those N gradients into one shared gradient that every replica applies, keeping all copies in lockstep forever.

Why averaging gradients = training on the union batch

The claim that makes DDP correct is exact, not hand-wavy. Write the loss over the full union batch of B = N·b examples as the mean per-example loss:

L(θ) = (1/B) Σ_{i=1..B} ℓ_i(θ)
        = (1/(N·b)) Σ_{r=1..N} Σ_{j in batch_r} ℓ_j(θ)

Gradients are linear — the gradient of a sum is the sum of gradients — so pull the outer sum out and regroup:

∇L = (1/(N·b)) Σ_r Σ_j ∇ℓ_j
     = (1/N) Σ_r [ (1/b) Σ_j ∇ℓ_j ]
     = (1/N) Σ_r g_r          (since g_r = (1/b) Σ_j ∇ℓ_j)

So the plain average of the local gradients, (1/N) Σ_r g_r, equals ∇L, the true gradient of the mean loss over all B examples. Averaging is not an estimate of the union-batch gradient — it is the union-batch gradient. This is why DDP training is mathematically equivalent to single-GPU training on a batch of size B, and why the loss curves line up. The one caveat: the loss must be a mean (not a sum) and every microbatch must be the same size b, or the per-replica weights 1/N are wrong and the average silently biases toward the larger shard.

All-reduce and the ring algorithm

The primitive that turns N local gradients into one shared average is all-reduce: it combines one array per device element-wise (here, sum) and leaves the same reduced array on every device. Fold in a 1/N and each replica holds the identical averaged gradient with no separate broadcast and no central parameter server — a symmetry that keeps the replicas bit-identical. The only question is cost: a naive ‘send everything to GPU 0, add, send it back’ scheme makes that one GPU a bottleneck moving (N-1)·D data. The ring algorithm removes it.

Ring all-reduce arranges the N GPUs in a logical ring and splits each device’s gradient array of D elements into N chunks of D/N. It runs in two phases, each of N-1 steps.

Phase 1 — reduce-scatter. In each step, every GPU sends one chunk to its right neighbour and receives a chunk from its left, adding the received chunk into its own copy. After N-1 steps, each GPU holds one chunk that has accumulated the sum from all N devices — the fully reduced chunks are scattered one per GPU. Phase 2 — all-gather. Now every GPU circulates its completed chunk around the ring; after another N-1 steps, every GPU has every fully-reduced chunk, i.e. the complete reduced array.

No single GPU ever handles more than a D/N chunk per step, and every link is busy in both directions at once — no node is a hub, bandwidth is used symmetrically. This is the algorithm NCCL uses on NVLink and InfiniBand.

The 2(N-1)/N communication cost

Count the bytes each GPU moves. In reduce-scatter, a GPU sends one D/N chunk in each of N-1 steps, for (N-1)/N · D elements sent. All-gather is the mirror image: another (N-1)/N · D. Total data sent per GPU:

V_send = 2 · (N-1)/N · D    (and the same amount received)

Stare at the factor 2(N-1)/N. As N grows it approaches 2 and stays there: 1.5 at 4 GPUs, 1.75 at 8, ~1.97 at 64, ~2.00 at 1000. The per-GPU transfer is essentially independent of the number of GPUs — it never exceeds 2D, versus the naive tree-to-root scheme where the root absorbs (N-1)·D, growing linearly. Ring all-reduce is in fact bandwidth-optimal: 2(N-1)/N · D is the least data any all-reduce can move per node, since each node must both export its contribution and import the final result. The price for large N is latency — the 2(N-1) sequential steps grow with N, so many tiny messages hurt, which is what bucketing fixes.

A worked example

Take a 1-billion-parameter model, gradients in fp32 (4 bytes each), on N = 8 GPUs joined by NVLink at roughly 200 GB/s of usable bandwidth per GPU.

D            = 1e9 params × 4 bytes   = 4.0 GB of gradient
factor       = 2 · (N-1)/N = 2 · 7/8      = 1.75
V_send/GPU   = 1.75 × 4.0 GB              = 7.0 GB
time         = 7.0 GB / 200 GB/s            ≈ 35 ms per all-reduce

Every GPU ships 7 GB and the whole gradient is averaged in about 35 ms. Scale to N = 64 and the ring factor rises only to 2·63/64 ≈ 1.97, so per-GPU transfer barely moves to ~7.9 GB — whereas a naive gather-to-one-GPU root would face 63 × 4 = 252 GB. That flatness is the whole reason ring all-reduce scales: an 8-GPU node and a 64-GPU cluster do nearly the same per-GPU work for the same model.

Bucketing and overlapping all-reduce with backward

Thirty-five milliseconds of communication per step would be crippling if the GPUs sat idle during it. They do not, thanks to a timing coincidence: backpropagation produces gradients layer by layer, from the output back to the input. The moment the last layer’s gradients are ready, they can be all-reduced while the backward pass still churns through earlier layers — communication overlaps computation instead of following it.

Firing one all-reduce per parameter tensor, though, would launch thousands of tiny messages, where the ring’s per-message latency (those 2(N-1) steps) dominates. So PyTorch DDP groups gradients into buckets — contiguous blocks of roughly 25 MB by default. Hooks detect when the last gradient in a bucket is ready and kick off one asynchronous all-reduce for the whole bucket, while backward keeps filling and dispatching the next one — so several buckets are in flight while the GPU still computes. Bucketing trades a little granularity for far fewer, fatter messages; overlap is close to free. Together they hide most of the communication under compute, so a well-tuned DDP step costs little more than the local backward pass alone.

Effective batch = local x N, and the learning rate

Because averaging gradients is exactly the union-batch gradient, DDP’s effective batch size is the sum of the local microbatches:

B_effective = b_local × N   (× grad_accumulation_steps, if used)

Eight GPUs at a local microbatch of 32 train as if on batches of 256. This is the main lever DDP gives you — near-linear throughput scaling — but it changes the optimization problem, so hyperparameters must move with it. The linear scaling rule is the standard start: multiply the effective batch by k and multiply the learning rate by k too, because a larger batch gives a lower-variance gradient estimate that tolerates a bigger step. The rule breaks down at very large batches (the step turns too aggressive and destabilizes), so it is paired with a learning-rate warmup that ramps the LR up over the first few hundred iterations. The mental model: adding GPUs in DDP grows the batch, not the step count, so the same epochs take fewer, larger, better-estimated steps at a higher learning rate.

The catch: DDP replicates everything

DDP is fast, but it is memory-blind: every one of the N GPUs stores a complete, redundant copy of the entire training state — nothing is sharded. For mixed-precision Adam, the classic per-parameter accounting is stark:

fp16 parameters            2 bytes
fp16 gradients             2 bytes
fp32 master parameters     4 bytes   (optimizer state)
fp32 Adam momentum (m)     4 bytes   (optimizer state)
fp32 Adam variance (v)     4 bytes   (optimizer state)
-----------------------------------
total                     16 bytes per parameter, on EVERY GPU

A 1-billion-parameter model needs ~16 GB of state, and DDP holds that full 16 GB on all eight GPUs — 128 GB of aggregate memory to store what is really only 16 GB of unique information, replicated eight times. The gradient communication is efficient, but the storage is pure redundancy: the 12 bytes of optimizer state per parameter are identical on every replica and never need to be. This is the wall DDP hits — not compute, not communication, but model-plus-optimizer state that will not fit on one GPU once models grow past a few billion parameters.

Why ZeRO and FSDP exist

The redundancy above is exactly what ZeRO (Zero Redundancy Optimizer) and its PyTorch cousin FSDP (Fully Sharded Data Parallel) attack. They keep DDP’s data-parallel semantics — different microbatch per GPU, gradients averaged into the same union-batch update — but they stop replicating the state. Instead they shard it across the N GPUs so each holds only a 1/N slice, cutting per-GPU state memory from 16 bytes/param toward 16/N.

ZeRO does this in escalating stages: stage 1 shards the optimizer states (the fattest 12 bytes), stage 2 also shards the gradients, and stage 3 (equivalently FSDP) shards the parameters themselves, materializing a layer’s full weights only transiently — via an all-gather — when that layer runs, then discarding them. The trade is more communication: gathering sharded parameters on the fly means ZeRO-3/FSDP move more data than DDP’s single gradient all-reduce per step. That is the fundamental exchange — DDP spends memory to save communication; FSDP spends communication to save memory — and it is why DDP stays the right, simple default whenever the full state fits on one GPU, with sharding reserved for the models that do not.