Data Parallelism (DP)

The Real Problem: One Batch, Too Slow

Section 11.1 made the case that a single GPU cannot hold a frontier model. But before any of the more exotic parallelism schemes — tensor, pipeline, expert, sequence — there is a simpler problem that data parallelism solves, and that you will meet on the very first day you own more than one GPU.

SGD's gradient is an average. With a batch of $B$ examples and loss $\mathcal{L}_i(\theta)$ on each example, one step of gradient descent is:

$\theta \leftarrow \theta - \eta \cdot \frac{1}{B} \sum_{i=1}^{B} \nabla \mathcal{L}_i(\theta).$

The sum is embarrassingly parallel — you can compute any partition of those $B$ terms anywhere and the result is the same as long as you add them at the end. So the moment you have a second GPU sitting next to the first, the obvious speed-up writes itself: give half the batch to each GPU, run forward and backward on the local half, then add the two gradient contributions and apply one shared update.

That idea — same model, sharded data, summed gradients — is data parallelism. It is the single most-deployed parallelism strategy in deep learning: the one you reach for first, the one every other parallelism scheme composes on top of, the one that makes 8-GPU nodes and 8 000-GPU clusters look architecturally similar from Python's point of view. And it is the reason a frontier run can process tens of millions of tokens per second instead of tens of thousands.

Intuition: Many Workers, One Identical Model

Picture a workshop. You are training a model — call it $\theta$ — and the training set is a giant bin of examples. With one worker, training is sequential: pick a handful of examples, compute the average gradient, nudge $\theta$, repeat.

Hire eight workers. Give each of them the same copy of the current model and a different handful of examples. Each worker computes their local mean gradient. They are now staring at eight different vectors pointing in eight slightly different directions — same model, different data, different gradients.

Here is the only subtle part of DP: the eight workers cannot just each apply their own gradient and walk away. If they did, they would end the step with eight different models and the workshop would split into eight uncoordinated trainings. So they pause. They announce their gradient vectors to each other. They sum them and divide by eight. Now every worker holds the same averaged direction. Every worker takes the same step. The eight copies of the model are once again identical, ready for the next iteration.

Three intuitive consequences flow from this picture:

1. DP is mathematically equivalent to single-GPU SGD on a bigger batch. Eight workers each at batch 8 is — at the gradient level — identical to one worker at batch 64. The reason this matters: you can predict what DP will do, in terms of loss curve and learning-rate sensitivity, by treating it as a batch scale-up.
2. The sync step is the bill. Forward and backward are perfectly parallel — N workers, N× the throughput, no coordination needed. The only thing that costs you is the gradient-averaging round, and that cost grows with the number of parameters in the model, not the number of examples in the batch. That is why a 70 B model on 1024 GPUs is communication-bound while a 70 M model on 1024 GPUs is compute-bound.
3. You cannot use DP to fit a bigger model. Every worker holds the entire model. Eight workers means eight identical copies of $\theta$. DP buys you more throughput, not more capacity. The moment your model stops fitting on one GPU, you need ZeRO/FSDP, tensor parallelism, or pipeline parallelism — DP alone will not save you. Sections 11.3–11.7 are the answer to that.

The Math: Mini-Batch SGD Across N Workers

Let the global mini-batch be $\mathcal{B} = \{(x_i, y_i)\}_{i=1}^{B}$ and partition it into $N$ disjoint shards $\mathcal{B}_r$ of size $B/N$ each, one per rank. The single-GPU gradient at parameter $\theta$ is:

$g(\theta) = \frac{1}{B} \sum_{i=1}^{B} \nabla_\theta \mathcal{L}(x_i, y_i; \theta).$

Rewrite the sum as a sum-over-ranks:

$g(\theta) = \frac{1}{B} \sum_{r=1}^{N} \sum_{(x_i, y_i) \in \mathcal{B}_r} \nabla_\theta \mathcal{L}(x_i, y_i; \theta) = \frac{1}{N} \sum_{r=1}^{N} g_r(\theta),$

where $g_r(\theta) = \frac{N}{B} \sum_{(x_i, y_i) \in \mathcal{B}_r} \nabla_\theta \mathcal{L}(x_i, y_i; \theta)$ is the local mean gradient computed on rank $r$. This identity is the entire mathematical content of data parallelism. The global gradient is the average of the local gradients provided each rank divides by its own shard size (so that $g_r$ is already a mean, not a sum).

The SGD update on every rank is then:

$\theta^{(t+1)} = \theta^{(t)} - \eta \cdot \frac{1}{N} \sum_{r=1}^{N} g_r(\theta^{(t)}).$

Three engineering consequences fall out of these two lines:

1. Synchrony is required. Every $g_r$ is evaluated at the same $\theta^{(t)}$. If rank 3 is one step behind, its $g_3$ is $\nabla \mathcal{L}(\cdot; \theta^{(t-1)})$ — a gradient for the wrong model. Async DP accepts this bias for throughput; sync DP refuses.
2. The communication primitive is sum, not max or min. We need $\sum_r g_r$. Then we divide locally by $N$. NCCL and other collective libraries expose this as AllReduce(SUM) — every rank ends up holding the same sum, no central server, no leader election.
3. The learning rate must scale. Replacing batch $B/N$ with batch $B$ reduces gradient variance by $\sqrt{N}$. The empirical rule (Goyal et al. 2017) is linear LR scaling: multiply $\eta$ by $N$, then warm it up linearly for the first ~5 % of steps to keep early training stable. Frontier runs (LLaMA, DeepSeek-V3) tune this empirically — the linear rule is the prior, not the answer.

AllReduce: The Sync Primitive That Makes DP Work

The math says we need $\sum_r g_r$, identical on every rank, in as few bytes per rank as possible. There are three natural ways to implement it; only one of them scales.

Naive parameter server: $O(N)$ bytes on the server

Pick one rank as the server. Every other rank sends its gradient to the server, the server adds them up, then broadcasts the sum back. Total bytes through the server: about $2 (N-1) |g|$. The server is a hot spot — every byte of every gradient flows through it, both in and out. Network and CPU at the server saturate the moment $N$ is more than a handful. This is what DistBelief used in 2012; nobody uses it for synchronous frontier training today.

Tree all-reduce: $O(\log N)$ latency

Arrange the ranks in a binary tree. Sum up the tree, broadcast down. Total bytes per rank: $2 |g| \log_2 N$ for small messages, but each link is sequential — the cost is dominated by latency, not bandwidth. Good for small payloads (latencies of microseconds), bad for whole-model gradient buffers.

Ring all-reduce: $2(N-1)/N \cdot |g|$ bytes per rank

This is the algorithm every modern DP framework uses, because it has a remarkable property: the bytes each rank must send and receive approach a constant (about $2 |g|$) as $N$ grows. Doubling the cluster does not double the traffic per GPU; it barely moves it. The algorithm has two phases, each running for $N - 1$ rounds:

1. Scatter-reduce. Split the gradient buffer on every rank into $N$ equal chunks. In round $s$, rank $r$ sends chunk index $(r - s) \bmod N$ to rank $r+1$, which adds it onto its own copy of that chunk. After $N - 1$ rounds, rank $r$ owns the full sum of exactly one chunk. Bytes per rank: $(N-1)/N \cdot |g|$.
2. All-gather. Same ring, no addition — each fully-summed chunk walks $N - 1$ more hops, overwriting each rank's stale copy. After another $N - 1$ rounds, every rank holds every fully-summed chunk. Bytes per rank: another $(N-1)/N \cdot |g|$.

Total bytes per rank: $B_{\text{ring}} = \frac{2(N-1)}{N} \cdot |g| \xrightarrow{N \to \infty} 2 |g|.$ Compare with the naive PS at $2(N-1) |g|$ on a single link. For $N = 1024$ ranks and a 70 B BF16 model ($|g| \approx 140 \, \text{GB}$), the ring sends about $280 \, \text{GB}$ per rank; the parameter server's central node would have to absorb $140 \cdot 1023 \approx 143 \, \text{TB}$. At 400 GB/s of network bandwidth (typical IB on a frontier cluster), that single link would need an hour per step.

Manual Numerical Walkthrough

Two by-hand calculations: one round of ring all-reduce on 4 GPUs with toy 4-element gradient vectors, and a wall-clock estimate of DP scaling on a 70 B model across 256 H100s.

Visualizing the Ring and the Scaling Wall

Two widgets. The first animates ring all-reduce step by step — watch the cyan chunks accumulate into the green fully-reduced chunks, then propagate around the ring in the all-gather phase. The second is a scaling calculator: pick a model, GPU count, network bandwidth, and precision, and see when DP's communication cost overtakes compute.

Plain Python: DP and Ring-AllReduce From Scratch

Before reaching for PyTorch DDP, let us build the whole thing in plain NumPy — including the ring all-reduce, by hand, no MPI. The toy model is a single linear layer; the toy "cluster" is a list of arrays in one process. The mechanics are byte-identical to what runs across 2048 H800s.

1import numpy as np
2
3# Tiny toy model:  y = w * x + b   with a squared-error loss.
4# We will simulate N=4 "GPUs" with disjoint mini-batches and run
5# one step of synchronous data-parallel SGD by hand.
6
7np.random.seed(0)
8N_GPUS, BATCH_PER_GPU, LR = 4, 8, 0.1
9
10def make_shard(gpu_id):
11    """Each rank gets a different slice of the global batch."""
12    x = np.random.randn(BATCH_PER_GPU).astype(np.float32) + gpu_id
13    y = 3.0 * x + 0.5 + 0.1 * np.random.randn(BATCH_PER_GPU)
14    return x, y
15
16# (A) Every rank starts with the SAME weights -- this is the DP invariant.
17w, b = np.float32(0.0), np.float32(0.0)
18
19# (B) Forward + backward on the LOCAL shard only. No comm yet.
20local_grads = []
21for g in range(N_GPUS):
22    x, y = make_shard(g)
23    y_hat = w * x + b
24    err   = y_hat - y                              # shape (B,)
25    g_w = (2.0 / BATCH_PER_GPU) * (err * x).sum()  # local mean gradient
26    g_b = (2.0 / BATCH_PER_GPU) * err.sum()
27    local_grads.append(np.array([g_w, g_b], dtype=np.float32))
28
29# (C) Ring all-reduce, implemented by hand in two phases.
30def ring_all_reduce(buffers):
31    """In-place ring all-reduce: every buffer ends up holding the sum."""
32    N = len(buffers)
33    bufs = [b.copy() for b in buffers]
34    chunk_size = bufs[0].size // N
35    chunks = [b.reshape(N, chunk_size) for b in bufs]
36
37    # Phase 1: scatter-reduce. After N-1 rounds, rank r owns the full
38    # sum of chunk (r - (N-1)) mod N.
39    for s in range(N - 1):
40        for r in range(N):
41            send_idx = (r - s)     % N
42            recv_idx = (r - s - 1) % N
43            dst = (r + 1) % N
44            chunks[dst][recv_idx] += chunks[r][send_idx]
45
46    # Phase 2: all-gather. Each fully-reduced chunk walks the ring.
47    for s in range(N - 1):
48        for r in range(N):
49            send_idx = (r - s + 1) % N
50            dst = (r + 1) % N
51            chunks[dst][send_idx] = chunks[r][send_idx]
52
53    return [c.reshape(buffers[0].shape) for c in chunks]
54
55# (D) Sync the gradients (sum), then average and apply the update.
56synced = ring_all_reduce(local_grads)
57mean_grad = synced[0] / N_GPUS                     # all ranks identical
58w -= LR * mean_grad[0]
59b -= LR * mean_grad[1]
60
61# (E) Sanity check vs the "ground truth" big-batch gradient.
62xs, ys = zip(*[make_shard(g) for g in range(N_GPUS)])
63xs, ys = np.concatenate(xs), np.concatenate(ys)
64yh = 0.0 * xs + 0.0
65big_grad_w = (2.0 / xs.size) * ((yh - ys) * xs).sum()
66big_grad_b = (2.0 / xs.size) * (yh - ys).sum()
67print(f"DP averaged grad : ({mean_grad[0]:+.4f}, {mean_grad[1]:+.4f})")
68print(f"Big-batch grad   : ({big_grad_w:+.4f}, {big_grad_b:+.4f})")
69print(f"Updated weights  : w = {w:+.4f}   b = {b:+.4f}")

The output of that script ends with two lines: DP averaged grad ≈ big-batch grad — the mathematical identity at the heart of DP, verified to floating point. If those two numbers ever disagree, DP is broken; finding why they disagree is one of the most common distributed-training debugging exercises (mis-sized shards, missing set_epoch, non-deterministic kernels, dropout-RNG drift).

PyTorch: Real DDP on a GPU Cluster

The plain-Python version was the mechanism. The PyTorch version is the same mechanism, except: the ring all-reduce lives in NCCL kernel code, the bucketing is automatic, the overlap with backward is hooked into autograd. Two extra lines around your existing training loop and you go from one GPU to one thousand.

1import os
2import torch
3import torch.distributed as dist
4import torch.nn as nn
5from torch.nn.parallel import DistributedDataParallel as DDP
6from torch.utils.data import DataLoader, DistributedSampler
7
8def setup():
9    # torchrun sets these env vars for every process.
10    dist.init_process_group(backend="nccl")          # H100/H800: nccl
11    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
12
13def train_one_step(model, batch, optimizer, criterion):
14    x, y = batch[0].cuda(non_blocking=True), batch[1].cuda(non_blocking=True)
15    optimizer.zero_grad(set_to_none=True)
16    logits = model(x)                                  # forward through DDP
17    loss   = criterion(logits, y)                       # local loss
18    loss.backward()                                     # AllReduce runs HERE,
19                                                        # bucketed & overlapped
20                                                        # with the backward pass.
21    optimizer.step()                                    # identical update on every rank
22
23def main():
24    setup()
25    rank = dist.get_rank()
26    world = dist.get_world_size()
27
28    model = MyTransformer().cuda()                      # SAME init via broadcast
29    model = DDP(model, device_ids=[int(os.environ["LOCAL_RANK"])],
30                gradient_as_bucket_view=True,           # zero-copy bucket views
31                static_graph=True)                      # enables bucket reuse
32    optim = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
33    crit  = nn.CrossEntropyLoss()
34
35    sampler = DistributedSampler(dataset, num_replicas=world, rank=rank,
36                                  shuffle=True, drop_last=True)
37    loader  = DataLoader(dataset, batch_size=4, sampler=sampler,
38                          num_workers=2, pin_memory=True, persistent_workers=True)
39
40    for epoch in range(num_epochs):
41        sampler.set_epoch(epoch)                        # reshuffle per epoch
42        for step, batch in enumerate(loader):
43            train_one_step(model, batch, optim, crit)
44            if rank == 0 and step % 100 == 0:
45                print(f"step {step}  loss {loss.item():.4f}")
46
47if __name__ == "__main__":
48    main()
49# Launch:  torchrun --nproc_per_node=8 --nnodes=4 train.py

Three numbers worth memorising for any DDP run:

At Massive Scale: Where Pure DP Falls Off

DP's simplicity makes it the universal first choice; the same simplicity dictates its ceiling. Three forces push frontier runs past pure DP and into the hybrid parallelism stack that the rest of Chapter 11 builds.

1. Memory: weights, gradients, master, Adam. For a 70 B model in BF16, every rank holds $140 \, \text{GB}$ of weights, the same again in gradients, $280 \, \text{GB}$ in FP32 master weights, and $560 \, \text{GB}$ in Adam moments — a total of $\approx 1.12 \, \text{TB}$ per rank, far more than any single GPU's 80–192 GB of HBM. Pure DP cannot run this model at all. ZeRO-1 shards the optimiser state, ZeRO-2 also shards the gradient, ZeRO-3 (and PyTorch's FSDP) shards the parameters themselves — at the cost of extra all-gather operations every forward and backward. Section 11.4 builds this in full.
2. Communication: $2|g|$ per step, every step. Even when memory fits, the all-reduce volume is fixed by model size. Doubling GPUs doubles compute throughput, but per-GPU all-reduce time barely moves — so the gap between compute and comm narrows until comm dominates. At that point you must either compress the gradients (PowerSGD, 1-bit Adam, FP8 collectives) or stop having every rank hold every parameter. Tensor parallelism (Section 11.5) attacks the latter — each layer's matrix is split across a few ranks, turning a big AllReduce into smaller in-shard reductions.
3. Pipeline bubbles for very deep models. Once a model is too big for one node, you can also slice it by layer (pipeline parallelism, Section 11.6). DP composes with that: across a 256-stage pipeline, each pipeline replica is its own data-parallel group running synchronous DP. Frontier runs are almost always "DP × TP × PP × EP" — a Cartesian product of parallelism dimensions with DP as the outermost loop.

What DeepSeek-V3 actually does

DeepSeek-V3 uses DP at the outermost level across roughly 64-rank data-parallel groups, with each group running a complex inner configuration of expert parallelism, sequence parallelism, and pipeline parallelism that the rest of the chapter unpacks. The DP all-reduce on the optimiser-state side is sharded ZeRO-style so that no rank holds the whole optimiser. The point is not that DP is obsolete at frontier scale — it is that it is one axis of a 4-D parallelism layout, and its specific contribution is what we characterised in this section: synchronous, replicated-model, ring-all-reduce gradient averaging.

Engineering Reality: Bucketing, Overlap, and the Tricks That Matter

Four engineering wins separate a textbook DP implementation from what production frameworks actually do. Each is worth roughly a 2× throughput improvement on a real cluster; combined they are the difference between "multi-GPU works" and "multi-GPU is the platform."

1. Gradient bucketing. Issuing one all-reduce per parameter is a disaster: a 7 B model has hundreds of thousands of parameters, each all-reduce has microseconds of kernel-launch overhead, and the total per-step overhead dwarfs the actual transfer. DDP bundles consecutive parameters into ~25 MB buckets and issues one all-reduce per bucket. The bucket size is tuned so the per-bucket overhead is small but the bucket itself is large enough to be bandwidth-bound. 25 MB on a 400 GB/s link is 60 µs of actual transfer; the launch overhead is roughly 10 µs — a healthy 6:1 ratio.
2. Backward-pass overlap. Buckets fire all-reduce as soon as they fill, while the backward pass is still computing earlier layers' gradients. On a transformer, this means the all-reduce on layer 80's gradients is in flight while CUDA is still computing the gradient for layer 60. By the time backward returns, almost all communication is done — DDP's step time is $\max(t_{\text{compute}}, t_{\text{comm}})$, not their sum. This is the single biggest reason DP scales further than its naive analysis suggests.
3. Mixed-precision collectives. Sending gradients in BF16 instead of FP32 halves the wire volume; sending them in FP8 (after appropriate scaling, see Section 10.3) halves it again. Modern NCCL exposes typed all-reduce, and frameworks (DeepSpeed, Megatron-LM, the DeepSeek stack) cast on the wire and accumulate in higher precision on receive. For a 70 B model, BF16 → FP8 collectives save 70 GB per step per rank. On 1024 ranks that is 70 TB of network traffic per step removed.
4. Sharded optimiser state (ZeRO-1). Even before touching the full ZeRO-3 / FSDP redesign, ZeRO-1 splits Adam's $m, v$ moments across DP ranks: rank $r$ owns the master copy of parameter shard $r$. The forward pass still uses replicated BF16 weights, but optimiser memory drops by $N_{\text{DP}}$. For a 70 B model on 64 DP ranks, that frees $\approx 550 \, \text{GB}$ of HBM per rank — enough to enable longer sequences or larger batches. Section 11.4 explains the full sharding hierarchy.

The failure modes you will actually see

The deepest lesson of this section. Data parallelism is the parallelism strategy with the strongest mathematical guarantees — it is provably equivalent to single-GPU SGD on a bigger batch, end of story. Every other parallelism technique we will meet — tensor, pipeline, expert, sequence — has a numerical caveat: a slightly different forward pass, a slightly different gradient, a slightly different floating-point accumulation order. DP has none. The price you pay for that purity is the per-step gradient all-reduce — bandwidth proportional to model size, paid every single step of training. The rest of Chapter 11 is the engineering of how to pay that price as cheaply as possible (FSDP, ZeRO) or how to avoid paying it for the entire model at once (TP, PP, EP). But DP is the baseline every other technique is measured against, and the one every other technique composes on top of.