ZeRO: Zero Redundancy Optimizer

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Memory Breakdown of Standard Data Parallelism
5. ZeRO Stage 1: Partition Optimizer States
6. ZeRO Stage 2: Partition Gradients
7. ZeRO Stage 3: Partition Parameters
8. Communication Analysis
9. ZeRO-Offload: CPU and NVMe
10. FSDP: PyTorch's ZeRO-3
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Compute the memory breakdown for mixed-precision training with AdamW.
2. Derive memory savings for each ZeRO stage.
3. Analyze the communication overhead vs standard data parallelism.
4. Explain ZeRO-Offload's GPU-CPU-NVMe memory hierarchy.
5. Compare ZeRO-3 and FSDP implementations.

Notation

• $\Psi$ — model parameters (count)
• $N_d$ — data parallel degree
• $K$ — optimizer state memory multiplier (12 for AdamW mixed-precision)

Core Intuition

Standard data parallelism replicates the entire model (parameters, gradients, optimizer states) on every GPU. For a 70B model: each GPU needs ~280 GB just for optimizer states — impossible. ZeRO partitions these redundant copies across GPUs: each GPU stores only $1/N_d$ of the states, gathering them on-demand via communication.

Memory Breakdown of Standard Data Parallelism

For mixed-precision training with AdamW and $\Psi$ parameters:

• Parameters (FP16): $2\Psi$ bytes
• Gradients (FP16): $2\Psi$ bytes
• Optimizer states:
  • FP32 parameters copy: $4\Psi$ bytes
  • FP32 momentum: $4\Psi$ bytes
  • FP32 variance: $4\Psi$ bytes
  • Total optimizer: $12\Psi$ bytes

Total per GPU: $2\Psi + 2\Psi + 12\Psi = 16\Psi$ bytes.

For 7B model: $16 \times 7 \times 10^9 = 112$ GB per GPU (without activations).

ZeRO Stage 1: Partition Optimizer States

Each GPU stores optimizer states for only $1/N_d$ of the parameters.

Memory per GPU: $2\Psi + 2\Psi + 12\Psi/N_d = 4\Psi + 12\Psi/N_d$.

For $N_d = 8$: $4\Psi + 1.5\Psi = 5.5\Psi$ (vs $16\Psi$). 2.9x reduction.

Communication: Same as standard DP (all-reduce gradients). Each GPU updates only its partition, then all-gathers the updated parameters.

ZeRO Stage 2: Partition Gradients

Additionally partition gradients: each GPU keeps gradients only for its optimizer partition.

Memory per GPU: $2\Psi + 2\Psi/N_d + 12\Psi/N_d = 2\Psi + 14\Psi/N_d$.

For $N_d = 8$: $2\Psi + 1.75\Psi = 3.75\Psi$. 4.3x reduction.

Communication: Replace all-reduce with reduce-scatter (each GPU receives reduced gradients only for its partition). Then all-gather updated parameters.

Total communication: same volume as standard DP ($2\Psi$ bytes per step).

ZeRO Stage 3: Partition Parameters

Partition everything: parameters, gradients, and optimizer states.

Memory per GPU: $16\Psi / N_d$.

For $N_d = 64$: $16\Psi/64 = 0.25\Psi$ per GPU.

A 70B model: only $0.25 \times 70 \times 10^9 \times 1 = 17.5$ GB per GPU. Fits on a single 80GB GPU with room for activations!

Communication: Parameters must be all-gathered for both forward and backward passes. Total: $3 \times 2\Psi/N_d \times N_d = 6\Psi$ bytes per step (3x more than standard DP).

Communication Analysis

ZeRO-3 tradeoff: $N_d\times$ less memory but $3\times$ more communication. Worthwhile when memory is the bottleneck (which it usually is for large models).

ZeRO-Offload: CPU and NVMe

ZeRO-Offload: Move partitioned optimizer states to CPU memory (much larger: 512GB–2TB).

ZeRO-Infinity: Extend to NVMe SSDs (essentially unlimited storage):

• Parameters: GPU or CPU.
• Optimizer states: CPU or NVMe.
• Gradients: GPU (compute) → CPU (accumulate).

Data flow: GPU computes forward/backward → gradients offloaded to CPU → CPU performs optimizer step → updated parameters prefetched back to GPU.

Overlap: Computation on micro-batch $n+1$ overlaps with CPU optimizer step for micro-batch $n$.

FSDP: PyTorch's ZeRO-3

Fully Sharded Data Parallelism: PyTorch-native implementation of ZeRO-3.

Key differences from DeepSpeed ZeRO-3:

• Uses PyTorch's distributed primitives (process groups).
• Finer-grained sharding (per-parameter, not per-layer).
• Better integration with PyTorch ecosystem (autograd, checkpointing).
• Supports mixed precision via native PyTorch AMP.

Sharding strategy:

• FULL_SHARD: ZeRO-3 (shard everything).
• SHARD_GRAD_OP: ZeRO-2 (shard gradients + optimizer).
• NO_SHARD: Standard DDP.

Common Pitfalls

Pitfall 1. Using ZeRO-3 with small batch sizes. The all-gather communication per step is $3\times$ standard DP. If batch/compute per step is small, communication dominates.

Pitfall 2. Not overlapping communication with computation. ZeRO-3 must prefetch the next layer's parameters while computing the current layer. Without overlap: training is 2-3x slower.

Pitfall 3. ZeRO-Offload with slow CPUs. The CPU optimizer step must keep up with GPU forward/backward. Insufficient CPU cores or slow memory bandwidth creates a bottleneck.

Summary

• ZeRO-1: Partition optimizer states → ~3x memory reduction, same communication.
• ZeRO-2: Additionally partition gradients → ~4x reduction, same communication.
• ZeRO-3/FSDP: Partition everything → $N_d\times$ reduction, 3x communication.
• ZeRO-Offload: Leverage CPU/NVMe for optimizer states when GPU memory is exhausted.
• FSDP = PyTorch-native ZeRO-3 with fine-grained sharding.

Exercises

Exercise 1. For a 13B model with $N_d=8$: compute per-GPU memory for standard DP, ZeRO-1, ZeRO-2, and ZeRO-3.

Exercise 2. Derive the communication volume for ZeRO-3 (show why it's $3 \times 2\Psi$ per step).

Exercise 3. For a training step taking 100ms compute on 8 GPUs with 400 GB/s inter-GPU bandwidth: compute the communication overhead for ZeRO-3.

Exercise 4. Design an offload strategy for training a 175B model on 8x A100-80GB GPUs (explain what goes where).

Exercise 5. Compare FSDP FULL_SHARD vs SHARD_GRAD_OP for a 7B model on 4 GPUs with 80GB each. Which is appropriate?