Sharded Data Parallelism: ZeRO Stages 1-3

"I kept a full copy of the optimizer state, just like all 511 of my colleagues. Between us we stored the same numbers 512 times, and then we wondered why we ran out of memory."

A Data-Parallel Worker, Hoarding Redundantly

The previous section split a single model across devices by cutting it into tensor and pipeline stages, which changes the structure of the computation and demands that the model be carved up by hand. That works, but it gives up the great virtue of data parallelism from Chapter 15: every worker runs the identical program on its own shard of the data, and the framework needs to know nothing about the model's internal structure. Sharded data parallelism asks whether we can keep that virtue and still fit a model far larger than one device, by attacking memory rather than computation. The answer is yes, and the idea is almost embarrassingly simple once stated: stop storing the same numbers on every worker.

Recall the memory bill from training. For a model with $\Phi$ parameters trained with mixed-precision Adam, each worker in plain data parallelism holds an fp16 copy of the parameters ($2\Phi$ bytes), an fp16 copy of the gradients ($2\Phi$ bytes), and the optimizer state, which for Adam is an fp32 master copy of the weights plus the first and second moments ($4\Phi + 4\Phi + 4\Phi = 12\Phi$ bytes). That is $16\Phi$ bytes per worker, and the $12\Phi$ of optimizer state dominates. With $K$ workers, the cluster stores the same $16\Phi$ bytes $K$ times over. Every worker needs its own slice at any given instant, but no worker needs the whole thing all the time. ZeRO exploits exactly that gap.

1. The Redundancy in Plain Data Parallelism Beginner

In data-parallel training each of the $K$ workers processes a different micro-batch, computes a local gradient, and the workers combine those gradients with an all-reduce so that all of them apply the identical update and stay in lockstep. Because they stay in lockstep, every worker's parameters are identical, every worker's averaged gradient is identical, and after the optimizer step every worker's optimizer state is identical too. The workers are running $K$ exact replicas of the same state. That replication is what makes the method trivial to reason about, and it is pure waste: of the $16\Phi$ bytes a worker stores, only $16\Phi / K$ of it is information unique to that worker. The rest is a copy of what its neighbors already hold.

The Zero Redundancy Optimizer, introduced by Rajbhandari and colleagues at Microsoft for the DeepSpeed library, observes that we can partition this replicated state across the workers so that worker $k$ is the sole owner of slice $k$, and reconstruct any full tensor only at the moment it is required for computation. The partition is purely a memory-placement decision; the arithmetic each worker performs on its own data is unchanged, which is why ZeRO preserves the data-parallel programming model exactly. The cost is that reconstructing a tensor that lives in pieces across the cluster requires moving those pieces, and freeing it afterward means moving them again next time. The whole design is a trade of memory for communication, and ZeRO offers it in three stages of increasing aggressiveness.

2. Three Stages: Sharding Optimizer State, Gradients, Parameters Intermediate

The three model-state components differ in how much memory they occupy and how often they are touched, so ZeRO shards them in order of payoff. Let $K$ be the number of data-parallel workers and write the per-device model-state memory at each stage. Plain data parallelism stores

Stage 1 (ZeRO-1) shards the optimizer state, the largest term, leaving the parameters and gradients replicated:

Stage 2 (ZeRO-2) additionally shards the gradients, since each worker only needs the gradient slice corresponding to the optimizer slice it owns:

Stage 3 (ZeRO-3) finally shards the parameters themselves, so that no full tensor is resident anywhere except transiently during its own computation:

The pattern is clean: each stage moves one more term from "replicated" to "divided by $K$." Stage 3 divides everything, so per-device memory falls like $1/K$ without bound, which is what makes trillion-parameter training possible on accelerators that individually hold only tens of gigabytes. Figure 16.4.1 shows the same progression as a picture: shaded blocks owned by one worker, faint blocks owned by neighbors.

3. The Communication: All-Gather to Borrow, Reduce-Scatter to Return Intermediate

Sharding the parameters means no worker holds a full layer's weights at rest, yet computing that layer's output needs all of its weights. ZeRO-3 resolves this with a just-in-time dance. Right before a layer runs in the forward pass, the workers all-gather its parameter shards so that every worker briefly holds the complete layer; the layer computes; then each worker frees the parts it does not own, dropping back to its $1/K$ slice. The backward pass all-gathers the same parameters again to compute gradients, then the freshly computed gradients are combined and partitioned in one step with a reduce-scatter, so that worker $k$ ends up holding exactly the summed gradient slice for the parameters it owns and can apply its slice of the optimizer step locally.

This is the moment the thesis thread of the whole book reappears. In Chapter 15 plain data parallelism combined gradients with a single all-reduce. A ring all-reduce is internally exactly a reduce-scatter followed by an all-gather, as built in Section 4.5. ZeRO does not invent a new collective; it splits the all-reduce back into its two halves and uses each half where the sharded layout needs it: reduce-scatter to land each gradient slice on its owner, all-gather to reconstruct parameters on demand. The same primitive from Section 1.3 of Chapter 1, scaled out one more way.

The communication bill follows directly. Stages 1 and 2 still move a reduce-scatter plus an all-gather worth of gradient data per step, which is the same $2\Phi$ of traffic as a plain all-reduce, so they shrink memory for free in communication terms. Stage 3 adds the parameter all-gathers: one in the forward pass and one in the backward pass, each moving $\Phi$, on top of the gradient reduce-scatter of $\Phi$. That is $3\Phi$ versus the $2\Phi$ of plain data parallelism, a $1.5\times$ increase. The trade is summarized in Table 16.4.1 and made concrete by the demo that follows.

4. Watching Memory Fall Like 1/K Intermediate

The accounting above is easy to state and easy to get wrong, so the code below computes it directly. It models a 1.5-billion-parameter network (GPT-2 XL scale) trained with mixed-precision Adam and reports the per-device model-state memory for plain data parallelism and ZeRO Stages 1 through 3 as the worker count $K$ grows, then the communication volume each method moves per step. There are no deep-learning libraries here, only the byte arithmetic of Section 2 and Section 3, so the numbers are exactly the formulas made visible.

import numpy as np

P = 1_500_000_000            # 1.5B parameters (GPT-2 XL scale)
b_fp16_param = 2             # fp16 parameter copy, bytes/param
b_fp16_grad  = 2             # fp16 gradient
b_opt        = 12            # Adam fp32 master + m + v = 4+4+4 bytes/param
GB = 1024**3

def per_device_gb(K, stage):
    param = b_fp16_param * P
    grad  = b_fp16_grad  * P
    opt   = b_opt * P
    if stage >= 1: opt   /= K     # ZeRO-1: shard optimizer state
    if stage >= 2: grad  /= K     # ZeRO-2: also shard gradients
    if stage >= 3: param /= K     # ZeRO-3: also shard parameters
    return (param + grad + opt) / GB

print("Per-device model-state memory (GB), P = 1.5e9 params, Adam mixed precision")
print(f"{'K':>4} | {'plain DP':>9} | {'ZeRO-1':>8} | {'ZeRO-2':>8} | {'ZeRO-3':>8}")
print("-" * 52)
for K in (1, 2, 8, 32, 128, 512):
    row = [per_device_gb(K, s) for s in (0, 1, 2, 3)]
    print(f"{K:>4} | {row[0]:>9.2f} | {row[1]:>8.2f} | {row[2]:>8.2f} | {row[3]:>8.2f}")

# Communication per step in units of Phi (one fp16 copy of all params).
# Plain DP / ZeRO-1 / ZeRO-2: gradient all-reduce = reduce-scatter + all-gather = 2*Phi.
# ZeRO-3: forward all-gather (Phi) + backward all-gather (Phi) + grad reduce-scatter (Phi).
Phi = b_fp16_param * P / GB
print("\nCommunication volume per step (GB moved per device, fp16 params)")
print(f"{'method':>10} | {'all-gather':>11} | {'reduce-scatter':>15} | {'total':>7}")
print("-" * 52)
for name, ag, rs in [("plain DP", Phi, Phi), ("ZeRO-1", Phi, Phi),
                     ("ZeRO-2", Phi, Phi), ("ZeRO-3", 2*Phi, Phi)]:
    print(f"{name:>10} | {ag:>11.2f} | {rs:>15.2f} | {ag+rs:>7.2f}")

Per-device model-state memory (GB), P = 1.5e9 params, Adam mixed precision
   K |  plain DP |   ZeRO-1 |   ZeRO-2 |   ZeRO-3
----------------------------------------------------
   1 |     22.35 |    22.35 |    22.35 |    22.35
   2 |     22.35 |    13.97 |    12.57 |    11.18
   8 |     22.35 |     7.68 |     5.24 |     2.79
  32 |     22.35 |     6.11 |     3.41 |     0.70
 128 |     22.35 |     5.72 |     2.95 |     0.17
 512 |     22.35 |     5.62 |     2.83 |     0.04

Communication volume per step (GB moved per device, fp16 params)
    method |  all-gather |  reduce-scatter |   total
----------------------------------------------------
  plain DP |        2.79 |            2.79 |    5.59
    ZeRO-1 |        2.79 |            2.79 |    5.59
    ZeRO-2 |        2.79 |            2.79 |    5.59
    ZeRO-3 |        5.59 |            2.79 |    8.38

Two facts jump out of Output 16.4.1. First, plain data parallelism never benefits from more workers in memory terms: the 22.35 GB line is constant, which is precisely why a 1.5-billion-parameter model with Adam will not train on a 16 GB accelerator under plain data parallelism no matter how large the cluster. Second, the stages separate by how much they help and how much they cost. ZeRO-1 and ZeRO-2 reduce memory substantially at no extra communication, so they are nearly free wins; ZeRO-3 reduces memory without bound but adds the $1.5\times$ communication seen in the lower table. The choice of stage is therefore a memory-versus-bandwidth decision: use the lowest stage that makes the model fit, because each higher stage trades away some communication headroom you might prefer to spend on bigger batches.

5. Choosing a Stage, and Where It Breaks Down Advanced

The stages form a ladder you climb only as far as you must. If the model fits under ZeRO-1, stop there and keep the bandwidth. If it does not, ZeRO-2 costs nothing more in communication and frees the gradient memory. Only when the parameters themselves will not fit do you reach for ZeRO-3 and accept the parameter all-gathers. The practical complication is that ZeRO-3's all-gathers happen layer by layer during the forward and backward passes, so they sit on the critical path in a way the single end-of-step all-reduce of plain data parallelism does not. Implementations hide this by prefetching the next layer's parameters while the current layer computes, overlapping communication with computation, but on a slow interconnect the overlap is imperfect and ZeRO-3 can leave accelerators idling on the network. This is the same communication-versus-computation tension quantified by the cost models of Chapter 3, now applied per layer rather than per step.

There is also a scaling subtlety. As $K$ grows the sharded terms vanish but the replicated terms do not. Under ZeRO-1 the per-device memory tends to $4\Phi$ as $K \to \infty$, the floor set by the still-replicated parameters and gradients; you can see it in Output 16.4.1 flattening near 5.6 GB. Only ZeRO-3, with nothing replicated, keeps falling. This is why frameworks pair ZeRO-3 with the tensor and pipeline parallelism of the previous section for the very largest models: ZeRO-3 handles the data-parallel dimension, and the other axes cut the per-layer activation and compute that sharding alone does not touch. Combining all three is the 3D-parallelism whose trade-offs Chapter 17 and the foundation-model training of this part assemble in full.

# ds_config.json: {"zero_optimization": {"stage": 3,
#                  "overlap_comm": true, "contiguous_gradients": true}}
import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(),
    config="ds_config.json")          # stage 1, 2, or 3 chosen here

for batch in dataloader:
    loss = model_engine(batch)        # forward; params all-gathered per layer
    model_engine.backward(loss)       # backward + gradient reduce-scatter
    model_engine.step()               # each worker updates only its shard

Sharded data parallelism gives us model-scale memory savings while keeping the one-program-per-worker simplicity that makes data parallelism a joy to operate. We have seen what each stage shards, that the collectives are the familiar reduce-scatter and all-gather, and that the memory falls like $1/K$ while ZeRO-3 alone pays a $1.5\times$ communication tax. The next section, Section 16.5, turns this design into running PyTorch code with FSDP, which builds its public API directly on the all-gather and reduce-scatter pair the thesis thread keeps returning to.