The fp16 weights of a 7.5B-parameter model fit in 15 GB. Training that model with data parallelism puts 120 GB of model state on every GPU — before counting a single activation. The gap isn't waste or a framework bug. It's arithmetic that the ZeRO paper wrote down in one line, and that line explains why every modern trainer (DeepSpeed ZeRO, PyTorch FSDP) is built around partitioning state you'd never see in the model checkpoint.
Where 16 bytes per parameter comes from
Mixed-precision training with Adam keeps, for every single parameter Ψ:
| State | Precision | Bytes/param |
|---|---|---|
| Parameter | fp16 | 2 |
| Gradient | fp16 | 2 |
| Master copy of parameter | fp32 | 4 |
| Adam momentum | fp32 | 4 |
| Adam variance | fp32 | 4 |
The ZeRO paper writes this as 2Ψ + 2Ψ + KΨ = 16Ψ bytes, with K = 12 for mixed-precision Adam. The fp16 weights you think of as "the model" are 12.5% of the model state. The optimizer — the thing that vanishes the moment training ends — is 75% of it.
Concretely: GPT-2 1.5B has fp16 weights of about 3 GB, but needs at least 24 GB of model state to train. That's why it wouldn't fit on a 32 GB V100 with plain data parallelism, even though inference needed a tenth of the card.
Why you can't just delete the fp32 copies
The obvious question: fp16 parameters already exist, so why keep a 4-byte fp32 copy of each one? Because the weight update is usually tiny relative to the weight, and fp16 can't absorb it.
The mixed-precision training paper gives two failure modes:
- Underflow. Any value with magnitude below 2⁻²⁴ becomes zero in fp16. Gradients multiplied by a small learning rate routinely land there.
- Swamping. Even when the update is representable, adding it to a much larger weight loses it: when the weight is at least 2048× larger than the update, the addition right-shifts the update's bits past fp16's 10-bit mantissa and the update rounds away entirely.
w + Δw == w, every step, silently.
This isn't theoretical. On their Mandarin speech model, updating the fp16 weights directly cost 80% relative accuracy versus the fp32 baseline; accumulating updates into an fp32 master copy matched fp32 training exactly. The same paper is why momentum and variance stay in fp32 too, and why tensor cores accumulate fp16 products into fp32 — some networks failed to match baseline accuracy when dot products accumulated in fp16.
So the 12 bytes of fp32 state per parameter are load-bearing. You can't shrink K. What ZeRO noticed is that you don't have to: you can stop replicating it.
The replication nobody questioned
Classic data parallelism gives every GPU a full replica of all 16Ψ bytes. Each rank computes gradients on its shard of the batch, the ranks all-reduce the gradients, and then every rank performs the identical optimizer update on identical state. With 64 GPUs, the momentum tensor exists in 64 identical copies, and 63 of them are pure redundancy.
ZeRO (Zero Redundancy Optimizer) partitions the three classes of state across the data-parallel group instead of replicating them. Each stage partitions one more class. For the paper's running example — Ψ = 7.5B, data-parallel degree Nd = 64:
| Stage | What's partitioned | Bytes/param | Memory per GPU |
|---|---|---|---|
| Baseline DP | nothing | 16Ψ | 120 GB |
| Stage 1 (Pos) | optimizer states | 4Ψ + 12Ψ/Nd | 31.4 GB |
| Stage 2 (Pos+g) | + gradients | 2Ψ + 14Ψ/Nd | 16.6 GB |
| Stage 3 (Pos+g+p) | + parameters | 16Ψ/Nd | 1.88 GB |
Stage 1 alone is roughly a 4× cut at large Nd, because the optimizer states were three-quarters of the bill. Stage 2 gets to 8×. Stage 3 scales linearly with the number of GPUs: double the cluster, halve the per-GPU footprint.
The part that sounds too good: communication
Partitioning state that the forward and backward passes need sounds like it should drown you in network traffic. The paper's accounting says otherwise.
Baseline data parallelism all-reduces gradients each step. A ring all-reduce is a reduce-scatter followed by an all-gather, so it moves 2Ψ data per step. Against that baseline:
- Stage 1 and Stage 2 move exactly 2Ψ — no increase at all. Instead of all-reducing, each rank reduce-scatters gradients (Ψ) so every rank ends up owning the final gradients for its partition, updates only its slice of the parameters, and the ranks all-gather the updated fp16 parameters (Ψ). Same total volume, 8× less memory.
- Stage 3 moves 3Ψ — a 1.5× increase. Parameters now have to be all-gathered on demand during the forward and backward passes (then discarded), plus the gradient reduce-scatter. That 50% overhead buys memory that scales with cluster size.
This is the sharpest tradeoff in the design: stages 1–2 are nearly free and should be your default; stage 3 is a deliberate purchase of capacity with bandwidth.
The same math, shipped as FSDP
PyTorch's FSDP is this exact scheme under different names, and the FSDP docs map cleanly onto the stages:
NO_SHARD: states "not sharded but instead replicated across ranks," gradients synced via all-reduce — baseline DP.SHARD_GRAD_OP: "gradients and optimizer states are sharded during computation" — ZeRO stage 2.FULL_SHARD: "parameters, gradients, and optimizer states are sharded"; FSDP "unshards (via all-gather) before the forward, reshards after the forward," and after backward "synchronizes and shards them (via reduce-scatter)" — ZeRO stage 3, including the 3Ψ communication pattern.HYBRID_SHARD: full sharding within a node, replication across nodes — keeping the expensive all-gathers on NVLink instead of the inter-node fabric.
What the 16Ψ formula doesn't cover
Model state is only one tenant of GPU memory, and the ZeRO paper is explicit about the rest:
- Activations dominate at scale: GPT-2 1.5B at sequence length 1K and batch size 32 wants about 60 GB of activations. Activation checkpointing cuts that to about 8 GB for roughly 33% recompute overhead — and a 100B model still needs ~60 GB even with checkpointing.
- Temporary buffers: a flattened fp32 all-reduce buffer for a 1.5B model is 6 GB on its own if sized to the full model.
- Fragmentation: the paper reports OOM in extreme cases with over 30% of memory still free, because short-lived and long-lived allocations interleave until no contiguous block fits a request.
So when you're estimating whether a model fits, start from the one-liner — 16 bytes × parameters, divided by whatever your sharding stage divides — then budget activations separately, because no amount of optimizer sharding touches them.
The durable lesson: the checkpoint you download is the smallest object in the training job. Everything expensive about training a model is state that never leaves the cluster — and ZeRO's insight was simply that identical copies of it on every GPU were a convention, not a requirement.