ZeRO and Sharded Optimizer States
A technique for training huge neural networks across many GPUs by splitting up the optimizer's memory-hungry bookkeeping instead of copying it identically onto every device.
Prerequisites: Adam and Adaptive Optimizers
Ordinary data-parallel training gives every GPU an identical, complete copy of the model's weights, gradients, and optimizer state (like Adam's running averages), even though each GPU only ever needs its own slice to do its share of the work — like ten cooks each preparing a full ten-course meal instead of splitting the courses. For a large model, that redundant state can dwarf the weights themselves, since Adam alone stores two extra numbers per parameter, so training can run out of memory on every GPU well before any device is doing useful extra work.
ZeRO (Zero Redundancy Optimizer) removes the redundancy in stages: stage 1 shards only the optimizer state across GPUs so each device holds just its 1/N slice, stage 2 also shards the gradients, and stage 3 shards the model weights themselves, briefly reconstructing each layer's full weights via communication only at the moment it's needed.
Take a 7-billion-parameter model trained in mixed precision across 8 GPUs. Adam's optimizer state alone needs roughly 56 GB total. Without sharding, all 56 GB sits on every GPU; with ZeRO stage 1 across 8 GPUs, each device holds about 7 GB — an 8x reduction on that component, freeing memory for larger batches or bigger models on the same hardware.
ZeRO eliminates the redundant, identical copies of optimizer state, gradients, and weights that ordinary data parallelism keeps on every GPU, sharding each instead so N devices together hold roughly 1/N of the memory any one of them would otherwise need.
Further reading
- Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2020)