ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He2020SC 2020
Read it on arxiv.org(opens in a new tab)Why this one
Read this after Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. Megatron shows how to split the math of a giant transformer. ZeRO shows the quieter problem that often stops you first: every GPU is carrying duplicate optimizer state, gradients, and parameters until memory, not flops, sets the ceiling. The mistake is to think ZeRO is just another distributed training trick. Its real value is the accounting. It teaches you to itemize model memory like an engineer, then remove redundancy stage by stage without turning the whole program into a hand-partitioned science project. If you are fine-tuning or pretraining large models, this is the paper that makes GPU OOM errors feel less like weather and more like a budget you can reason about.
What to take away
- Optimizer states can take more memory than the parameters, especially with Adam.
- ZeRO partitions optimizer state, gradients, and parameters across data-parallel workers.
- The paper gives a practical memory ledger before it gives a distributed systems design.
Reads with
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
megatron splits transformer compute before zero removes replicated training state
- Mixed Precision Training
precision cuts the tensor sizes that zero then accounts for and shards
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
trillion-parameter sparse models need zero-style state sharding to fit training