PyTorch Distributed Data Parallel

"They wrapped me in one line of code and told everyone I was easy. Behind the curtain I was registering hooks, packing buckets, and racing the backward pass on every step. I never complained, because that is the job."

A DistributedDataParallel Wrapper Keeping Up Appearances

By this point in the chapter you can derive data-parallel training from first principles. The gradient of an average loss is itself an average, so splitting a batch across $K$ workers and averaging their gradients reproduces the single-machine gradient exactly, a fact established back in Section 1.1. The combine step is an all-reduce, the collective built in Chapter 4, and the previous section showed how to hide its cost by overlapping it with the backward pass. What remains is to stop writing all of that yourself. PyTorch DistributedDataParallel (DDP) is the tool that packages the entire chapter into production code, and using it well is mostly a matter of understanding what it does on your behalf so you can reason about the cases where it needs help.

1. The Anatomy of a DDP Job Beginner

A DDP training script is built from five moving parts, and once you can name them the rest is detail. First, each process joins a process group so the workers can find each other and run collectives. Second, you wrap your model in DistributedDataParallel, which broadcasts the initial weights from rank 0 and registers the backward hooks that will do the gradient averaging. Third, you give the data loader a DistributedSampler so each process reads a disjoint slice of the dataset rather than all of it. Fourth, you launch the whole thing with torchrun, which spawns one process per GPU and sets the environment variables (rank, world size, master address) that the process group reads. Fifth, you guard checkpoint saving so that only rank 0 writes the file, because all replicas hold identical weights and writing from all of them at once would be wasteful at best and corrupt the file at worst.

The code below is the skeleton of essentially every single-node-multi-GPU DDP job. It needs a real multi-GPU machine and the NCCL backend to run, so we show it illustratively here and reproduce its core behavior in pure Python later in the section; read it for shape, not to execute.

# Save as train_ddp.py, launch with:  torchrun --nproc_per_node=4 train_ddp.py
import os, torch, torch.nn as nn, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def main():
    dist.init_process_group(backend="nccl")          # 1. join the K-worker group
    rank       = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])        # which GPU on this machine
    torch.cuda.set_device(local_rank)                 # pin this process to one GPU

    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])       # 2. wrap: broadcast + hooks

    sampler = DistributedSampler(train_set)           # 3. disjoint shard per rank
    loader  = DataLoader(train_set, batch_size=64, sampler=sampler)
    opt     = torch.optim.AdamW(model.parameters(), lr=3e-4)

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)                      # reshuffle shards each epoch
        for x, y in loader:
            x, y = x.cuda(local_rank), y.cuda(local_rank)
            opt.zero_grad()
            loss = loss_fn(model(x), y)
            loss.backward()                           # hooks fire all-reduce here
            opt.step()                                # every rank already averaged

        if rank == 0:                                 # 5. save once, from rank 0
            torch.save(model.module.state_dict(), "ckpt.pt")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

The striking thing about Code 15.6.1 is how little of it is distributed-systems code. There is no explicit all-reduce, no manual gradient gathering, no bucket management; loss.backward() looks exactly as it does on one machine. That is the entire point of DDP, and it is the reason it is the default tool for the data-parallel case. The synchronization that earlier sections spelled out by hand has been folded into the backward pass itself, which is exactly the overlap strategy the previous section motivated.

2. Why DDP Replaced DataParallel Beginner

PyTorch once shipped an easier-looking class called DataParallel (note: no "Distributed"). It ran in a single process and, on each step, scattered the input batch across the local GPUs, replicated the model to each, gathered the outputs back to the primary GPU, and computed the backward there. It needed no launcher and no process group, which made it tempting, and it is the wrong tool for almost every real job. Because it lives in one Python process, it is throttled by the Global Interpreter Lock, so the single driver thread becomes a bottleneck that the extra GPUs cannot relieve. It re-replicates the model to every GPU on every forward pass, which is pure overhead. It routes all outputs and the loss through one primary GPU, creating a memory and bandwidth hotspot that scales badly. And it cannot cross machine boundaries at all.

DistributedDataParallel fixes each of these by construction. It runs one process per GPU, so there is no shared interpreter lock and no single driver thread. It replicates the model exactly once, at construction, and thereafter only exchanges gradients. It never funnels activations through a primary device; each process keeps its own forward and backward entirely local and the only cross-device traffic is the gradient all-reduce. And it works identically across one machine or a thousand, because the process group abstraction does not care whether ranks share a host. The practical guidance from the PyTorch team is blunt and worth internalizing: use DistributedDataParallel, not DataParallel, even on a single machine with multiple GPUs.

3. Correctness Details DDP Quietly Enforces Intermediate

For the averaged gradient to equal the single-machine gradient, several invariants must hold, and DDP enforces or assumes each of them. Every replica must start from identical weights, which DDP guarantees by broadcasting rank 0's parameters and buffers when you wrap the model; this is why you do not need to seed every process identically yourself. Each process owns exactly one GPU, the configuration DDP is tuned for. The gradient is averaged, not merely summed, so that the effective update matches a single batch of size $K$ times the per-process batch; DDP divides the all-reduced sum by the world size for you. With $K$ processes each using a local batch of $B$, the global batch is $KB$, and the synchronized gradient is

which is exactly the mean gradient over the global batch when the shards are equal in size. Equal shard sizes matter: the DistributedSampler pads the last shards so every rank sees the same number of batches, because a rank that runs out of data early would leave the others waiting forever at the next all-reduce. The one place this exactness leaks is batch normalization, whose statistics are computed per process over only the local batch of $B$ examples, not the global $KB$. When that local batch is small, swapping the layers for SyncBatchNorm (which all-reduces the mean and variance across replicas) restores the single-machine semantics at the cost of an extra collective.

4. Gradient Accumulation and no_sync Intermediate

Sometimes the global batch you want is larger than the GPUs can hold even after splitting it $K$ ways. The standard remedy is gradient accumulation: run several forward and backward passes, summing gradients in the .grad buffers, and step the optimizer only once every few micro-batches. Under DDP this creates a subtlety. By default DDP all-reduces on every backward pass, so naive accumulation would synchronize gradients on every micro-batch, paying the communication cost many times for a single optimizer step. That is wasteful: you only need the averaged gradient on the step where you actually call optimizer.step(). DDP exposes the no_sync() context manager precisely for this. Inside it, the backward hooks still accumulate gradients locally but skip the all-reduce; you run all but the last micro-batch inside no_sync(), then run the final micro-batch outside it so a single all-reduce synchronizes the accumulated total.

# Accumulate over `accum` micro-batches, all-reduce only on the last one.
for i, (x, y) in enumerate(loader):
    is_last = (i + 1) % accum == 0
    ctx = model.no_sync() if not is_last else contextlib.nullcontext()
    with ctx:                                     # skip all-reduce until the end
        loss = loss_fn(model(x), y) / accum       # scale so the mean is correct
        loss.backward()                           # local grad accumulation
    if is_last:
        opt.step(); opt.zero_grad()               # one all-reduce already fired

This is the deep link back to communication-efficient optimization. Skipping the all-reduce on intermediate micro-batches is the same idea as the local-update schemes of Section 10.5: do more local computation between synchronizations so the network cost is amortized over more work. Gradient accumulation is the simplest, exact instance of that trade, because the local steps here are merely accumulating into one gradient rather than taking independent optimizer steps, so the result is identical to a single large batch.

5. Reproducing DDP's Contract in Pure Python Intermediate

The clearest way to demystify DDP is to build its core contract without any GPUs or NCCL. The demonstration below defines a tiny MiniDDP wrapper around a plain parameter vector. It holds $K$ replicas initialized identically (DDP's broadcast from rank 0), and on every step it computes each replica's gradient on its own data shard, then averages the gradients across replicas by walking them in fixed-size buckets and all-reducing one bucket at a time, exactly as DDP's backward hooks do. After training, it checks two things that DDP promises: that all replicas remain bit-for-bit identical, and that the synchronized trajectory matches an ordinary single-process run over the full batch.

import numpy as np

rng = np.random.default_rng(7)
N, d, K, STEPS, LR = 4096, 16, 4, 30, 0.05

X = rng.standard_normal((N, d))                  # one shared regression problem
w_true = rng.standard_normal(d)
y = X @ w_true + 0.1 * rng.standard_normal(N)
w_init = np.zeros(d)

def full_grad(w, Xb, yb):                         # mean-squared-error gradient
    n = len(yb)
    return (2.0 / n) * (Xb.T @ (Xb @ w - yb))

# Reference: ordinary single-process training over the whole batch.
w_ref = w_init.copy()
for _ in range(STEPS):
    w_ref -= LR * full_grad(w_ref, X, y)

class MiniDDP:
    """On each step, bucket the per-replica gradients and all-reduce (sum then
    divide by world size) so every replica applies the identical averaged grad,
    exactly as torch DDP does inside its backward hook."""
    def __init__(self, w0, world_size, bucket_size):
        self.world_size, self.bucket_size = world_size, bucket_size
        self.replicas = [w0.copy() for _ in range(world_size)]   # identical init

    @staticmethod
    def _all_reduce_buckets(grads, bucket_size):
        P = grads[0].size
        out = [np.empty_like(g) for g in grads]
        for start in range(0, P, bucket_size):                   # one all-reduce
            stop = min(start + bucket_size, P)                   # per bucket
            stacked = np.stack([g[start:stop] for g in grads])
            averaged = stacked.sum(axis=0) / len(grads)          # SUM / world_size
            for r in range(len(grads)):
                out[r][start:stop] = averaged
        return out

    def step(self, shards, lr):
        local = [full_grad(self.replicas[r], X[s], y[s])         # local grad only
                 for r, s in enumerate(shards)]
        synced = self._all_reduce_buckets(local, self.bucket_size)
        for r in range(self.world_size):
            self.replicas[r] -= lr * synced[r]

shards = np.array_split(np.arange(N), K)                         # disjoint shards
ddp = MiniDDP(w_init, world_size=K, bucket_size=5)               # 16 params, bk 5
for _ in range(STEPS):
    ddp.step(shards, LR)

spread = max(np.max(np.abs(ddp.replicas[r] - ddp.replicas[0])) for r in range(K))
err    = np.max(np.abs(ddp.replicas[0] - w_ref))
print("world size K              :", K)
print("steps                     :", STEPS)
print("max spread across replicas:", f"{spread:.2e}")
print("max |DDP - single-process|:", f"{err:.2e}")
print("final loss (single)       :", f"{np.mean((X @ w_ref - y) ** 2):.6f}")
print("final loss (mini-DDP)     :", f"{np.mean((X @ ddp.replicas[0] - y) ** 2):.6f}")

world size K              : 4
steps                     : 30
max spread across replicas: 0.00e+00
max |DDP - single-process|: 2.22e-16
final loss (single)       : 0.045429
final loss (mini-DDP)     : 0.045429

The two reported quantities are the two guarantees DDP makes. A spread of exactly zero across replicas means the workers never drift apart, because they apply the same averaged gradient every step from the same starting point. A difference of $2.2 \times 10^{-16}$ from the single-process run, the smallest representable gap in double precision, means the bucketing and averaging changed nothing about the answer. Real DDP adds GPUs, NCCL, and overlap with the backward pass, but the contract it enforces is precisely the one this forty-line wrapper demonstrates: many replicas, one trajectory, identical to training on one machine.

model = DDP(model, device_ids=[local_rank])   # the entire chapter, in one line
# ...then train normally; loss.backward() does the bucketed all-reduce for you.

6. When the Wrapper Is Not Enough Advanced

DDP solves the data-parallel case completely, and the data-parallel case is the one where the whole model fits on one GPU and only throughput binds. The moment the model itself no longer fits, DDP has nothing to offer, because it replicates the full model to every process. That is the boundary where this chapter ends and the next begins: sharded data parallelism and model parallelism, which split the parameters themselves across devices rather than replicating them. The collectives change too, from a single all-reduce to the reduce-scatter and all-gather pair that sharded methods rely on, a transition foreshadowed when this book first introduced collective communication. Within its lane, though, DDP is close to optimal, and the right instinct is to reach for it first and only graduate to heavier machinery when a memory ceiling, not a throughput ceiling, forces the move.

You now have the practical tool that packages this entire chapter. The exact gradient identity, the synchronous step, the ring all-reduce, and the overlap with backpropagation all live inside a single wrapper that lets you write almost-normal training code. The next section turns to Horovod, an alternative framework that brought the same all-reduce-based data parallelism to TensorFlow, Keras, and PyTorch alike, and that pioneered some of the overlap techniques DDP now uses, so that you can recognize the shared idea wearing a different library's clothes.