Explain ZeRO stages 1, 2 and 3. How much memory does each stage leave per GPU for a 70B model?

TL;DR: Mixed-precision Adam holds 16 bytes per parameter: bf16 weights (2), bf16 gradients (2), fp32 master weights (4), fp32 first moment (4) and fp32 second moment (4). ZeRO shards these across the n data-parallel ranks in three stages: stage 1 shards the 12 bytes of optimizer state, stage 2 also shards the 2 bytes of gradients, stage 3 also shards the 2 bytes of weights. For a 70B on 64 GPUs that is 296 GB, 157 GB and 18 GB per rank respectively, so only stage 3 fits on an 80 GB card.

How to approach it

Write the 16-byte decomposition first, because every stage is a statement about which of those five terms is divided by n. Then define each stage as "what is still replicated" rather than "what is sharded", since the replicated part is what limits memory. Compute the per-rank bytes for one model and one world size, then state what each stage costs in communication, because stage 3 is the one that changes the collective pattern.

A strong answer

A typical situation: a team wants to fine-tune Llama 3.1 70B on a handful of 8 × H100 nodes with plain data parallelism, and the first attempt runs out of memory before the first step. The arithmetic says it had to.

inputs:  N = 70.6e9 parameters
         mixed-precision Adam state per parameter:
           bf16 weights          2 B
           bf16 gradients        2 B
           fp32 master weights   4 B
           fp32 Adam m           4 B
           fp32 Adam v           4 B
                                16 B total

static training state = 70.6e9 × 16 B = 1.13e12 B ≈ 1,130 GB

sanity: one H100 has 80 GB and a node has 640 GB, so the replicated state does not fit on a
        node, let alone a card. Data parallelism without sharding cannot train this model.

ZeRO (Zero Redundancy Optimizer, from DeepSpeed) observes that under data parallelism every one of the n ranks holds an identical copy of those 1,130 GB, and that the redundancy can be removed one term at a time. ZeRO and FSDP has the mechanics; here is what each stage keeps per rank.

Stage 1 shards the optimizer state. Each rank owns the fp32 master weights and both Adam moments for 1/n of the parameters, and updates only those. Weights and gradients stay replicated, so the forward and backward passes are unchanged and the all-reduce of gradients is unchanged. After the update, each rank broadcasts (all-gathers) its slice of the updated bf16 weights.

Stage 2 also shards the gradients. Instead of all-reducing the full gradient so every rank has all of it, the ranks reduce-scatter it so each rank receives only the reduced gradient for the 1/n of the parameters whose optimizer state it owns. The total bytes on the wire are unchanged from stage 1 (a reduce-scatter is half of an all-reduce, and the all-gather of updated weights is the other half), but no rank ever holds a full gradient.

Stage 3 also shards the weights. Each rank holds only 1/n of every layer's parameters at rest. Before a layer runs forward, the ranks all-gather that layer's weights, use them, and free them; the same all-gather happens again in backward. This is the stage that changes the per-step traffic: the weights are gathered twice per step on top of the gradient reduce-scatter, so the total is about 3 × model bytes per rank versus the 2 × model bytes of a plain all-reduce.

per-rank bytes per parameter, n data-parallel ranks:
  stage 0 (DDP):  2 + 2 + 12       = 16
  stage 1:        2 + 2 + 12/n
  stage 2:        2 + (2 + 12)/n   = 2 + 14/n
  stage 3:        16/n

for N = 70.6e9:
  n     stage 1        stage 2        stage 3
   8    2+2+1.5 = 5.5  → 388 GB     2+1.75 = 3.75 → 265 GB     2.0   → 141 GB
  16    4.75           → 335 GB     2.875         → 203 GB     1.0   →  71 GB
  32    4.375          → 309 GB     2.4375        → 172 GB     0.5   →  35 GB
  64    4.1875         → 296 GB     2.21875       → 157 GB     0.25  →  18 GB

sanity: stages 1 and 2 never get below the 4 B/param of the replicated bf16 weights and
        gradients, which is 282 GB for a 70B; neither fits an 80 GB card at any n.
        Stage 3 at n = 16 is 71 GB, which is inside the card but leaves nothing for
        activations; n = 32 (four nodes) is the realistic floor for this model.

The table also explains when the earlier stages are worth using. For a model whose 16 bytes per parameter fits on the card at n = 1, stages 1 and 2 remove memory pressure at no communication cost and are close to free. Stage 1 alone cuts a 7B model's state from 112 GB to 28 GB plus 84/n, which on 8 GPUs is 38.5 GB, enough to turn an out-of-memory into a run. Stage 3 is the tool for models that do not fit at all, and its extra all-gather traffic is the price.

A second cost of stage 3 is the shape of the traffic, not the volume. The all-gather happens per layer, right before that layer needs its weights, so it must be prefetched or the GPU waits. FSDP's forward_prefetch and backward_prefetch and DeepSpeed's stage3_prefetch_bucket_size exist for this; a stage 3 run with prefetch off can lose a third of its throughput to waiting.

Decision: stage 3 (or FSDP, its PyTorch equivalent) for any model whose training state is larger than a card, at the smallest n that leaves room for activations; stage 1 or 2 for models that fit, as a memory cushion. Reversal condition: if the network cannot carry three model-sizes of traffic per step within the compute time, stage 3 across nodes loses to tensor or pipeline parallelism, and the answer becomes a 3D layout.

The reversal condition: a model that fits comfortably in stage 2, where stage 3's all-gather of parameters on every forward buys memory you did not need and costs throughput you did. Model Memory Footprint is the arithmetic that decides which stage you are in. nvidia-smi --query-gpu=memory.used per rank tells you which stage you are actually running versus which one you configured.

What interviewers probe next

• "Why fp32 master weights at all?" bf16 has 8 bits of mantissa, so a weight update smaller than about 1/256 of the weight rounds to zero; the fp32 copy accumulates small updates and is rounded to bf16 for the forward pass.
• "Does stage 3 change the math?" No; every rank computes the same forward and backward it would under DDP, it just fetches the weights first. Results are bitwise identical to DDP up to reduction order.
• "Where do activations go in this table?" On top of it, and they scale with micro-batch and sequence rather than with n, which is why the 18 GB at n = 64 is a floor and the real number per rank is 40 to 60 GB.
• "Can you shard the optimizer state without DeepSpeed?" PyTorch's ZeroRedundancyOptimizer is stage 1; FSDP with SHARD_GRAD_OP is stage 2 and FULL_SHARD is stage 3.

Common mistakes

• Quoting "ZeRO-3 makes memory 1/n" and forgetting that activations are not sharded by n at all.
• Believing stage 3 increases communication by a large factor; it is about 1.5× the plain all-reduce, and it is overlappable.
• Treating stage 2 as a memory win for models that do not fit at stage 1; the replicated 4 bytes per parameter of weights and gradients are the same in both, and for a 70B that is 282 GB before anything else.
• Forgetting the Adam moments are fp32 and computing 8 bytes per parameter.

Key takeaways

• 16 B per parameter: 2 weights + 2 gradients + 4 master + 4 m + 4 v. A 70B carries 1,130 GB of static state.
• Stage 1 shards the 12 B of optimizer state; stage 2 adds the gradients; stage 3 adds the weights. Per rank: 2 + 2 + 12/n, 2 + 14/n, 16/n.
• For a 70B only stage 3 fits an 80 GB card, and only at n ≥ 16; n = 32 is the practical floor with activations.
• Stage 3 trades an all-reduce for a reduce-scatter plus two all-gathers, about 1.5× the bytes, which must be prefetched to stay hidden.