ZeRO Stages: Sharding Optimizer State, Gradients, and Parameters

Data parallelism replicates everything. Each of $P$ devices holds a complete copy of the parameters, a complete gradient buffer, and a complete optimiser state, then exchanges gradients so all copies stay identical. The replication is what makes the scheme simple and what caps the model size: the per-device memory requirement is the full model’s requirement, no matter how many devices are added.

Sharding the replicated state removes that cap. The state is divided across the same $P$ ranks that already hold replicas, so the per-device memory for the sharded terms falls by a factor of $P$ while the batch is still split for computation. The cost is additional communication at the boundaries where a rank needs state it does not own.

What stage 1 partitions

The first stage shards the optimiser state: the first and second moments, and the full-precision master weights if mixed precision is in use. With $P$ ranks, rank $i$ owns the slice of each buffer corresponding to parameters $[i \cdot N/P, (i+1) \cdot N/P)$.

The forward and backward passes still use complete parameter copies on every rank, so the computation is unchanged and the gradients arrive complete on every rank after the ordinary all-reduce. The difference appears at the update: each rank applies the optimiser only to its own slice, using the slice of the gradient it received, and updates only its own master weights. The reduced memory is the optimiser term, which for Adam in mixed precision is 12 of the 18 bytes per parameter accounted for in the memory budget. Stage 1 therefore removes two thirds of the state while leaving parameters and gradients replicated.

At the end of the update, each rank holds fresh values for its slice and stale values for everything else. An all-gather of the updated parameters restores the replicas. The communication is one all-gather of $N$ parameters per step, added to the all-reduce that was already there.

Stage 2 also partitions gradients

The second stage shards the gradients as well. The reduction is changed from all-reduce to reduce-scatter: instead of every rank receiving the full sum, each rank receives the summed slice it is responsible for. The total traffic is the same as all-reduce, since all-reduce is itself reduce-scatter followed by all-gather, but the all-gather is now only needed for parameters.

The saving is the gradient term, 4 bytes per parameter, at no additional communication beyond what stage 1 required. The mechanism is economical because the reduction and the sharding are the same collective: the reduce-scatter both sums across ranks and leaves each rank with a distinct slice.

There is a subtlety in the buffer layout. A reduce-scatter writes its output into a buffer sized for one slice per rank, so the persistent gradient storage is $N/P$ parameters rather than $N$. Frameworks that pre-allocate gradient buffers for the full parameter set at construction time keep the full allocation alive even under stage 2, and the saving never materialises. The allocation has to follow the partitioning; if the full buffer is retained, the stage is enabled in name only.

A second subtlety concerns accumulation. When gradients are summed over several micro-batches before the reduction, the partial sums are per-rank and complete, so the reduce-scatter happens once per optimiser step rather than once per micro-batch. The traffic is unchanged; what changes is that the partial buffers must hold the full parameter set for the duration of the window, which raises the peak by exactly the amount stage 2 was meant to remove. Accumulation and gradient sharding therefore conflict at the buffer level, and the usual resolution is to reduce per micro-batch and accumulate the reduced slices instead.

Stage 3 partitions parameters

The third stage shards the parameters themselves. No rank holds a complete layer at rest. Before a layer’s forward computation, its parameters are all-gathered from the owning ranks; after the computation, they are released. The same happens for the backward pass, and again for the update.

This is a different kind of change. Stages 1 and 2 move bookkeeping off the critical path; stage 3 makes parameter residency dynamic, which means the memory profile is a function of the schedule rather than a fixed allocation. Peak memory is the largest layer’s gathered parameters plus the working set of everything else, not the sum over all layers.

The communication cost rises accordingly. Each parameter is gathered once per forward and once per backward, so the traffic is roughly $2N$ per step in addition to the gradient reduction. That is a large volume, and it is why stage 3 is only worthwhile when the model does not otherwise fit, and why it is normally combined with mechanisms that reduce the number of gathers.

# Per-layer gather and release, the shape of a stage-3 forward pass.
for layer in model.layers:
    full = all_gather(layer.shard)      # N_layer bytes from all ranks
    out = layer.forward(full, inp)
    layer.shard = None                  # release before the next layer
    full = None
    inp = out

The prefetch matters more than the gather itself. A gather issued immediately before the layer runs leaves the device idle for its duration; a gather issued one layer ahead overlaps with the current layer’s arithmetic. Choosing the prefetch depth is a memory-versus-latency trade, because prefetched layers occupy memory that the released ones would have freed.

Communication volume compared

For $P$ ranks and $N$ parameters, the per-step communication per rank across the three stages is approximately:

Stage	Collective	Volume per rank
Baseline	all-reduce of gradients	$2N(P-1)/P$
Stage 1	all-reduce + all-gather of parameters	$2N(P-1)/P + N(P-1)/P$
Stage 2	reduce-scatter + all-gather of parameters	$N(P-1)/P + N(P-1)/P$
Stage 3	reduce-scatter + 2 gathers per step	$\approx 3N(P-1)/P$

Stage 2 is notable for costing no more than the baseline all-reduce despite halving the gradient memory: the reduce-scatter and the parameter all-gather together move $2N(P-1)/P$, exactly what the original all-reduce moved, while the stage-1 arrangement needed an additional gather on top of it. Stage 3 roughly triples the volume, which is the price of not holding the model.

These are asymptotic counts. At small $P$ the per-message latency dominates and the stage with more collectives loses more than the table suggests, which is one reason sharding is not used on two devices when replication fits.

Where the shards live and what happens on failure

Sharded state is distributed, so a rank failure loses the slice it owned. Checkpointing must therefore write either a gathered copy, which requires every rank to participate, or a per-rank shard set with enough metadata to reconstruct the partitioning later. The second is cheaper and is what most implementations do, but it makes the checkpoint dependent on the world size and the sharding layout: loading it at a different $P$ requires reshuffling, which is a real operation and not a file copy.

Stage 3 makes this sharper, because at rest a rank holds only fragments. A checkpoint taken mid-step, when some layers are gathered and others are not, captures an inconsistent view. The safe point is between steps, with all shards released and the optimiser state consistent.

When each stage is the right choice

Stage 1 applies when the optimiser state is the largest term and the model otherwise fits. This is the common case for models in the low billions of parameters on devices with tens of gigabytes: the state is three times the weight footprint, so sharding it recovers the majority of the problem at a communication cost of one extra gather.

Stage 2 applies when the gradient buffer is also significant, which for models near the device capacity is usually true. The absence of extra cost relative to stage 1 makes it the default choice whenever stage 1 is warranted, provided the implementation supports the fused reduce-scatter.

Stage 3 applies when parameters, gradients, and state together exceed capacity, or when a very long sequence makes activations dominant and the parameter term must be minimised to leave room. Its gather traffic means it is normally paired with a topology-aware process group: gathering across a fast intra-node interconnect is far cheaper than across nodes, so the shard group is usually confined to a node and a separate mechanism handles the cross-node dimension.

The interaction with pipeline parallelism is the subtle part. A pipeline stage already holds only a subset of layers, so sharding within a stage divides a smaller model; the two mechanisms multiply, and the ordering of the process groups determines whether the expensive collectives stay intra-node. That composition is the subject of planning a parallelism layout.