ZeRO and FSDP

1. Intuition

Start from the 16 bytes per parameter of mixed-precision AdamW: 2 for bf16 weights, 2 for bf16 gradients, 12 for optimizer state (an fp32 master copy plus Adam's $m$ and $v$). Under plain data parallelism (DDP) every GPU holds all 16.

Stage 1: shard the optimizer state. The optimizer state is the biggest piece (12 of 16 bytes) and is only used in the optimizer step, not in forward or backward. So split it: GPU $r$ keeps the Adam state for slice $r$ of the parameters only. After backward, instead of all-reducing the gradients, reduce-scatter them, so GPU $r$ receives only the summed gradient for slice $r$. It updates slice $r$ and then all-gathers the updated weights so every GPU again has the full model for the next forward.

Stage 2: shard the gradients too. GPU $r$ only ever needs gradient slice $r$, so it can free the rest as soon as the reduce-scatter is done. Same communication as stage 1.

Stage 3: shard the weights. Each GPU stores only its slice of the weights. Before computing a layer, the GPUs all-gather that layer's weights, use them, and free them. Backward gathers them again. Memory per GPU now falls as $1/N$ for everything, but weights cross the network twice per step instead of once. PyTorch calls this FSDP (fully sharded data parallel).

What sharding does not touch. Activations are not sharded by ZeRO: each GPU still processes its own micro-batch. Combine ZeRO with gradient checkpointing when activations are the problem.

2. Visualization

Try it (predict first, then check):

1. On the default 7B-style model with 8 GPUs, predict which strategy first fits under 80 GB. Then switch the GPU memory to 24 GB and predict again.
2. Choose ZeRO-1 and slide the GPU count from 8 to 256. Predict where the bar stops shrinking and which colours make up the floor.
3. Pick the 70B-style model and ZeRO-3. Predict the fewest GPUs that fit 80 GB, then read it in the readout.
4. Switch precision from bf16 + AdamW to fp32 + AdamW. Predict whether the optimizer segment grows or shrinks, and why the total still grows.

3. The math

Symbols

Memory per GPU

Activations (not sharded) come on top. As $N \to \infty$, ZeRO-1 tends to $4P$ and ZeRO-2 to $2P$: replicated pieces set a floor. Only ZeRO-3 keeps falling.

Worked example: 7B-style model on 8 GPUs

$P = 6{,}738{,}149{,}376$, so $2P = 13.48$ GB (decimal GB, rounded).

1. DDP: $16P = 8 \times 13.48 = 107.81$ GB. Plus 0.822 GB activations: 108.63 GB. Matches the fine-tuning lesson.
2. ZeRO-1: weights $2P = 13.48$, gradients $13.48$, optimizer $12P/8 = 80.86/8 = 10.11$. Sum $13.48 + 13.48 + 10.11 = 37.06$, plus 0.822: 37.88 GB.
3. ZeRO-2: weights $13.48$, gradients $13.48/8 = 1.68$, optimizer $10.11$. Sum $25.27$, plus 0.822: 26.09 GB.
4. ZeRO-3: $107.81 / 8 = 13.48$, plus 0.822: 14.30 GB.

So ZeRO-1 alone takes the job from "does not fit an 80 GB GPU" to "fits with room for a bigger micro-batch", a 65% drop.

Communication per GPU per step

A ring all-reduce of $T$ bytes sends $2\frac{N-1}{N}T$ per GPU; reduce-scatter and all-gather each send $\frac{N-1}{N}T$ (collective operations). With bf16 weights and gradients, $T = 2P$ for either.

ZeRO-1 and ZeRO-2 do exactly what a ring all-reduce does internally, a reduce-scatter and an all-gather; the only change is that the all-gather carries updated weights instead of summed gradients. For the 7B-style model on 8 GPUs: DDP sends $1.75 \times 13.48 = 23.58$ GB per GPU per step, and ZeRO-3 sends $1.5 \times 23.58 = 35.38$ GB.

Worked example: when does a 70B-style model fit?

$P = 68{,}975{,}329{,}280$ (from the previous lesson), so $16P = 1{,}103.61$ GB; activations at 2,048 tokens with checkpointing are 3.255 GB.

• ZeRO-3 on 8 GPUs: $1{,}103.61 / 8 = 137.95$, plus 3.255 = 141.21 GB. Does not fit.
• On 16 GPUs: $1{,}103.61 / 16 = 68.98$, plus 3.255 = 72.23 GB. Fits 80 GB.
• ZeRO-1 can never fit: its floor is $4P = 275.9$ GB.

4. Implementation

Build it from scratch

We simulate $N$ ranks in one process and train a tiny linear regression with Adam both ways. Adam is elementwise, so it can update any slice of the parameters independently: that is the property ZeRO relies on.

def adam(p, g, m, v, t, lr=0.05, b1=0.9, b2=0.999, eps=1e-8):
    """One Adam step on any slice of parameters. Elementwise, so it works on a shard."""
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g * g
    p = p - lr * (m / (1 - b1**t)) / (np.sqrt(v / (1 - b2**t)) + eps)
    return p, m, v

Predict: would ZeRO-1 still be exact with an optimizer that normalizes by the norm of the whole gradient, like LAMB?

The two collectives. Reduce-scatter gives rank $r$ the average of chunk $r$ across ranks; all-gather hands every rank every chunk.

def reduce_scatter_mean(grads, n):
    """Rank r receives the average over ranks of chunk r only."""
    chunks = [np.array_split(g, n) for g in grads]
    return [np.mean([chunks[src][r] for src in range(n)], axis=0) for r in range(n)]

def all_gather(shards):
    """Every rank receives every shard, concatenated in rank order."""
    full = np.concatenate(shards)
    return [full.copy() for _ in shards]

Predict: how many floats does each rank receive from reduce_scatter_mean for $P = 10$ and $N = 4$?

The ZeRO-1 step. Each rank allocates Adam state only for its slice. After the local backward, reduce-scatter, update the slice, all-gather.

        if zero:
            g_shard = reduce_scatter_mean(grads, n_ranks)       # comm half 1
            new = []
            for r in range(n_ranks):
                w_shard = np.array_split(w[r], n_ranks)[r]
                p, m[r], v[r] = adam(w_shard, g_shard[r], m[r], v[r], t)
                new.append(p)
            w = all_gather(new)                                 # comm half 2
        else:
            g = np.mean(grads, axis=0)                          # all-reduce (average)
            for r in range(n_ranks):
                w[r], m[r], v[r] = adam(w[r], g, m[r], v[r], t)

Predict: after 50 steps, how different are the DDP and ZeRO-1 weights?

The memory formulas are a few lines. k counts how many of (optimizer, gradients, weights) the strategy shards:

def per_gpu(c, strategy, n, precision="bf16 mixed + AdamW", tokens=2048):
    P = n_params(c)
    bw, bg, bo = PRECISION[precision]
    k = STRATEGIES.index(strategy)          # 0 = DDP, then how many of (optimizer, grads, weights) are sharded
    return dict(
        weights=bw * P / (n if k >= 3 else 1),
        grads=bg * P / (n if k >= 2 else 1),
        optimizer=bo * P / (n if k >= 1 else 1),
        activations=activation_bytes(c, tokens, 4 if precision.startswith("fp32") else 2),
    )

Predict: with fp32 weights and AdamW (no separate master copy, so 4 + 4 + 8 bytes), what is ZeRO-1's floor per parameter?

Use the library

PyTorch FSDP2 shards per module; each fully_shard call becomes one all-gather unit:

from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
for block in model.layers:
    fully_shard(block, mp_policy=mp)                       # ZeRO-3 per transformer block
fully_shard(model, mp_policy=mp, reshard_after_forward=True)

reshard_after_forward=True frees gathered weights after forward (ZeRO-3); False keeps them until backward finishes, close to ZeRO-2. The older FullyShardedDataParallel wrapper names these ShardingStrategy.FULL_SHARD, SHARD_GRAD_OP, and HYBRID_SHARD (shard within a node, replicate across nodes). DeepSpeed selects the stage in its config:

{ "zero_optimization": { "stage": 2, "overlap_comm": true }, "bf16": { "enabled": true } }

DeepSpeed also offers ZeRO-Offload and ZeRO-Infinity, which move optimizer state (and optionally weights) to CPU memory or NVMe, trading speed for capacity.

Run it

uv run code/27-training-at-scale/zero_sharding.py

Key lines:

  max |w_DDP - w_ZeRO1| after 50 steps = 0.00e+00
  Adam state floats across all ranks: DDP 80, ZeRO-1 20  (P = 10, N = 4)
 N = 8 GPUs, bf16 mixed precision + AdamW (16 bytes/param)
  strategy  weights    grads    optim   activ  total GB  comm GB vs DDP
  DDP         13.48    13.48    80.86   0.822    108.63    23.58   1.00x
  ZeRO-1      13.48    13.48    10.11   0.822     37.88    23.58   1.00x
  ZeRO-2      13.48     1.68    10.11   0.822     26.09    23.58   1.00x
  ZeRO-3       1.68     1.68    10.11   0.822     14.30    35.38   1.50x
  ZeRO-3 on  16 GPUs:  72.23 GB per GPU  fits in 80 GB

The weights match DDP exactly, and the Adam state across all ranks shrinks from $N \times 2P$ floats to $2P$.

5. Engineering

Picking a stage.

• Fits with DDP? Use DDP.
• Fits with ZeRO-1 or ZeRO-2? Use them: same communication as DDP, so usually no slowdown. ZeRO-2 is a common default for fine-tuning 7B to 13B models on one 8-GPU server.
• Otherwise ZeRO-3 / FSDP. Budget for 1.5 times DDP's traffic, and keep it overlapped with compute (prefetch the next layer's all-gather while the current layer computes).

Hybrid sharding. Inter-node links are slower than NVLink. Sharding within a node (or a group of nodes) and replicating across groups keeps most traffic on fast links, at the cost of a higher memory floor. Choose the group size from the memory formula: the shard group must be big enough for $16P / N_{\text{group}}$ to fit.

Wrap granularity. In FSDP, each wrapped unit is gathered whole. Wrap per transformer block: wrapping the whole model as one unit gathers every weight at once and loses the memory benefit; wrapping tiny modules creates many small, latency-bound collectives.

Peak memory is higher than the formula. ZeRO-3 temporarily holds a full gathered layer (or two, with prefetch), communication buffers, and the allocator's fragmentation. Leave 10 to 20% headroom and measure with torch.cuda.max_memory_allocated.

Checkpoints are sharded too. Each rank saves its own slice of optimizer state (and weights, under ZeRO-3). Use the framework's distributed checkpoint format (torch.distributed.checkpoint, DeepSpeed's), which records which rank owns what and can reshard when you resume on a different number of GPUs.

Gradient clipping needs a global norm. Under ZeRO-2 and 3 no rank has the whole gradient. Each rank computes the sum of squares of its slice, an all-reduce adds them, and every rank applies the same factor. FSDP's clip_grad_norm_ does this; a per-rank norm would give each slice a different factor.

Common mistakes

• Expecting ZeRO-1 or ZeRO-2 memory to keep falling with more GPUs: the replicated weights and gradients are a $4P$ or $2P$ floor.
• Expecting ZeRO to shrink activations: it does not. Use checkpointing or smaller micro-batches.
• Writing ZeRO-1's memory as $2P + 2P + 8P/N$ (forgetting the fp32 master copy): it is $4P + 12P/N$.
• Saving rank 0's optimizer state only and calling it a checkpoint: the other ranks' slices are lost.
• Running ZeRO-3 across slow inter-node links with no prefetch or hybrid sharding: GPUs wait on all-gathers.
• Clipping by a local per-rank norm under sharded gradients.

6. Knowledge check

Practice exercise. A synthetic 13B-style model has $P = 13 \times 10^9$ parameters, trained with bf16 + AdamW on $N = 16$ GPUs. Compute per-GPU training state (no activations) for DDP, ZeRO-1, ZeRO-2, and ZeRO-3, and say which fit 24 GB GPUs.

Implementation challenge. Extend the script's train function with a stage=2 mode in which each rank keeps only its own gradient slice after the reduce-scatter (free the rest), and count the gradient floats held per rank. Then add global-norm gradient clipping that works on sharded gradients.

g_shard = reduce_scatter_mean(grads, n_ranks)
grads = None                                             # ZeRO-2: full gradients freed
held = [len(s) for s in g_shard]                         # P/N floats per rank
sq = sum(float(np.sum(s * s)) for s in g_shard)          # all-reduce of local sums of squares
scale = min(1.0, max_norm / (np.sqrt(sq) + 1e-6))
g_shard = [s * scale for s in g_shard]                   # same factor on every rank

Summary

• DDP stores $16P$ bytes on every GPU. ZeRO-1 shards the 12 bytes of optimizer state ($4P + 12P/N$), ZeRO-2 adds gradients ($2P + 14P/N$), ZeRO-3 / FSDP adds weights ($16P/N$).
• A 7B-style model on 8 GPUs: 108.63, 37.88, 26.09, and 14.30 GB per GPU including activations.
• ZeRO-1 and ZeRO-2 send the same bytes as DDP (reduce-scatter + all-gather = all-reduce); ZeRO-3 sends about 1.5 times as much because weights are gathered for forward and again for backward.
• Only ZeRO-3 keeps shrinking with more GPUs; a 70B-style model first fits 80 GB cards at 16 GPUs.
• Activations are not sharded; pair ZeRO with gradient checkpointing, sharded checkpoints, and global-norm clipping.

Next: Self-attention

Related: Scaling and distributed training · Gradient checkpointing · LLM inference

Go deeper, code-first: AI Engineering from Scratch, ZeRO stage 1 over real processes.

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): ZeRO Optimizer State Sharding.

Interview angle