Fill in the DeepSpeed config keys that select a ZeRO stage
The key is `zero_optimization.stage`, value 1, 2, or 3. Higher stages shard more training state across GPUs at the cost of more cross-GPU communication.
Imagine eight friends carrying a giant tent to a campsite. The smallest version of the trick is each friend just carries an equal share of the tent poles, and when it is time to pitch, they pass poles around as needed. A bigger version of the trick has them split not just poles but also the assembly instructions, so each person knows their slice. The biggest version splits the tent fabric itself, so no one person ever carries the whole thing. Each step saves more weight on each friend's back, but it also means more passing things back and forth at setup time. The DeepSpeed config has one knob that names which of these three sharing levels you want.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
DeepSpeed's ZeRO is the technique that made it possible to train models with tens of billions of parameters on commodity 8-GPU nodes. The core idea is simple: in plain data-parallel training every GPU holds a redundant copy of the model weights, gradients, and optimizer state. ZeRO observes that this redundancy is wasteful and shards each of those pieces across the data-parallel ranks instead.
The knob that picks how aggressively to shard lives in the DeepSpeed JSON config at the path zero_optimization.stage with integer values 1, 2, or 3. Each value moves one more piece of state from per-rank redundant to sharded across ranks. The benefit is memory; the cost is communication, because sharded state has to be gathered together at the moments when the math needs it.
This deep dive walks through what each stage shards, what communication pattern each stage adds, the practical inflection points where the next stage stops being worth it, and the composition rules with offload, pipeline parallel, and gradient checkpointing.
The three stages and what each shards
Stage 1 shards optimizer state across the data-parallel ranks. For Adam this is the heaviest single memory term: two fp32 moments per trainable parameter, costing 8 bytes per parameter combined. With N ranks each holds 8/N bytes per parameter of moment storage, then applies its slice of the optimizer step locally. The optimizer step now executes in parallel across ranks rather than redundantly.
Stage 2 keeps the stage 1 sharding and additionally shards the gradients. Each rank holds only the gradient slices that match its optimizer slice. The collective operation that produces these sharded gradients is a reduce-scatter, which has the same total bandwidth as the all-reduce stage 1 used but produces sharded outputs instead of replicated outputs. Memory savings stack on top of stage 1; communication cost is essentially unchanged.
Stage 3 additionally shards the parameters themselves. The model weights are split across ranks, and any layer that needs full weights gathers them via an all-gather just before its forward computation, then drops the gathered copy as soon as the layer is done. Backward gather is symmetric. This is what makes 70B-parameter training viable on 80 GB GPUs that could not hold even one copy of the weights, let alone the gradients and moments.
The memory savings are real and large. For a 13B model with Adam in bf16, full data-parallel needs roughly 200 GB per rank. Stage 1 cuts that to about 100 GB. Stage 2 cuts it further to about 75 GB. Stage 3 brings it down to about 30 GB per rank on 8 GPUs, which is the difference between fitting on H100 and not fitting at all.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- Hugging Face Accelerate and TRL ship DeepSpeed config presets at stages 1, 2, and 3 for users fine-tuning Llama 4 Maverick on multi-GPU nodes.
- Microsoft's DeepSpeed examples repository uses stage 3 with parameter offload to train models that do not fit on any single GPU.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhy is the stage 1 to stage 2 transition almost free in communication terms?
Compare the collective operations. Stage 1 already does an all-reduce of gradients; stage 2 converts that to a reduce-scatter of the same bandwidth, which produces sharded outputs as a side effect.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Setting stage 3 by default without measuring throughput. The all-gather of parameters every step can dominate runtime on slow interconnects and erase the memory savings.
60 second bullets to scan on the way to the call.
Exact key path
zero_optimization.stageand its three integer valuesWhat gets sharded at stage 1 versus stage 2 versus stage 3
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.