ZeRO Optimization › ZeRO Stage 2: Gradient Partitioning

ZeRO Stage 2 extends Stage 1 by also partitioning the gradient tensors. In Stage 1, gradients are still replicated in full before the Reduce-Scatter operation: each GPU holds the complete averaged gradient tensor until after the optimizer step. Stage 2 observes that this full replication is also unnecessary. After the Reduce-Scatter averaging, each GPU only needs the gradients corresponding to its owned parameter shard to perform the optimizer update. There is no reason to keep the other shards' gradients around.

Stage 2 uses a Reduce-Scatter as the fundamental communication primitive after the backward pass. The distinction from Stage 1 lies in when the gradient memory is freed:

1. Each GPU starts with its local gradient tensors computed from its data shard.
2. The Reduce-Scatter operation averages the gradients across all GPUs and simultaneously distributes them so that GPU $i$ ends up with the averaged gradient for parameter shard $i$ only.
3. Each GPU runs the optimizer step on its shard using the received gradient shard and its stored optimizer state shard.
4. An All-Gather restores the full updated parameters everywhere.

The key difference from Stage 1 is that between the Reduce-Scatter and the All-Gather, the full gradient tensor no longer exists on any single GPU. Each GPU holds only $1/N_d$ of the total gradient. This saves $2\Psi$ bytes per GPU compared to Stage 1. The communication volume is unchanged: the same Reduce-Scatter and All-Gather operations occur, just with the gradient memory freed earlier.

The reason gradient sharding works correctly is the same reason Stage 1 works: after the optimizer step, each GPU reconstructs the full parameters via All-Gather. The gradient for shard $i$ is only ever needed by GPU $i$ (to update its owned parameters), so there is no reason any other GPU needs to retain it. The gradient is a means to an end, and once the optimizer step has consumed it, it can be discarded.

This also has a practical benefit for implementations: the gradient tensor for shard $j$ can be freed as soon as the Reduce-Scatter for that shard completes, even before all shards have finished their Reduce-Scatter. In practice, implementations bucket the Reduce-Scatter operations so that gradient memory is released progressively as each bucket's communication completes, reducing peak memory below the formula's worst case.

$$M_{\text{Stage 2}} = 2\Psi + \frac{2\Psi + 12\Psi}{N_d}$$

where:

• $2\Psi$: 16-bit parameter replica (still a full copy, needed for forward and backward passes)
• $(2\Psi + 12\Psi) / N_d$: gradient shard plus optimizer state shard for the owned portion only

With $N_d = 64$ GPUs and a 7B model, Stage 2 per-GPU memory is approximately 14 GB (parameters) plus 0.9 GB (sharded gradients and optimizer states), totaling roughly 14.9 GB. This fits on a single A100 with room to spare, compared to 112 GB for the baseline. The memory savings relative to Stage 1 come entirely from eliminating the $2\Psi$ full gradient replica.

An important practical consideration: Stage 2 is fully compatible with gradient accumulation. When accumulating gradients across multiple micro-batches before the optimizer step, the partial gradients are accumulated locally on each GPU without any cross-GPU communication. The Reduce-Scatter only fires when the full accumulation is complete and the optimizer step is about to begin. This means the memory profile and communication pattern are unchanged when using gradient accumulation, which is commonly used to simulate large effective batch sizes on hardware with limited memory.

Stage 2 is also compatible with mixed-precision training in the standard way. The forward and backward passes use fp16 parameters and gradients, and the optimizer step converts the sharded gradient to fp32 before updating the fp32 optimizer states and master parameters. The fp16 parameter replica is then updated from the fp32 master, and the All-Gather broadcasts the updated fp16 parameters to all devices.