CapstonePython

ZeRO Optimizer State Sharding

ZeRO stage 1 shards optimiser state across N ranks, so each rank holds 1/N of the Adam moments and the fp32 master copy. Adam stores two moment estimates per parameter, both in float32, which gives a 7B-parameter model 56 GB of optimiser state, the largest single allocation in the training stack. After the local step the updated parameter shards broadcast back, every rank reconstructs the full model, and the next step begins. Memory for that state drops linearly with N.

TIME: ~90 minPREREQUISITES: Phase 19 Track C lessons 42-49

Learning Objectives

Shard optimiser state (first moment, second moment, fp32 master copy) across N ranks so each rank owns 1/N.
Use reduce_scatter to deliver each rank only its shard of the gradient (summed, then divided by world size), then allgather to send the updated parameter shards back to every rank.
Compute the per-rank memory table for vanilla DDP and ZeRO stage 1, and read the stage 2 and stage 3 rows as design numbers.
Defend the choice of stage 1 vs stage 2 vs stage 3 on model size and bandwidth budget.

The Problem

ZeRO stage 1 shards optimiser state: each rank keeps 1/N of the Adam moments and the fp32 master copy instead of a full replica of each.

It starts from what vanilla DDP does, which is replicate everything. Parameters, gradients, and optimiser state sit in full on every rank. For a 7B-parameter model in fp16 that means 14 GB of parameters, 14 GB of gradients, and 84 GB of optimiser state per rank (28 GB of fp32 master copy plus 56 GB of Adam moments). The optimiser state is the largest term and the easiest to shard, because it is only touched during the step, not during forward or backward.

A stage 1 step goes like this. After backward, ZeRO reduce_scatters the gradient instead of allreducing the full thing and stepping locally, so each rank receives only its shard of the gradient (summed across ranks, then divided by world size in the code to give the mean). The rank applies the optimiser step to its shard of the master parameters. The updated parameter shards then allgather back, and every rank has the full model for the next forward. Optimiser memory drops by N. Wire traffic per step matches DDP by bandwidth, since one reduce_scatter plus one allgather equals one allreduce. Memory drops and throughput holds.

Warm-up0/3

1. In plain data-parallel training with Adam, what does each GPU hold a full copy of?Warm-up

2. An allreduce with a sum across 4 ranks finishes. What does each rank hold?Warm-up

3. Adam keeps extra running statistics for every parameter it updates. What are they?Warm-up

Check: The Problem0/3

1. What is ZeRO stage 1?Define

2. Why does ZeRO shard optimiser state before parameters?Explain

3. What is the per-step communication pattern of ZeRO stage 1?Explain

The Concept

Mermaid diagram

Stages of ZeRO

ZeRO comes in three stages. Each one shards another piece of training state across ranks, and the table shows what that costs in memory and communication.

Stage	What is sharded	Memory per rank	Comm per step
DDP	nothing	params + grads + optim	1x allreduce
ZeRO-1	optimiser state	params + grads + optim/N	1x reduce_scatter + 1x allgather
ZeRO-2	optim + grads	params + grads/N + optim/N	1x reduce_scatter + 1x allgather
ZeRO-3	optim + grads + params	params/N + grads/N + optim/N	1x allgather per layer + 1x reduce_scatter per layer

Stage 1 is the cheapest win because optimiser state dominates the budget. Stage 2 needs gradient-shard accumulation logic, with the same bandwidth. Stage 3 (FSDP) pays per-layer comm on each forward and backward in exchange for the parameter-shard memory drop. code/main.py implements stage 1 only. The ZeRO-2 and ZeRO-3 rows describe the design and are not implemented there.

The memory math, real numbers

For a model with P parameters trained with Adam in mixed precision:

Term	Vanilla	ZeRO-1	Why
fp16 params	2P bytes	2P bytes	needed for forward
fp16 grads	2P bytes	2P bytes	needed for backward
fp32 master copy	4P bytes	4P/N bytes	only the optim uses it
fp32 first moment	4P bytes	4P/N bytes	only the optim uses it
fp32 second moment	4P bytes	4P/N bytes	only the optim uses it
Total	16P bytes	4P + 12P/N bytes

The optimiser-state terms add up to 12P bytes (master copy 4P, first moment 4P, second moment 4P). The two moments alone are 8P, which is where the 56 GB for a 7B model in the motto comes from. All three together are 84 GB, and the master copy alone is 28 GB.

At N=8, vanilla needs 16P and ZeRO-1 needs 5.5P, a 65% drop. At N=64, vanilla is still 16P and ZeRO-1 needs 4.19P, a 74% drop.

Why reduce_scatter beats allreduce-then-shard

Reduce_scatter sums a tensor across ranks and hands each rank only one chunk of the result. Allreduce gives every rank the full summed gradient, but a rank that owns only shard r has no use for the other (N-1)/N of it. Reduce_scatter delivers exactly the shard each rank owns. Allreduce is reduce_scatter plus allgather, so the parameter-shard allgather later in the step stands in for the second half. Net wire is identical to DDP, and memory is divided.

Check: The Concept0/12

1. What does ZeRO-2 add over ZeRO-1?Define

2. What does ZeRO-3 shard, and what communication does it need?Define

3. What is the master copy in mixed-precision ZeRO training?Define

4. What does reduce_scatter do to a tensor held on every rank?Define

5. Why does ZeRO-1 use the same per-step bandwidth as DDP?Explain

6. Why is reduce_scatter a better fit for ZeRO-1 than an allreduce followed by keeping one slice?Explain

7. Why does ZeRO-1 still keep 4P bytes of fp16 parameters and gradients on every rank?Explain

8. Using the memory math, what is the per-rank memory drop of ZeRO-1 versus vanilla DDP for Adam mixed precision on 8 ranks?Apply

9. Using ZeRO-1 per-rank memory of 4P + 12P/N against 16P for vanilla DDP, at how many ranks does ZeRO-1 need exactly half of vanilla (8P)?Apply

10. The memory math gives a 65% drop at N=8 and 74% at N=64. Why does multiplying the rank count by 8 add so little?Analyze

11. With ZeRO-1 on 8 ranks the optimiser state fits, but the fp16 parameters themselves no longer fit on one rank. Which design fits, and what does it cost?Design

12. Parameters fit and optimiser state is already sharded, but gradient memory is now the problem, and you want to keep DDP-level bandwidth. Which choice fits?Design

Build It

code/main.py implements:

flat_param_numel(module), gather_flat_params(module) and scatter_flat_to_params(module, flat), which count a model's parameters, pack them into one contiguous fp32 vector, and copy a flat vector back into the parameters. gather_flat_grads(module) does the same for gradients. The flat layout is what makes sharding by rank a simple slice, and shard_bounds pads the vector so it divides evenly by world size.
ZeroOptimizer(module, world_size, rank, lr) that owns the rank's shard of the fp32 master copy and Adam moments.
step() that runs reduce_scatter (sum) on the flat gradient, divides the result by world size so each rank holds the mean gradient for its shard, applies Adam to that shard, and allgathers the updated fp32 master shards back into the full model.
A demo that trains a 3-layer MLP for 20 steps on four torch.distributed gloo ranks. It prints per-step loss for rank 0 and rank 3, the final parameter norm and optim_shard_bytes for every rank, and one memory_table at the end.

The demo runs in fp32 end to end, so no fp16 appears in the code path: gradients are cast to fp32 and the allgather moves fp32 master shards. The fp16 figures in the memory math describe the mixed-precision design the table models, and the module docstring is written in those terms. The docstring also points at lesson 77 for a loss comparison, but the demo runs no DDP baseline.

Run it:

bash

python3 code/main.py

It prints per-step loss, the per-rank shard sizes, and a memory table. The table is arithmetic on a parameter count, not a measurement. Its header says optimiser memory, but each row is the total per-rank bytes (params + grads + optimiser state), with vanilla DDP's full copy against ZeRO-1 holding 1/N of the optimiser state.

Check: Build It0/12

1. What does `gather_flat_params(module)` do in main.py?Define

2. Why does `step()` divide the reduce_scattered gradient by the world size?Explain

3. Why does `shard_bounds` pad the flat vector when the parameter count does not divide by the world size?Explain

4. What does `shard_bounds(10, 4, 3)` return as (start, end, chunk)?Apply

5. A ZeroOptimizer wraps a model with 10 scalar parameters in total on 4 ranks. What does `shard_bytes()` return?Apply

6. `memory_table(1000, 4)` is called. What does it print for ZeRO-1?Apply

7. Four ranks hold the gradient values 2, 4, 6 and 8 for one element, and that element lies in rank 1's shard. What does rank 1's `local_grad` hold for it after `step()` runs reduce_scatter and the division?Apply

8. A parameter received no gradient, so `p.grad` is None. What does `gather_flat_grads` put in the flat vector for it?Apply

9. `memory_table` prints a ZeRO-1 figure well above 12P/N bytes at N=4. Why?Analyze

10. After a ZeRO-1 run, every rank should print the same final parameter norm, but rank 2 prints a different one. The step pads the flat vector and trims the gathered result. Where do you look first?Analyze

11. A teammate says ZeRO-1 cannot save memory because `module.parameters()` still holds the whole model on every rank. What is wrong with that argument?Analyze

12. The demo trains in fp32 end to end, yet memory_table reports a vanilla figure of 16P. A reader expects the demo's real memory to match it. What is the wrong assumption?Analyze

Production patterns in the wild

Three patterns harden ZeRO enough to ship.

Checkpoint the shards. ZeRO-1 splits optimiser state across ranks, so the checkpoint has to record which rank owns what. Lesson 80 builds the sharded checkpoint manifest that resumes a ZeRO run on the same world size. Without it, the saved state can't be read back at restart.

Pair it with mixed precision. ZeRO is a mixed-precision technique, and the fp32 master copy is what gets sharded. Running ZeRO without mixed precision pays the memory tax on the fp32 master without the corresponding fp16 forward win. In practice, runs pair ZeRO with autocast or bf16 weights. The demo here skips that and runs fp32 throughout.

Start at stage 1. The comm is identical to DDP by bandwidth, the memory savings are linear in N, and the only cost is the bookkeeping for the optimiser shard. Stacks usually default to stage 1 unless parameter shard memory is also a problem. Then stage 2 or 3 trades comm for memory.

Check: Production patterns in the wild0/3

1. Why must a ZeRO checkpoint record which rank owns which shard?Explain

2. Why should ZeRO be paired with mixed precision?Explain

3. You train a 7B model on 8 ranks. Parameters and gradients fit, but the Adam state runs out of memory. Which change fits best?Design

Use It

Where ZeRO shows up:

DeepSpeed ZeRO. The reference implementation. deepspeed_config.json selects stage 1/2/3 and partition sizes.
PyTorch FSDP. The PyTorch-native equivalent. ShardingStrategy.SHARD_GRAD_OP is ZeRO-2, and FULL_SHARD is ZeRO-3.
HuggingFace Accelerate. Wraps both DeepSpeed and FSDP under a uniform config.

Check: Use It0/1

1. PyTorch FSDP's `FULL_SHARD` strategy shards parameters, gradients and optimiser state. Which ZeRO stage is that?Explain

Ship It

Lesson 79 (pipeline parallel) shards along the orthogonal axis. Instead of sharding optimiser state across the same model, it shards layers across ranks. Lesson 81 composes DDP + ZeRO in the end-to-end demo.

Exercises

Extend to ZeRO-2 by sharding gradients: each rank only stores the gradient for its shard, achieved by zeroing out the non-shard portion after backward.
Add a memory profiler that prints actual fp32 byte usage on rank 0 versus the formula prediction.
Measure the per-step wall-clock time of vanilla DDP versus ZeRO-1 and decompose into forward, backward, comm.
Implement gradient clipping under ZeRO-1: the L2 norm must be computed across all shards via allreduce of the local norm squared.
Implement a "naive ZeRO" with allreduce instead of reduce_scatter, measure the wire-time difference. Defend the reduce_scatter choice with numbers.

Key Terms

Term	What people say	What it actually means
ZeRO-1	"Shard the optimiser"	Each rank holds 1/N of fp32 master + Adam moments
ZeRO-2	"Shard grads too"	Each rank also drops the non-shard gradients after reduce_scatter
ZeRO-3	"Shard params"	Each rank holds 1/N of fp16 params; allgather per layer in forward
Master copy	"fp32 weights"	The high-precision parameter copy the optimiser updates
Reduce_scatter	"Split the sum"	Deliver each rank only its shard's summed gradient

Author's Note

My bet is that the per-rank optim_shard_bytes line deserves a closer read than the memory table. memory_table is arithmetic on a parameter count. optim_shard_bytes comes from shard_bytes(), which counts the master, m and v tensors the rank actually holds, at 4 bytes each. I'd want the measured number sitting next to the formula before trusting either.

I'd start at stage 1 and move up only when parameter or gradient memory becomes the problem. code/main.py implements stage 1 alone, so ZeRO-2 is left to exercise 1, and the stage 3 trade of per-layer comm for memory stays on paper here.

If something breaks, I'd check the shard bounds first. shard_bounds pads when the parameter count doesn't divide by world size, and step() trims the allgathered vector back to self.total. The final parameter norm printed per rank is the quick check that every rank rebuilt the same model.