ZeRO — the Zero Redundancy Optimizer, introduced with DeepSpeed — is the answer to a wasteful fact about ordinary data-parallel training: every GPU stores its own complete copy of the model’s parameters, gradients, and optimizer states, even though the copies are identical after each step. On N GPUs you pay for the same memory N times. ZeRO’s insight is that data parallelism does not actually need that redundancy: it can partition those training states across the GPUs, so each holds only a 1/N slice, and gather what it momentarily needs over the network. This piece is the hub for the ZeRO family. It builds the 2 + 2 + K byte-per-parameter memory model, walks the three partitioning stages (P_os, P_g, P_p), derives the per-stage memory formula, works a concrete example, weighs the communication cost, and places ZeRO against model parallelism. The three stages each get their own deep-dive; here we map the whole territory.
The redundancy hiding in data parallelism
Standard data parallelism is the simplest way to train on many GPUs: replicate the entire model on every device, give each a different slice of the batch, run forward and backward independently, then all-reduce the gradients so every replica applies the same update. It is easy to reason about and scales throughput almost linearly. But it is brutally memory-inefficient. Each of the N GPUs holds a full copy of the parameters, a full copy of the gradients, and a full copy of the optimizer states — and after the synchronized update those copies are bit-for-bit identical.
That is pure redundancy. The model-state memory does not shrink as you add GPUs; it is duplicated N times over. So data parallelism gives you more compute but not more capacity: a model that does not fit on one GPU still does not fit on a hundred. ZeRO attacks exactly this. It keeps the data-parallel execution model — each GPU still processes its own micro-batch through the whole network — but stops storing N identical copies of the training state. Instead it splits that state across the GPUs and reconstructs pieces on demand.
The memory model: 2 + 2 + K bytes per parameter
To reason about the savings you need a memory model. Modern large-model training uses mixed precision with Adam, and for a model of Ψ parameters the per-parameter memory splits into three buckets:
fp16 parameters : 2Ψ bytes
fp16 gradients : 2Ψ bytes
optimizer states (K) : KΨ bytes
-------------------------------------
total : (2 + 2 + K) · Ψ bytesThe first 2Ψ is the working copy of the weights in fp16; the second 2Ψ is the fp16 gradients. The interesting term is K, the optimizer state per parameter. Mixed-precision Adam keeps an fp32 master copy of the weights (4 bytes) plus the fp32 first moment / momentum (4 bytes) and second moment / variance (4 bytes) — so K = 4 + 4 + 4 = 12. Those fp32 states, not the fp16 weights, are the memory hogs: they are 12Ψ of the 16Ψ total. That imbalance is precisely why ZeRO partitions the optimizer states first.
The vanilla baseline: 16Ψ on every GPU
Plug K = 12 into the model and vanilla data parallelism costs (2 + 2 + 12)Ψ = 16Ψ bytes of model-state memory on each GPU, independent of how many GPUs you have. For a 1.5-billion-parameter model that is 16 × 1.5e9 = 24 GB — already tight on a 32 GB card once activations are added. Scale to 7.5 billion parameters and it is 16 × 7.5e9 = 120 GB, which fits on no single accelerator made.
Adding GPUs does not help, because every one independently needs the full 120 GB. This is the wall ZeRO is built to break: the 16Ψ is fixed, dominated by the 12Ψ optimizer term, and completely redundant across the data-parallel group. Divide that 16Ψ by the number of GPUs instead of replicating it, and the hardware that could not hold the model at all holds it comfortably. That division is what the three ZeRO stages deliver, one bucket at a time.
Stage 1 — P_os: partition the optimizer states
ZeRO Stage 1, written P_os (partition optimizer states), targets the fattest bucket. The 12Ψ bytes of fp32 master weights, momentum, and variance are split into N equal shards; GPU i owns only shard i. Parameters and gradients are still replicated in full, but each GPU is now responsible for updating only its 1/N slice of the weights.
The per-GPU memory becomes 2Ψ + 2Ψ + 12Ψ/N. As N grows the optimizer term collapses toward zero and the cost approaches 4Ψ — a 4× reduction from the 16Ψ baseline. Crucially, Stage 1 costs nothing in extra communication: a standard data-parallel step already moves 2Ψ of gradient/parameter traffic, and by reorganizing that same traffic as a reduce-scatter of gradients plus an all-gather of the updated weights, P_os keeps total communication identical. It is the cheapest, safest first move — explored fully in the Stage 1 deep-dive (tm_zero_1_math).
Stage 2 — P_g: partition the gradients too
Stage 2, P_os+g, adds gradient partitioning (P_g) on top of Stage 1. The observation is that once GPU i only updates its 1/N shard of parameters, it only ever needs the gradients for that shard. So there is no reason to keep the full 2Ψ gradient buffer on every device. During the backward pass, as each layer’s gradients are produced they are reduce-scattered to the owning GPU and the rest are discarded.
Now per-GPU memory is 2Ψ + (2 + 12)Ψ/N = 2Ψ + 14Ψ/N, which approaches 2Ψ for large N — an 8× reduction. Like Stage 1, P_g preserves the baseline communication volume of 2Ψ, because the reduce-scatter it uses moves no more data than the all-reduce it replaces. This is the reason Stage 2 is the popular default: it roughly doubles the savings of Stage 1 for free in bandwidth terms. The full treatment lives in the Stage 2 deep-dive (tm_zero_2_math).
Stage 3 — P_p: partition the parameters too
Stage 3, P_os+g+p, takes the final step and partitions the parameters themselves (P_p). Now nothing is fully replicated: each GPU permanently stores only its 1/N slice of the fp16 weights, its slice of the gradients, and its slice of the optimizer states. When a layer is needed for the forward or backward pass, its parameters are all-gathered from the owners just in time, used, and then freed again.
Per-GPU model-state memory drops to the clean (2 + 2 + 12)Ψ/N = 16Ψ/N — a reduction that is linear in the number of GPUs. With enough devices the model state per GPU becomes arbitrarily small, which is what makes trillion-parameter training feasible. The catch is communication: because parameters must be gathered twice (once in forward, once in backward) on top of the gradient reduce-scatter, Stage 3 raises total volume to about 3Ψ, roughly 1.5× the baseline. This stage is functionally equivalent to PyTorch FSDP; see the deep-dive (tm_zero_3_math).
The per-stage memory formula and a worked example
Collecting the three stages into one table, with N the data-parallel degree and K = 12:
baseline (DP) : 2Ψ + 2Ψ + KΨ = 16Ψ
Stage 1 (P_os) : 2Ψ + 2Ψ + KΨ/N → 4Ψ (~4x)
Stage 2 (P_g) : 2Ψ + (2 + K)Ψ/N → 2Ψ (~8x)
Stage 3 (P_p) : (2 + 2 + K)Ψ/N = 16Ψ/N → 0 (Nx)Now make it concrete with a 7.5B-parameter model on N = 64 GPUs, so Ψ = 7.5e9:
baseline : 16 × 7.5 = 120 GB / GPU (fits nowhere)
Stage 1 : 4×7.5 + 12×7.5/64 = 31.4 GB / GPU
Stage 2 : 2×7.5 + 14×7.5/64 = 16.6 GB / GPU
Stage 3 : 16 × 7.5 / 64 = 1.9 GB / GPUThe progression is stark: a model that fit on no GPU at 120 GB drops to 1.9 GB per device under Stage 3 — a 64× cut matching the GPU count exactly. Even Stage 1 alone turns the impossible into a comfortable fit.
The communication trade
Memory is not free; the currency you spend is network traffic. The elegant part of ZeRO is how little the first two stages cost. A baseline data-parallel step moves 2Ψ of data per step (the gradient all-reduce, counted as a reduce-scatter plus an all-gather). Stages 1 and 2 restructure that same 2Ψ without adding to it — you get the 4× and 8× memory wins at zero extra communication.
Stage 3 is where you finally pay. Since no GPU holds the full parameter set, the layers must be all-gathered during the forward pass and again during the backward pass, adding roughly 2Ψ of parameter traffic on top of the Ψ gradient reduce-scatter — about 3Ψ total, or 1.5× the baseline. That extra 50% is usually a bargain for the linear memory reduction it buys, but it means Stage 3 rewards fast interconnects and computation overlap. The practical rule: reach for the lowest stage that makes your model fit, because each step up trades bandwidth headroom for capacity.
Where ZeRO sits versus model parallelism
ZeRO is not the only way to train a model too big for one GPU; the classic alternative is model parallelism — tensor parallelism (splitting individual matrix multiplies across GPUs) and pipeline parallelism (splitting layers into stages). These genuinely divide the model, but they demand intrusive code changes, communicate inside every layer at high frequency, and scale cleanly only within a tightly-coupled node.
ZeRO’s appeal is that it achieves comparable memory savings while keeping the data-parallel programming model intact. Each GPU still runs the entire network on its own micro-batch; ZeRO only changes where the training state lives, not how the computation is expressed — so it drops into existing models with almost no rewriting and scales across nodes. The two approaches are also complementary: large-scale runs routinely combine ZeRO with tensor and pipeline parallelism (so-called 3D parallelism) to push past what either achieves alone — ZeRO making data parallelism memory-efficient, model parallelism splitting the compute graph, different axes of the same scaling problem.
The stages as a family: which one to reach for
ZeRO is best understood as a dial, not a switch. Stage 1 (P_os) shards the optimizer states for a 4× cut at no communication cost — the safe default when the fp32 states are your only problem. Stage 2 (P_g) adds gradient sharding for 8×, still bandwidth-neutral, and is the sweet spot for most workloads. Stage 3 (P_p) shards parameters too for a full N× linear reduction, at a 1.5× communication premium — the tool for the very largest models.
This overview is deliberately a map, not the terrain. Each stage has its own mechanics — which collectives fire and when, how reduce-scatter and all-gather interleave with compute, and the tuning knobs — covered in the dedicated deep-dives for Stage 1, Stage 2, and Stage 3. Start here for the 2 + 2 + K model; follow the pointers for a specific stage.