ZeRO Optimization › ZeRO Stage 3: Parameter Partitioning

ZeRO Stage 3 completes the partitioning by distributing the parameters themselves across the GPUs. In Stages 1 and 2, each GPU still holds a full 16-bit replica of all parameters for the forward and backward passes. This $2\Psi$ bytes is the last significant replicated component. Stage 3 eliminates it, reducing every component of model state to $1/N_d$ per GPU.

The challenge with partitioning parameters is that the forward and backward passes need parameters for every layer. If GPU $i$ only permanently stores a $1/N_d$ shard, how does it execute the forward pass through layers whose parameters live on other GPUs?

Stage 3 exploits the sequential structure of the forward pass: a transformer processes layers one at a time, and at any given moment, only the parameters for the current layer are actively needed. Stage 3 uses this property to fetch parameters on demand:

1. Each GPU permanently owns only $1/N_d$ of the total parameters, distributed across all layers.
2. Before layer $l$'s computation begins, all GPUs participate in an All-Gather to collect the complete parameter set for layer $l$ from whichever GPUs own those shards.
3. The computation for layer $l$ proceeds normally, using the fully gathered parameters.
4. Once layer $l$'s contribution to the activations (forward) or gradients (backward) has been computed, the gathered parameters are discarded. The GPU reverts to holding only its permanent shard.

The backward pass mirrors this: backpropagation traverses layers in reverse, fetching each layer's parameters via All-Gather before computing gradients, then immediately discarding the fetched parameters. A Reduce-Scatter ships the local gradient contribution to the GPU that owns each parameter shard, which then performs the optimizer update.

This "just-in-time" parameter fetching means that at any moment, a GPU holds:

1. Its permanently-owned parameter shard, which is $1/N_d$ of all parameters
2. The parameters for the layer currently being computed, fetched transiently via All-Gather
3. Its permanently-owned gradient shard and optimizer state shard

The transient parameter buffer for the current layer is bounded by the size of the largest single layer, which is typically much smaller than $\Psi$. For a transformer model, the largest layer (usually the MLP FFN) has roughly $4 \times d_\text{model}^2$ parameters. For a 7B model with $d_\text{model} = 4096$, that is about $4 \times 4096^2 = 67 \text{ million}$ parameters, or 0.13 GB at fp16, which is negligible compared to the savings.

When all components are partitioned, the per-GPU memory for model state becomes:

$$M_{\text{Stage 3}} = \frac{16\Psi}{N_d}$$

where $16\Psi$ is the total model state (fp16 params + fp16 grads + fp32 master + fp32 Adam moments). Every component scales with $N_d$. This is a perfect linear reduction in per-GPU memory with the number of GPUs.

For a 7B model with 64 GPUs:

$$\frac{16 \times 7 \times 10^9 \times 1 \text{ byte}}{64} = \frac{112 \text{ GB}}{64} = 1.75 \text{ GB}$$

Stage 3 makes it theoretically possible to train arbitrarily large models given enough GPUs, since the memory per GPU decreases linearly with $N_d$. Doubling the GPU count halves the per-GPU model state memory. This is a fundamentally different scaling property than Stages 1 and 2, where the $2\Psi$ parameter replica creates a floor that does not decrease with more GPUs.

Stage 3's memory efficiency comes at a communication price. The total communication volume relative to standard data parallelism is:

• Baseline (Standard DP): 1 All-Reduce equivalent, total volume $2\Psi$ bytes per step
• Stage 1: 1 All-Reduce equivalent, total volume $2\Psi$ bytes (Reduce-Scatter + All-Gather for gradients/parameters)
• Stage 2: 1 All-Reduce equivalent, total volume $2\Psi$ bytes (same operations, gradient memory freed earlier)
• Stage 3: 3 All-Reduce equivalents, total volume $6\Psi$ bytes per step

Stage 3's forward pass requires one All-Gather to fetch parameters layer by layer (total volume $2\Psi$). The backward pass requires a second All-Gather for parameters (another $2\Psi$) plus a Reduce-Scatter for gradients ($2\Psi$). The full step therefore transfers $6\Psi$ bytes versus $2\Psi$ for the baseline. This 3x communication overhead is manageable on high-bandwidth interconnects such as NVLink (600 GB/s bidirectional on A100 nodes) or InfiniBand (200-400 Gb/s between nodes), but can become a significant bottleneck on slower networks.

The communication-compute overlap is critical to Stage 3's practical performance. Modern implementations prefetch the parameters for the next layer while computing the current layer, so the All-Gather for layer $l+1$ proceeds in parallel with the computation for layer $l$. If the network is fast enough to transfer a layer's parameters before the GPU finishes computing with the previous layer, the communication overhead is effectively hidden. This prefetching only works for sequential architectures where the execution order is predictable at scheduling time.