Data Parallelism & ZeRO

We're now on the training side of the book, and chapter 2 already delivered the verdict: an 8B model needs 120 GiB just for its fixed state, so it does not fit on one 80 GiB GPU, and neither does anything larger. Every technique from here to chapter 11 is a way of splitting something across devices. The four things you can split — the batch, the tensors, the layers, the sequence — give the four parallelism strategies, and this chapter is the first and simplest: split the batch.

But "simplest" hides the entire cost model of distributed training, because splitting the batch means the gradients have to be reconciled every step, and reconciling gradients means moving them over the slow links from chapter 1. So we have to understand the communication before the memory savings make sense.

Data parallelism, and why the obvious implementation is bad#

Give every GPU a full copy of the model. Feed each a different slice of the batch. Each computes gradients on its slice; you average the gradients across all GPUs; everyone applies the same averaged update and stays in sync. That's data parallelism, and the averaging step is the whole story.

PyTorch's DataParallel (single process, multiple threads, one GPU as coordinator) is the version you should never use — the Python GIL serializes the threads, and the coordinator GPU becomes a scatter/gather bottleneck that gets worse with device count.

DistributedDataParallel (DDP) is the real one: one process per GPU, no GIL contention, and — crucially — the gradient averaging is done with a ring all-reduce that has no coordinator and no central bottleneck. DDP also overlaps that communication with the backward pass: as soon as a layer's gradients are ready, they start reducing while earlier layers are still computing. Gradients are bucketed (grouped into ~25 MB chunks) so the reduction launches on whole buckets rather than per-tensor. Done well, most of the communication hides entirely behind compute.

Understanding why the ring is optimal is worth the detour, because the cost formula it produces is the one every later chapter's communication is measured against.

Deriving the ring all-reduce cost#

You have $N$ GPUs, each holding a gradient vector of $X$ bytes, and you want every GPU to end up holding the sum. The naive approach — everyone sends to one GPU, which sums and broadcasts back — moves $O(NX)$ bytes through one link. The ring makes every link carry the same load and finishes in bytes independent of $N$ per GPU.

Arrange the GPUs in a logical ring, each with a left and right neighbor. Split each GPU's vector into $N$ chunks. Then two phases:

Reduce-scatter ($N-1$ steps). In step $k$, GPU $i$ sends one chunk to its right neighbor and receives a different chunk from its left, adding what it receives into its own copy. Choreographed correctly, after $N-1$ steps each GPU holds the fully-summed value of one chunk — a different chunk on each GPU. Bytes sent per GPU: $(N-1) \times X/N$.

All-gather ($N-1$ steps). Now propagate those completed chunks around the ring so everyone gets all of them. Same pattern, same cost: $(N-1) \times X/N$.

Total bytes each GPU sends:

As $N \to \infty$ this approaches $2X$, a constant. Doubling your GPU count doesn't increase per-GPU communication volume — it's the property that makes DDP scale to thousands of GPUs. Bandwidth stays constant; only latency grows, linearly in the number of steps, which is why very large rings eventually prefer tree or hierarchical collectives that trade bandwidth for fewer hops.NCCL picks the algorithm (ring, tree, or hybrid) per message size and topology automatically. The ring formula is the one to keep in your head; the others are refinements on the latency term.

The problem DDP doesn't solve#

DDP splits the batch, which shrinks the activation memory per GPU. It does nothing for the fixed state — every GPU still holds a full copy of the parameters, gradients, and optimizer state. That's the 16 bytes/param from chapter 2, replicated $N$ times across $N$ GPUs. For an 8B model that's 120 GiB per GPU, and DDP alone can never train it, no matter how many GPUs you add.

The replication is pure waste. Every GPU stores identical optimizer state it mostly doesn't need at any given moment. ZeRO — the Zero Redundancy Optimizer — deletes that redundancy by sharding the fixed state across the data-parallel group instead of replicating it, and reconstructing pieces on demand.Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020. PyTorch's FSDP is the same idea with a different API; the accounting below applies to both. It comes in three stages, and the right way to understand them is to track both memory and communication at each one, because that trade is the whole decision.

ZeRO, stage by stage#

Recall the fixed budget per parameter under mixed-precision Adam: 2 bytes bf16 weights + 2 bytes bf16 gradient + 12 bytes optimizer state (fp32 master + two fp32 moments) = 16. The 12 bytes of optimizer state is three-quarters of it, which tells you where to start.

Stage 1 — shard the optimizer state. Each GPU keeps the full weights and gradients but only $1/N$ of the optimizer state. Since the optimizer step for a parameter only needs that parameter's state, each GPU updates its own shard of the weights, then an all-gather shares the updated weights. Memory drops from 16 to $4 + 12/N$ bytes/param. Communication: unchanged from DDP. This stage is close to free, and there is essentially no reason not to use it.

Stage 2 — also shard the gradients. A GPU only needs the full gradient for the parameters whose optimizer state it owns. So replace the all-reduce with a reduce-scatter: each GPU receives only the summed gradients for its shard. Memory drops to $2 + 14/N$ bytes/param. Communication: still unchanged — you were already going to do reduce-scatter + all-gather; you just don't all-gather the gradients you're about to throw away.

Stage 3 — shard the parameters too. Now each GPU permanently holds only $1/N$ of the weights. To run the forward pass through a layer, the GPUs all-gather that layer's parameters just in time, use them, and discard the non-local shards immediately. The backward pass all-gathers them again. Memory becomes $16/N$ bytes/param — linear in GPU count, so 64 GPUs give a 64× reduction and the model effectively has no fixed memory floor. Communication: ~1.5× DDP, because parameters are now gathered twice (forward and backward) on top of the gradient reduce-scatter.

Put together:

For 8B on 8 GPUs the fixed per-GPU state goes 120 → 41 → 28 → 15 GiB across the stages; on 64 GPUs, ZeRO-3 takes it to under 2 GiB. Watch it move on the memory calculator — the ZeRO slider there is exactly this table.

FSDP: the same idea, and the knobs that matter#

PyTorch's Fully Sharded Data Parallel is ZeRO-3 with first-class framework support, and its configuration surface is where the abstract trade becomes concrete decisions:

• Sharding unit. You wrap the model into FSDP units — usually one per transformer block. The unit is the granularity of an all-gather: the whole unit's parameters are gathered, used, and freed together. Too coarse and the transient gathered memory spikes; too fine and you drown in small collectives.
• Prefetching. FSDP can all-gather layer $L+1$'s parameters while layer $L$ is still computing, hiding the communication. This is the single most important setting for throughput and is why a well-tuned FSDP run isn't much slower than DDP despite moving far more data.
• Reshard after forward. Free the gathered parameters immediately after the forward pass and re-gather them for backward (saves memory, costs a second gather), or keep them resident (the reverse trade). This is a memory/comm dial you set per run.

Two more things you'll actually hit#

Gradient accumulation. Independent of sharding, and the cheapest way to grow the effective batch: run several micro-batches, summing gradients, and only all-reduce + step once every $k$ of them. It divides communication frequency by $k$ and lets you hit a large target batch size that wouldn't fit in memory at once. The one correctness note: scale the loss by $1/k$ so the accumulated gradient is a mean, not a sum, and make sure DDP doesn't synchronize on the intermediate micro-batches (no_sync()), or you throw away the savings.

Synchronous vs asynchronous. Everything above is synchronous — every GPU waits at the all-reduce barrier each step, so all replicas train on identical weights. Asynchronous schemes (parameter servers, stale-gradient methods) remove the barrier so no GPU ever waits, but they train on stale weights, which degrades the per-step learning efficiency enough that they've largely lost for dense LLM training. The barrier is a real cost, and paying it is still the right call — a fact worth remembering when someone proposes removing it.