Fill in the DeepSpeed config keys that select a ZeRO stage

Fill in blank·Easy·4.0 · 0·~1 min·Asked atAnduril Robinhood Together Ai·Relevant atCoreweave Databricks Lambda Labs Runway

Attempt it

In a DeepSpeed JSON config the ZeRO stage is selected by the key . , where the value is one of 1, 2, or 3. Higher stages shard more training state across GPUs at the cost of more -GPU communication.

**DeepSpeed** is a library that adds memory-saving training tricks on top of PyTorch. The most used trick is called **ZeRO**, which shards the training state across multiple GPUs so each GPU holds only a slice. The knob that controls how aggressively ZeRO shards lives in the DeepSpeed JSON config under the key `zero_optimization.stage`. You set its value to 1, 2, or 3 depending on how much memory you need to save. Stage 1 shards the optimizer state across the GPUs in your data-parallel group (GPUs that each train a model copy). This is the biggest single saving for Adam, which keeps two fp32 moments per parameter. Stage 2 also shards the gradients. Each GPU holds only the gradient slices that match its optimizer state slice. Stage 3 also shards the parameters themselves. The model weights are split across GPUs, and the missing slices are pulled in on demand during the forward pass. Each step up saves more memory per GPU. The cost, mostly at stage 3, is more communication between GPUs, because sharded pieces must be gathered when the math needs them. On fast NVLink interconnects this is cheap. On slower PCIe or network links it can bottleneck training.

**The selector key is `zero_optimization.stage`** with values 1, 2, or 3. The DeepSpeed JSON config nests it as `{"zero_optimization": {"stage": 2, ...}}`. Each stage shards more of the training state across the data-parallel ranks at the cost of more cross-GPU communication. Stage 1 shards optimizer state only. For mixed-precision Adam this is the dominant memory term: 12 bytes per parameter for the fp32 master weights and the two fp32 moments. Each rank owns its slice of the moments and applies its slice of the optimizer step. Stage 2 adds gradient sharding. Each rank holds only the gradients that match its optimizer slice, reducing peak training memory further. The reduce-scatter operation that produces these sharded gradients overlaps with the backward pass. Stage 3 adds parameter sharding. The model weights themselves are split, and a parameter is gathered just before any layer that needs it, then dropped after. This is what makes 70B-class training viable on 80 GB GPUs, spread across enough of them. The practical pick depends on memory pressure and interconnect speed. Stage 2 is the common default for mid-size fine-tuning where it fits. Stage 3 unlocks the largest models but adds the most communication overhead, so it benefits from NVLink or NVSwitch fabrics. On slower interconnects, stage 3 can be slower than stage 2 even when memory allows it.

TL;DR

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.

Memory aid

Easy to grasp

Imagine eight friends carrying a giant tent to a campsite. In the smallest version of the trick, each friend carries an equal share of the tent poles. 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.

Key concepts

DeepSpeed ZeRO sharding Data parallel

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 when the math needs it.

The three stages and what each shards

Stage 1 shards optimizer state across the data-parallel ranks. For mixed-precision Adam this is the heaviest single memory term: fp32 master weights plus two fp32 moments, 12 bytes per trainable parameter. With N ranks each holds 12/N bytes per parameter of optimizer state, 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 also shards the gradients. Each rank holds only the gradient slices that match its optimizer slice. A reduce-scatter produces these sharded gradients, and an all-gather then shares the updated parameters. Together they move the same volume as one all-reduce, so communication matches stage 1 and plain data parallel. Memory savings stack on top of stage 1.

Stage 3 also shards the parameters themselves. The model weights are split across ranks. Any layer that needs full weights gathers them via an all-gather just before its forward computation. It 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 mixed-precision Adam (about 16 bytes per parameter: 2 for bf16 weights, 2 for gradients, 12 for fp32 master weights and moments), full data-parallel needs roughly 208 GB per rank before activations. On 8 GPUs, stage 1 cuts that to about 72 GB, stage 2 to about 49 GB, and stage 3 to about 26 GB per rank. That is the difference between fitting on H100 with room for activations and not fitting at all.

The communication cost and where each stage breaks even

Offload: CPU and NVMe as overflow tiers

Composition with other parallelism strategies

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 multi-GPU fine-tuning.
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. A data-parallel all-reduce is a reduce-scatter plus an all-gather. Stage 2 reduce-scatters the gradients, so each rank keeps only its slice, then all-gathers the updated parameters: the same total volume.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

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.stage and its three integer values
What gets sharded at stage 1 versus stage 2 versus stage 3

Primary sources. Browse if you want the original framing.