Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

ZeRO Optimization › ZeRO Stage 1: Optimizer State Partitioning

Michael BrenndoerferJanuary 21, 20265 min read

Language AI HandbookAll chapters

Part 31 of 67Training Infrastructure

Chapter 225ZeRO Optimization › ZeRO Stage 1: Optimizer State Partitioning

Reading section 3 of 14 · ZeRO Stage 1: Optimizer State Partitioning · 5 minSection 3 of 14 · 5 min

Section 3 of 14 · 5 min read

ZeRO Optimization › ZeRO Stage 1: Optimizer State PartitioningOne section · Sections in this chapter3 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

ZeRO Stage 1 targets the largest single component of model state: the Adam optimizer states. In standard data parallelism, all $N_d$ GPUs hold identical copies of these tensors even though each GPU applies the optimizer update to the full set of parameters. Every GPU ends up doing the same work and storing the same state as all others. Stage 1 eliminates this by assigning each GPU responsibility for updating a $1/N_d$ shard of the parameters, so it only needs to store the optimizer states for that shard.

How the Partitioning WorksLink Copied

Before the training step begins, parameters are replicated across all GPUs as usual. The forward pass and backward pass proceed identically to standard data parallelism: each GPU processes its local data shard, computes a loss, and runs backpropagation to produce gradient tensors. Up to this point, Stage 1 is indistinguishable from baseline data parallelism.

After the backward pass, Stage 1 diverges. Instead of each GPU independently applying the optimizer update to the full parameter set (which would require holding the full optimizer state), Stage 1 performs a Reduce-Scatter operation:

Each GPU computes gradients for its local data shard as normal.
A Reduce-Scatter operation is performed: the gradients are averaged across all GPUs, and each GPU ends up with the averaged gradient for only its assigned parameter shard, rather than the full averaged gradient tensor.
Each GPU applies the optimizer update to its owned shard using only the optimizer states for that shard. Because each GPU owns a different shard, the $N_d$ GPUs collectively perform one complete optimizer step across all parameters without any single GPU needing the full optimizer state.
An All-Gather operation distributes the updated parameters back to all GPUs, restoring the full parameter replica everywhere.

The forward pass of the next iteration uses the complete, freshly-updated parameters that every GPU now holds. The key observation is that the communication operations in steps 2 and 4 together have the same total volume as the All-Reduce in standard data parallelism. A Reduce-Scatter followed by an All-Gather is mathematically equivalent to an All-Reduce, just decomposed into two phases.

Why This Works: Communication EquivalenceLink Copied

Standard data parallelism uses an All-Reduce to average gradients across all $N_d$ GPUs. An All-Reduce sends and receives $2\Psi$ bytes per GPU (in the ring-reduce formulation, it sends $(N_d - 1)/N_d \times \Psi$ bytes and receives the same, which approximates to $2\Psi$ for large $N_d$ ). ZeRO Stage 1's Reduce-Scatter + All-Gather sends and receives the same total volume. The memory savings come for free in terms of communication cost.

This communication equivalence is what makes Stage 1 a nearly free win. You replace the optimizer's most expensive memory component with a sharded version that costs nothing extra in network traffic. The only overhead is that the Reduce-Scatter and All-Gather must be implemented correctly and potentially use different collective implementations than a standard All-Reduce, but modern communication libraries handle this transparently.

Memory SavingsLink Copied

The optimizer states (first moment, second moment, and 32-bit parameter copy) account for $12\Psi$ bytes in mixed-precision Adam. Stage 1 reduces this to $12\Psi / N_d$ bytes distributed across $N_d$ GPUs, so each GPU holds:

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

where:

$2\Psi$ : 16-bit parameter replica (still a full copy, needed for forward and backward passes)
$2\Psi$ : 16-bit gradient replica (still a full copy before Reduce-Scatter)
$12\Psi / N_d$ : optimizer states for the owned shard only

For $N_d = 64$ GPUs, the optimizer state contribution drops from 84 GB per GPU to about 1.31 GB per GPU for a 7B model. The total memory per GPU drops from 112 GB to roughly 29.31 GB. Stage 1 alone can turn an infeasible training job into a feasible one, at essentially no increase in communication cost compared to standard data parallelism. For many practical training runs on relatively modern hardware, Stage 1 is sufficient to fit the desired model.

In[6]:

Code

def zero_stage1_memory(params_billions, num_gpus):
    """Memory per GPU under ZeRO Stage 1 (optimizer state partitioning)."""
    params = params_billions * 1e9
    # fp16 params + fp16 grads still fully replicated
    replicated = (params * 2 + params * 2) / 1e9
    # optimizer states partitioned across num_gpus
    optimizer_per_gpu = (params * 4 + params * 4 + params * 4) / num_gpus / 1e9
    return replicated + optimizer_per_gpu


def zero_stage2_memory(params_billions, num_gpus):
    """Memory per GPU under ZeRO Stage 2 (optimizer state + gradient partitioning)."""
    params = params_billions * 1e9
    # only fp16 params still fully replicated
    replicated = (params * 2) / 1e9
    # gradients + optimizer states partitioned
    partitioned_per_gpu = (
        (params * 2 + params * 4 + params * 4 + params * 4) / num_gpus / 1e9
    )
    return replicated + partitioned_per_gpu


def zero_stage3_memory(params_billions, num_gpus):
    """Memory per GPU under ZeRO Stage 3 (full partitioning of all model states)."""
    params = params_billions * 1e9
    # everything partitioned
    partitioned_per_gpu = (
        (params * 2 + params * 2 + params * 4 + params * 4 + params * 4)
        / num_gpus
        / 1e9
    )
    return partitioned_per_gpu


# Baseline: no ZeRO (standard data parallelism)
def baseline_memory(params_billions):
    return memory_profiles[params_billions]["total_gb"]


gpu_counts = [1, 4, 8, 16, 32, 64, 128]
model_7b_memories = {
    "Baseline": [baseline_memory(7)] * len(gpu_counts),
    "Stage 1": [zero_stage1_memory(7, nd) for nd in gpu_counts],
    "Stage 2": [zero_stage2_memory(7, nd) for nd in gpu_counts],
    "Stage 3": [zero_stage3_memory(7, nd) for nd in gpu_counts],
}

Out[7]:

Console

Per-GPU memory (GB) for a 7B parameter model
  GPUs    Baseline    Stage 1    Stage 2    Stage 3
----------------------------------------------------
     1      112.0G     112.0G     112.0G     112.0G
     4      112.0G      49.0G      38.5G      28.0G
     8      112.0G      38.5G      26.2G      14.0G
    16      112.0G      33.2G      20.1G       7.0G
    32      112.0G      30.6G      17.1G       3.5G
    64      112.0G      29.3G      15.5G       1.8G
   128      112.0G      28.7G      14.8G       0.9G

The progression shows how each stage supports larger models at the same hardware budget. Stage 3 with 64 GPUs reduces per-GPU memory from 112 GB to 1.75 GB for a 7B model, leaving ample headroom for activations and residuals.