ZeRO, FSDP & Sharded Training
Eliminating memory redundancy in distributed training by sharding optimizer states, gradients, and model parameters.
The Memory Redundancy Problem in DDP
Standard Distributed Data Parallelism (DDP) is memory inefficient:
In DDP across $N = 64$ GPUs, every single GPU stores an identical copy of:
- Model Parameters: $1\times$ Memory.
- Gradients: $1\times$ Memory.
- Adam Optimizer States: $4\times$ Memory ($16\text{ bytes per param}$ for FP32 master weights, momentum, and variance).
Duplicating this $6\times$ memory overhead identically across 64 GPUs wastes massive VRAM capacity.
ZeRO (Zero Redundancy Optimizer - Microsoft DeepSpeed, 2020) and Fully Sharded Data Parallel (PyTorch FSDP) eliminate memory redundancy by sharding state memory across GPUs.
┌─────────────────────────────────────────────────────────────┐
│ 1. ZeRO STAGE 1: OPTIMIZER STATE SHARDING (ZeRO-1) │
│ Shards Adam optimizer states across N GPUs. Cuts RAM by 4x! │
├─────────────────────────────────────────────────────────────┤
│ 2. ZeRO STAGE 2: GRADIENT SHARDING (ZeRO-2) │
│ Shards optimizer states AND gradients across N GPUs. │
├─────────────────────────────────────────────────────────────┤
│ 3. ZeRO STAGE 3 / FSDP: PARAMETER SHARDING (ZeRO-3) │
│ Shards optimizer states, gradients, AND model parameters! │
│ 100% Zero Memory Redundancy! Memory scales 1/N perfectly! │
└─────────────────────────────────────────────────────────────┘
How ZeRO Stage 3 / FSDP Works Just-in-Time
If model parameters are sharded so GPU 0 holds only $1/N$ of the model, how does GPU 0 run a forward pass through Layer 15?
Using Just-In-Time (JIT) Parameter Prefetching via AllGather:
FORWARD PASS AT LAYER 15:
1. GPU 0 executes AllGather ──► Fetches Layer 15 parameters from peer GPUs just-in-time!
2. Compute Layer 15 Forward Pass activations.
3. IMMEDIATELY DROP Layer 15 parameters from GPU 0 memory!
BACKWARD PASS AT LAYER 15:
1. GPU 0 executes AllGather ──► Fetches Layer 15 parameters again just-in-time!
2. Compute Layer 15 Backward Pass gradients.
3. Execute ReduceScatter ──► Send computed gradients to their assigned shard owner GPUs!
4. IMMEDIATELY DROP Layer 15 parameters from GPU 0 memory!
Stage Comparison Matrix
| Strategy | Sharded Components | VRAM Reduction per GPU | Communication Volume |
|---|---|---|---|
| Standard DDP | None (100% Duplicate) | $1\times$ (Baseline) | Low (AllReduce Gradients) |
| ZeRO-1 | Optimizer States | $4\times$ Reduction | Low (Same as DDP) |
| ZeRO-2 | Optimizer States + Gradients | $5\times$ Reduction | Low (Same as DDP) |
| ZeRO-3 / FSDP | Optimizer + Gradients + Parameters | $N\times$ Linear Reduction! | $+50%$ Communication Overhead |
DeepSpeed vs PyTorch FSDP
- DeepSpeed ZeRO-3 (Microsoft): Independent open-source library integrated into HuggingFace Accelerate. Includes CPU Offloading (
ZeRO-Offload) to spill optimizer states to host RAM. - PyTorch FSDP (Native PyTorch): Built-in native PyTorch API (
torch.distributed.fsdp). Offers lower overhead and cleaner integration for PyTorch 2.x training pipelines.
Say this out loud
ZeRO and PyTorch FSDP eliminate memory redundancy in data parallel training by sharding optimizer states, gradients, and model parameters across GPUs. ZeRO 3 / FSDP shards 100 percent of model states, using AllGather communication to fetch layer parameters just-in-time during forward and backward passes, enabling training of massive LLMs with linear VRAM reduction.
Followups to expect
- What is ZeRO-Offload / FSDP CPU Offloading? Offloading sharded optimizer states and gradient updates to CPU System RAM, allowing 13B parameter fine-tuning on a single NVidia RTX 4090 GPU.
- How does FSDP handle Activation Checkpointing? FSDP wraps activation checkpointing around individual sharded FSDP modules, discarding intermediate layer activations to achieve maximum memory efficiency.
Check yourself
What primary memory inefficiency in standard Distributed Data Parallelism (DDP) does ZeRO / FSDP eliminate?