ParallelismOptimization track
ZeRO & FSDP
Read it in Chapter 4 · checkpoint +50 XP →
Shard optimizer state, gradients and parameters across data-parallel ranks instead of replicating them.
Plain Data parallelism (DDP) wastes memory: every rank stores identical copies of everything. ZeRO (and PyTorch's FSDP) shards it:
| Stage | Sharded | Per-GPU state (mixed-precision Adam) |
|---|---|---|
| 0 | nothing | 16 B/param |
| 1 | optimizer states | 4 + 12/N B/param |
| 2 | + gradients | 2 + 14/N B/param |
| 3 / FSDP | + parameters | 16/N B/param |
The catch with stage 3: parameters must be all-gathered before each layer's forward and backward, every micro-batch. For a 70B model that's ~140 GB moved twice per micro-batch. On NVLink that's fine. Over a 400 Gb/s NIC it can dominate the step.
Practical guidance
- ZeRO-1 is nearly free and almost always worth it.
- ZeRO-3/FSDP is excellent within a node or on fast fabrics, and for fitting a model that otherwise wouldn't fit.
- At large scale, mix it: TP inside the node, ZeRO-1 or hybrid sharding (shard within a node, replicate across nodes).
Use it in a mission
Related
Sources & further reading
- ZeRO (Rajbhandari et al., 2020)
- Distributed Training (T. Tsugriy) — parallelism strategies, collectives, memory
Written for LazyLab in our own words; sources are cited for depth and verification.