ZeRO Stage 1/2/3

ZeRO (Zero Redundancy Optimizer) is the single most important memory optimization for training. It eliminates redundant memory across data-parallel ranks by sharding optimizer states, gradients, and parameters — without changing model semantics or convergence.

---

1. The Redundancy Problem

In standard data parallelism, every GPU holds a complete copy of everything:

Component	Per GPU	8 GPUs total	Redundancy
Parameters (BF16)	14 GB	112 GB	8×
Gradients (BF16)	14 GB	112 GB	8×
Optimizer (FP32)	84 GB	672 GB	8×
Total	112 GB	896 GB	8×

For a 7B model on 8 GPUs: 896 GB of aggregate memory, but only 112 GB of unique data. 784 GB is pure waste.

---

2. ZeRO Stage 1: Shard Optimizer States

What’s sharded: Optimizer states ($m$, $v$, master weights) — each GPU holds 1/N of the optimizer.

Communication added: AllGather the updated parameters after each optimizer step.

Before (standard DP)	After (ZeRO-1)
Per GPU: 2Ψ + 2Ψ + 12Ψ = 16Ψ	Per GPU: 2Ψ + 2Ψ + 12Ψ/N
7B, 8 GPUs: 112 GB per GPU	7B, 8 GPUs: 14 + 14 + 10.5 = 38.5 GB

Memory savings: 112 → 38.5 GB per GPU (2.9× reduction). Extra communication: One AllGather of parameters per step (~14 GB).

Why Optimizer First?

Optimizer states are 75% of training memory but are only needed during the optimizer step (once per training step). During forward and backward, they’re idle. Sharding them is high-ROI with minimal communication cost.

Worked Example: 13B Model on 8 GPUs

Parameters: 13B. Static memory = 16 \ 13B = 208 GB.

Component	Stage 0 (DDP)	Stage 1	Stage 2	Stage 3
Params (BF16)	26 GB	26 GB	26 GB	26/8 = 3.25 GB
Grads (BF16)	26 GB	26 GB	26/8 = 3.25 GB	3.25 GB
Optimizer (FP32)	156 GB	156/8 = 19.5 GB	19.5 GB	19.5 GB
Per GPU	208 GB	71.5 GB	48.75 GB	26 GB
Fits on 80 GB?	No	Yes (tight)	Yes	Yes (plenty)

• Stage 0: 208 GB per GPU — doesn’t fit. Need at least 3 GPUs just for static memory.
• Stage 1: 71.5 GB — fits on 80 GB, with 8.5 GB left for activations.
• Stage 2: 48.75 GB — comfortable, room for larger batch size.
• Stage 3: 26 GB — maximum flexibility for batch size and sequence length.

---

3. ZeRO Stage 2: + Shard Gradients

What’s sharded: Optimizer states AND gradients.

Each GPU accumulates gradients only for its assigned parameter partition. After backward, a ReduceScatter (instead of AllReduce) distributes the averaged gradient shards.

	ZeRO-1	ZeRO-2
Per GPU	2Ψ + 2Ψ + 12Ψ/N	2Ψ + 2Ψ/N + 12Ψ/N
7B, 8 GPUs	38.5 GB	14 + 1.75 + 10.5 = 26.25 GB

Communication: Same as standard DP (ReduceScatter + AllGather ≈ AllReduce). No extra communication vs ZeRO-1!

---

4. ZeRO Stage 3: + Shard Parameters

What’s sharded: Everything — optimizer states, gradients, AND parameters.

Each GPU holds only 1/N of the parameters. Before each layer’s forward/backward, the full parameters are gathered (AllGather), used, and discarded.

	ZeRO-2	ZeRO-3
Per GPU	2Ψ + 2Ψ/N + 12Ψ/N	2Ψ/N + 2Ψ/N + 12Ψ/N = 16Ψ/N
7B, 8 GPUs	26.25 GB	112/8 = 14 GB

Communication: AllGather parameters before each layer (forward AND backward) + ReduceScatter gradients. This is 1.5× the communication of standard DP.

Breaking it down:

• DDP (standard DP): AllReduce = 2Ψ bytes per step
• ZeRO-3: Forward AllGather (Ψ) + Backward AllGather (Ψ) + ReduceScatter (Ψ) = 3Ψ bytes per step
• Ratio: 3Ψ / 2Ψ = 1.5×

The ZeRO Stages Summary

Stage	Shards	Memory per GPU	Communication vs DP	Best for
0 (DP)	Nothing	16Ψ	1× (AllReduce)	Small models that fit
1	Optimizer	4Ψ + 12Ψ/N	~1×	Most training (low overhead)
2	+ Gradients	2Ψ + 14Ψ/N	1×	Same comm as DP, less memory
3	+ Parameters	16Ψ/N	~1.5×	Large models that don’t fit otherwise

---

5. DeepSpeed ZeRO Configuration

# ZeRO Stage 1 (most common)
ds_config = {
    "zero_optimization": {
        "stage": 1,
    },
    "bf16": {"enabled": True},
    "optimizer": {
        "type": "AdamW",
        "params": {"lr": 1e-4, "weight_decay": 0.01}
    }
}

# ZeRO Stage 2
ds_config["zero_optimization"] = {
    "stage": 2,
    "allgather_bucket_size": 5e8,       # bucket size for AllGather
    "reduce_scatter": True,              # use ReduceScatter (always True)
    "reduce_bucket_size": 5e8,
}

# ZeRO Stage 3
ds_config["zero_optimization"] = {
    "stage": 3,
    "param_persistence_threshold": 1e5,  # params < this stay replicated
    "prefetch_bucket_size": 5e7,         # prefetch next layer's params
    "sub_group_size": 1e9,               # partition granularity
}

Stage Selection Guide

Scenario	Recommended	Why
Model fits with standard DP	ZeRO-0 or 1	Minimal overhead
OOM with standard DP, fits with optimizer sharding	ZeRO-1	Best memory/comm tradeoff
Still OOM, or want larger batch size	ZeRO-2	Same comm as DP, less memory
Model doesn’t fit even with sharded optimizer	ZeRO-3	Only option, 3× comm cost
Extremely large model (>100B per GPU)	ZeRO-3 + offloading	See CPU & NVMe Offloading

---

6. ZeRO vs Tensor Parallelism

Both reduce per-GPU memory, but they work differently:

	ZeRO-3	Tensor Parallelism
What’s split	Full copies, gathered on-demand	Actual matrix dimensions
Communication	AllGather per layer (parameter size)	AllReduce per layer (activation size)
Compute efficiency	Each GPU does full layer compute	Each GPU does partial layer compute
Network requirement	Can use InfiniBand (lower BW)	Needs NVLink (high BW, low latency)
Max practical degree	Hundreds of GPUs	8 GPUs (NVLink group)

In practice: Use TP within a node (NVLink) and ZeRO across nodes (InfiniBand). They’re complementary, not competing.

For PyTorch’s native ZeRO-3 implementation, see FSDP.

---

Self-Check

• Calculate per-GPU memory for each ZeRO stage for a 13B model on 16 GPUs
• Explain why sharding optimizer states first is highest-ROI
• Compare communication overhead of ZeRO stages 1, 2, 3 vs standard DP
• Choose the right ZeRO stage for a given model/cluster configuration
• Explain why ZeRO-3 and TP are complementary (not competing)

---

References

• ZeRO: Memory Optimizations (Rajbhandari et al., 2019)
• DeepSpeed Documentation