Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

ZeRO Optimization › Worked Example: Training State Partitioning

Michael BrenndoerferJanuary 21, 20266 min read

Language AI HandbookAll chapters

Part 31 of 67Training Infrastructure

Chapter 225ZeRO Optimization › Worked Example: Training State Partitioning

Reading section 9 of 14 · Worked Example: Training State Partitioning · 6 minSection 9 of 14 · 6 min

Section 9 of 14 · 6 min read

ZeRO Optimization › Worked Example: Training State PartitioningOne section · Sections in this chapter9 / 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

Let's trace through a concrete Stage 2 training step with $N_d = 4$ GPUs and a toy model to make the mechanics concrete. This example uses a simplified SGD optimizer instead of Adam to keep the arithmetic clear, but the partitioning logic is identical.

In[12]:

Code

import numpy as np

# Toy model: 8 parameters total, each GPU owns 2 parameters
N_d = 4
total_params = 8
params_per_gpu = total_params // N_d

# Initial parameters (same on all GPUs - full replica before forward pass)
np.random.seed(42)
params_global = np.array(
    [0.1, -0.3, 0.5, 0.2, -0.1, 0.4, -0.2, 0.3], dtype=np.float32
)

# Each GPU processes different data, producing different local gradients
# (In reality, backprop on different mini-batches)
local_grads = {
    0: np.array(
        [0.04, -0.02, 0.06, 0.01, -0.03, 0.05, -0.01, 0.02], dtype=np.float32
    ),
    1: np.array(
        [0.02, -0.04, 0.03, 0.05, -0.01, 0.02, -0.04, 0.03], dtype=np.float32
    ),
    2: np.array(
        [0.06, -0.01, 0.04, 0.02, -0.02, 0.06, -0.02, 0.01], dtype=np.float32
    ),
    3: np.array(
        [0.02, -0.03, 0.05, 0.04, -0.04, 0.01, -0.03, 0.04], dtype=np.float32
    ),
}

# Step 1: Average gradients across GPUs (what Reduce-Scatter computes as average)
averaged_grads = np.mean([local_grads[i] for i in range(N_d)], axis=0)

# Step 2: Reduce-Scatter - each GPU receives the averaged gradient for its shard
owned_grad_shards = {
    i: averaged_grads[i * params_per_gpu : (i + 1) * params_per_gpu]
    for i in range(N_d)
}

# Step 3: Each GPU applies optimizer update to its owned parameter shard
# Simple SGD update (lr=0.1) for illustration
lr = 0.1
owned_param_shards = {
    i: params_global[i * params_per_gpu : (i + 1) * params_per_gpu]
    - lr * owned_grad_shards[i]
    for i in range(N_d)
}

# Step 4: All-Gather - reconstruct full updated parameters on all GPUs
params_updated = np.concatenate([owned_param_shards[i] for i in range(N_d)])

Out[13]:

Console

=== ZeRO Stage 2 Training Step (4 GPUs, 8 Parameters) ===

Initial parameters: [ 0.1 -0.3  0.5  0.2 -0.1  0.4 -0.2  0.3]
Averaged gradients: [ 0.035 -0.025  0.045  0.03  -0.025  0.035 -0.025  0.025]

After Reduce-Scatter - each GPU owns gradient shard:
  GPU 0 owns params [0, 1]: grad_shard = [ 0.035 -0.025]
  GPU 1 owns params [2, 3]: grad_shard = [0.045 0.03 ]
  GPU 2 owns params [4, 5]: grad_shard = [-0.025  0.035]
  GPU 3 owns params [6, 7]: grad_shard = [-0.025  0.025]

After optimizer step on each shard:
  GPU 0: updated shard = [ 0.0965 -0.2975]
  GPU 1: updated shard = [0.4955 0.197 ]
  GPU 2: updated shard = [-0.0975  0.3965]
  GPU 3: updated shard = [-0.1975  0.2975]

After All-Gather - full updated params: [ 0.0965 -0.2975  0.4955  0.197  -0.0975  0.3965 -0.1975  0.2975]

Expected (naive full-param update): [ 0.0965 -0.2975  0.4955  0.197  -0.0975  0.3965 -0.1975  0.2975]
Results match: True

This trace reveals something important: the final parameter values are identical to what a naive full-parameter update would produce. ZeRO is mathematically equivalent to standard data parallelism. The partitioning changes where computation happens and what tensors each GPU stores, but not the values that result. Each GPU performs the optimizer step only on its own shard, but because the averaged gradients are the same regardless of which GPU computes them, the resulting parameter updates are identical to a centralized update.

This mathematical equivalence is what makes ZeRO safe to use without any adjustment to hyperparameters, learning rate schedules, or convergence expectations. Training a model with ZeRO Stage 3 on 64 GPUs produces the same model as training it with Stage 1 on 4 GPUs or with standard data parallelism on a single GPU (given the same effective batch size and random seeds). ZeRO is purely an implementation optimization, not a change to the training algorithm.

Visualizing the Partitioning StructureLink Copied

Let's visualize how the parameter shards, gradient shards, and optimizer state shards map to GPUs under each ZeRO stage, to build an intuitive picture of what each device is responsible for.

Out[14]:

Visualization

Grid diagram showing tensor ownership per GPU under ZeRO Stage 1, 2, and 3 for parameters, gradients, and optimizer states. — Visualization of which tensors each GPU owns under each ZeRO stage for a 4-GPU example with 8 parameters. Each colored cell represents a parameter shard owned by a specific GPU. Under Stage 1, parameters and gradients remain fully replicated across all GPUs (all cells the same color), while optimizer states are partitioned. Stage 2 additionally partitions gradients. Stage 3 partitions all three components, with each GPU owning a unique shard of every tensor type.

The diagram makes the partitioning structure concrete. Under Stage 1, every GPU holds the same full replica of parameters and gradients (shown in gray), but each GPU owns a distinct colored shard of the optimizer states. Stage 2 additionally shards the gradients. Stage 3 shards everything: each color appears in exactly one quarter of each row, and no two GPUs hold any overlapping permanent state. The All-Gather operations can be understood as temporarily merging the colors within a row before computation proceeds.