ZeRO Optimizer — AI Glossary

Reference details

Related terms

• DeepSpeed
• FSDP (Fully Sharded Data Parallel)
• Gradient Checkpointing

How it works

1. Ordinary data-parallel training ends each step with an all-reduce over gradients, so there is already a full pass of gradient-sized traffic in the baseline.
2. Stage one shards only the optimizer states, and because each worker now updates just its own slice, the all-reduce is replaced by a reduce-scatter plus an all-gather of the updated weights.
3. Those two collectives move the same total volume as the all-reduce they replace, which is why the first stage delivers the largest memory saving for essentially no additional communication.
4. Stage two shards the gradients as well, and since the reduce-scatter was already producing per-worker gradient slices, this is a bookkeeping change rather than new traffic.
5. Stage three shards the parameters, and this one is genuinely different: full layers must now be gathered before use and released after, adding a parameter-sized all-gather per layer per direction that did not exist before.
6. Offload extends the ladder past the third stage by moving shards to host memory or disk, trading a further order of magnitude of bandwidth for capacity — so there are four rungs rather than three, and only the first two are cheap.

Example: Stage one removes three quarters of it

Mixed-precision training with Adam holds about 16 bytes per parameter: 2 for the fp16 weights, 2 for the gradients, and 12 for the fp32 master copy plus the two optimizer moments. The optimizer states alone are 12 of those 16 — 75% of the per-parameter memory. Stage 1 shards exactly that, so it delivers the largest single reduction and costs the least communication. Stage 2 adds gradient sharding, stage 3 adds the parameters and buys the most memory at the highest network cost. Work down the stages: take stage 1, measure, and only escalate when the model still does not fit.

Commonly confused with

ZeRO is the technique; DeepSpeed is the library that introduced it and PyTorch FSDP is a native implementation of the same idea. It is also not model parallelism: the computation is not distributed, only the storage — full layers are reconstructed on demand and then discarded, which is why the communication pattern is all-gather rather than the per-layer synchronisation tensor parallelism needs.

When to use it

Reach for it when:

• Training where optimizer state, not the weights themselves, is what exceeds memory
• Starting at stage 1 and escalating only as far as the model actually requires
• Fast interconnect, where the gather traffic of the higher stages is affordable

Reach for something else when:

• Parameter-efficient fine-tuning, which removes the optimizer-state problem outright
• Jumping to stage 3 by default — you pay maximum communication for memory you may not need
• Inference, which has entirely different memory characteristics

Primary source

• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arXiv)

Continue your research

Use these connected ModelRefs sections to compare alternatives, inspect implementation paths, and review the evidence and governance boundaries relevant to ZeRO Optimizer — AI Glossary.

• Browse models
• Review providers
• Inspect benchmarks
• Explore workflows
• Learn AI concepts
• Read decision guides
• Review trust and methodology
• Start a decision path