ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Every GPU stores a full copy of the optimizer. Partition it across the GPUs instead.

Standard data parallelism copies the entire model state onto every GPU, so most of a cluster's memory holds duplicates. ZeRO splits that state into shards, gathers each piece only when a GPU needs it, and shrinks per-GPU memory by up to the number of GPUs, at almost no extra communication.

A 1.5-billion-parameter model is three gigabytes of weights. Training it takes twenty-four. The other twenty-one are mostly the optimizer's fp32 bookkeeping, and standard data parallelism copies all of it onto every GPU you own.

For most of the last decade the thing that capped the size of a neural network was not how fast your GPUs could multiply matrices and not how much data you had. It was memory. A modern accelerator has a fixed amount of it, tens of gigabytes, and if the numbers a training step needs to keep alive do not fit, the step does not run. You can wait for more compute and gather more data, but the memory on the card is fixed.

The wall is closer than it looks, because training a model costs far more memory than storing it. Take GPT-2 at 1.5 billion parameters. Its weights in 16-bit precision are about 3 GB. Yet you cannot train it on a 32 GB GPU with a standard setup, because the weights are the small part. Training also has to hold the gradients, and, if you use a good optimizer, a set of running statistics for every single parameter, all of it in higher precision than the weights themselves. Add it up and a 1.5B model needs about 24 GB just for these model states before a single activation is stored. Most of the memory a training step uses is not the model.

The usual way to train on many GPUs makes this worse in one exact way. Data parallelism, the workhorse of distributed training, puts a complete replica of the model on every GPU and feeds each a different slice of the batch. Sixteen GPUs means sixteen identical 24 GB copies. Your cluster has plenty of aggregate memory, but you cannot use it to fit a bigger model, because every GPU is holding the same thing. Adding GPUs buys you throughput and never buys you room.

ZeRO, from the team behind Microsoft's DeepSpeed, is a plain fix for that specific waste. The name stands for Zero Redundancy Optimizer, and the idea is in the name: stop storing identical copies. Partition the model states so each GPU keeps only its own $1/N_d$ slice (one shard per GPU, with $N_d$ the number of data-parallel GPUs), and fetch the rest from whoever owns it, exactly when a step needs it and not before. The computation stays ordinary data parallelism, every GPU still runs the whole model on its own batch, but the memory bill drops in proportion to how many GPUs you have. A trillion-parameter model that would need 16 terabytes of model states fits, in principle, across 1024 GPUs at 16 GB each.

A few ideas explain the method: where a training step's memory actually goes, why data parallelism duplicates all of it, how to partition it in three stages, and why partitioning barely raises the communication bill. None is difficult on its own. Together they turn a hard memory limit into one that shrinks as you add GPUs.

The wall is memory, not compute

Before any fix, it helps to see the bill itemized, because the fix turns on the exact shape of that bill. During training, memory splits into two piles. The first, which the paper calls model states, is the optimizer states, the gradients, and the parameters, one entry per weight in the model. The second, the residual states, is everything else: the activations saved during the forward pass, temporary buffers, and memory lost to fragmentation. For large models the first pile dominates, so that is where ZeRO starts. The next section counts those bytes exactly.

Where the sixteen bytes go

The state of the art for training large models on NVIDIA hardware is mixed-precision training, and the memory story starts there. The point of mixed precision is speed: the big matrix multiplies run in 16-bit ($\text{fp16}$), which is where the tensor cores are fast and where activations and gradients take half the memory. But 16-bit is too coarse to accumulate the tiny updates an optimizer makes. Add a small step to a large weight in fp16 and the step can round away to nothing, so training stalls. The standard remedy keeps a second, full-precision ($\text{fp32}$) master copy of every weight. The fast forward and backward passes use a freshly rounded fp16 copy, and the optimizer applies its update to the fp32 master, where a small change is never rounded away.

So the weights are stored twice, once in fp16 for compute and once in fp32 for accurate accumulation. Now add the optimizer. The dominant choice, Adam, keeps two running statistics for every parameter: a momentum term, the decaying average of past gradients, and a second term, the decaying average of past squared gradients, which scales each weight's step by how noisy its gradient has been. Both are kept in fp32 for precision. Count the bytes per parameter for a model with $\Psi$ weights:

Two bytes for the fp16 parameters, two for the fp16 gradients, and twelve for the optimizer's fp32 records: the master weight copy, the momentum, and the second-moment term. The paper writes the optimizer-state total as$K\Psi$ and calls $K$ the memory multiplier, with $K = 12$ for mixed-precision Adam. Everything lands on sixteen bytes per parameter.

One grouping in (1) trips people up. The fp32 master weight copy sits inside the twelve, alongside momentum and variance. It is part of the optimizer's bookkeeping, not a fourth line item on top of the sixteen. The total is exactly $16\Psi$, never $20\Psi$. And a related count: Adam itself keeps only two states per weight, momentum and the second moment. That third four-byte block in$K$ is the fp32 master copy, which is an artifact of mixed precision, not something Adam stores. Get this wrong and every number downstream, the memory ladder and the trillion-parameter projection, drifts.

Now the sixteen bytes have a clear message. For a 1.5B model, $16 \times 1.5 = 24$ GB, of which the fp16 weights are just 3 GB. The remaining 21 GB, seven-eighths of the bill, is the optimizer's fp32 records and the gradients. The heaviest single chunk, the twelve bytes of optimizer states, is the fp32 material that never even touches the fast forward and backward passes. ZeRO exploits that imbalance: the biggest, most-duplicated part of the memory is the part a GPU touches least often during a step.

The multiplier $K$ depends on the optimizer, and it sets the size of the first win. Adam's two extra fp32 states give $K = 12$. Plain SGD with momentum keeps one velocity term, so$K = 8$; vanilla SGD keeps none, so $K = 4$, just the master copy. The richer the optimizer, the more fp32 bookkeeping it drags along, and the more there is to reclaim by not duplicating it.

Data parallelism keeps a copy on every GPU

To see the redundancy plainly, look at what data parallelism actually does. Every GPU holds a full replica of all$16\Psi$ bytes. Each processes a different micro-batch through that replica, runs a full forward and backward, and at the end the GPUs average their gradients so they all take the same optimizer step and stay in sync. The computation is efficient and coarse-grained, which is why data parallelism scales well. Its memory does not scale at all: the same $16\Psi$ bytes are stored $N_d$ times over. The paper calls this the redundancy: a cluster of 64 GPUs holds 64 identical copies of a model's state when a single copy would serve. The payoff of removing it is direct: the memory those duplicates occupy becomes free room for a larger model or a larger batch.

Toggle the figure below from DP to ZeRO and drag the GPU count. Under DP, per-GPU memory stays pinned at 120 GB no matter how many GPUs you add, because each holds a full copy. Under ZeRO, the same colored model state is cut into disjoint slices, one per GPU, and per-GPU memory falls as $120/N_d$. Watch the redundancy counter go from$\times N_d$ to $\times 1$:

Here ZeRO is easy to confuse with the other way people fit big models: model parallelism, the approach in Megatron-LM. Model parallelism splits the computation itself, chopping each layer's matrix multiplies across GPUs so each does a piece of every layer. That works, but it makes GPUs talk constantly, several all-reduces per layer in both directions, and the talk is only cheap inside a single node where the interconnect is fast (NVSwitch at about 300 GB/s). Cross a node boundary onto ordinary InfiniBand at about 12.5 GB/s and efficiency collapses. The paper measured a 40B model split with Megatron across two nodes running at about 5 TFlops per V100, under 5% of the hardware's peak.

ZeRO does something different, and that difference is misread more often than anything else about the paper. ZeRO does not split the computation. Every GPU still runs the entire forward and backward on its own batch, exactly as in data parallelism. Only the stored state is partitioned, and the missing pieces are gathered just before they are needed and released right after. No layer's matmul is ever divided. So ZeRO is data parallelism with the memory redundancy removed, not model parallelism, and it composes with model parallelism rather than replacing it: run both and the per-GPU reduction multiplies to $N_d \cdot N_m$ for data-parallel degree $N_d$ and model-parallel degree$N_m$.

Partition, do not replicate

The model states come in three kinds, and they are not equally expensive, so ZeRO partitions them in three stages, biggest chunk first. Each stage is cumulative, and each cuts more of the per-GPU bill. Recall the split from (1): the optimizer states are twelve of the sixteen bytes, the gradients two, the fp16 parameters two.

Stage one, $P_{os}$: partition the optimizer states. Group the twelve bytes of fp32 records into $N_d$ equal shards and let GPU $i$ own only shard $i$. Each GPU updates just its $1/N_d$ of the parameters, then the GPUs all-gather the updated weights so everyone has the full model again for the next forward pass. The optimizer states are the largest block, and, used only at the update, the least often touched, so partitioning them first reclaims the most memory for the least trouble. Per-GPU memory drops from$16\Psi$ to:

The four bytes that stay resident are the fp16 parameters and gradients; the twelve fp32 bytes shrink by$N_d$. As $N_d$ grows the second term vanishes and the memory approaches $4\Psi$, a$4\times$ reduction from $16\Psi$.

Stage two, $P_{os+g}$: also partition the gradients. Once GPU $i$ only updates its own parameter shard, it only needs the gradients for that shard. So instead of every GPU keeping every gradient, each gradient is reduced onto just the GPU responsible for it, and the rest can be dropped as soon as they are summed. Per-GPU memory becomes:

Now only the fp16 parameters ($2\Psi$) stay fully resident, and everything else, the two gradient bytes plus the twelve optimizer bytes, shrinks by $N_d$. Large $N_d$ takes this toward $2\Psi$, an$8\times$ reduction.

Stage three, $P_{os+g+p}$: also partition the parameters. Each GPU now stores only its$1/N_d$ of the weights too. When a layer's weights are needed for the forward or backward pass, the GPUs all-gather that layer from its owners, use it, and throw it away. Nothing is left fully replicated, and the per-GPU memory becomes simply:

Stage three alone has no floor. Every stage before it left a resident remainder ($4\Psi$, then$2\Psi$) that capped how far the memory could fall. Stage three removes even that, so the per-GPU memory is just the model divided by the GPU count. Double the GPUs and you halve the memory, without limit. That is what makes a trillion parameters conceivable: enough GPUs, and each one's share shrinks to whatever it needs to be.

Step the figure through the three stages and drag the GPU count. Watch the optimizer states (the three fp32 blocks) collapse first, then the gradients, then the parameters, and read the live GB off the total. At $N_d = 64$it reproduces Table 1 of the paper exactly:

The concrete ladder for the paper's running example, a 7.5B model on 64 GPUs, tells the story in four numbers. Standard DP: 120 GB per GPU. Partition the optimizer states and it falls to 31.4 GB. Add the gradients and it is 16.6 GB. Add the parameters and it is 1.9 GB, small enough to leave most of a 32 GB card free for activations and a healthy batch. From 120 GB to under 2, on the same hardware, by refusing to store anything twice.

The headline factors deserve one qualifier. The paper advertises $4\times$ and $8\times$ for the first two stages, and those are the limits as $N_d \to \infty$, when the resident $4\Psi$ and$2\Psi$ dominate the shrinking term. At a finite $N_d = 64$ the realized reductions are a little under: $120/31.4 \approx 3.8\times$ and $120/16.6 \approx 7.2\times$. Only stage three is exact, because it has no residual term at all: at $N_d = 64$ it is precisely $64\times$. The rounding is small, but if you check the table with a calculator, that is why it does not land on a clean 4 and 8.

Why the partition is almost free

The obvious worry with all this fetching is communication. If GPUs now have to ship each other pieces of the model, surely we have traded a memory problem for a bandwidth problem. The surprising result, and the reason ZeRO works, is that stages one and two add no communication at all, and stage three adds only fifty percent. To see why, you first need the yardstick: how much does plain data parallelism already communicate?

Data parallelism averages gradients with an all-reduce, and a bandwidth-optimal all-reduce is not a single magic operation. It is two steps. First a reduce-scatter: the GPUs sum their gradients but each keeps only its own $1/N_d$ slice of the summed result, so each GPU sends and receives about $\Psi$ elements. Then an all-gather: each GPU shares its reduced slice with everyone until all hold the full averaged gradient, another $\Psi$ elements moved. Reduce-scatter plus all-gather is an all-reduce, and the total traffic per GPU is:

Two details keep the number exact. The $2$ comes from the two-phase structure, reduce-scatter then all-gather, and has nothing to do with fp16 being two bytes; $\Psi$ here counts elements, not bytes. And the exact per-GPU volume is $2\,\tfrac{N_d-1}{N_d}\,\Psi$, a hair under $2\Psi$, which the paper rounds up. What matters is that $2\Psi$ is the fixed cost data parallelism was already paying every step, and ZeRO has to match it.

Now watch stage two, $P_{os+g}$, clear it exactly. The gradients no longer need a full all-reduce, because each GPU only updates its own parameter shard and so only needs the summed gradients for that shard. A single reduce-scatter delivers that: $\Psi$ moved. Each GPU updates its shard, and then one all-gather shares the updated parameters with everyone for the next forward pass: another $\Psi$. Total: $2\Psi$, the same as (5). ZeRO reaches an $8\times$ memory reduction while moving not one byte more than ordinary data parallelism already did.

# one ZeRO-DP step, stage P_os+g: gradients and optimizer states sharded.
# Nd GPUs; GPU i owns shard i of the params, grads, and optimizer state.
for micro_batch in local_data:          # each GPU: its own slice of the batch
    loss = model(micro_batch)           # FULL forward, replicated on every GPU
    loss.backward()                     # FULL backward, replicated on every GPU
reduce_scatter(grads)                   # move Psi: GPU i is left with grad shard i
optimizer.step(param[i], grad[i])       # update only your 1/Nd of the parameters
all_gather(param)                       # move Psi: every GPU gets the full params
# total moved: Psi + Psi = 2*Psi, exactly standard data parallelism.

The word to be careful with is "free." Stages one and two are free in the sense that they add nothing over the data-parallel baseline, both cost $2\Psi$, not in the sense that they communicate nothing. The baseline was never zero. What is remarkable is that the memory win, eight-fold, comes with no communication penalty at all over the setup you were already running.

Stage three is the only one that costs more, and it is easy to see where the extra comes from. Now that the parameters are also partitioned, they have to be gathered when they are needed, and they are needed twice: once on the way down in the forward pass and once on the way up in the backward pass. So the parameters get all-gathered in the forward ($\Psi$), all-gathered again in the backward ($\Psi$), and the gradients still get reduce-scattered ($\Psi$):

A fifty percent increase in communication in exchange for cutting memory by the full factor of $N_d$, without a floor. Set that trade against what it unlocks, a model bounded only by the number of GPUs you can gather: fifty percent more traffic buys a memory ceiling that rises with every GPU you add.

The stage-three gather also answers an obvious objection: if every GPU ends up with the full parameters in order to compute, how is that any less memory-hungry than just replicating them? Because the parameters are never all present at once. They are gathered one layer at a time, that layer computes, and its weights are discarded before the next layer is gathered. Peak live parameter memory is one layer's worth, not the whole model's. Press Play below and watch the sweep:

The communication ledger below lays the three schemes side by side. The dashed line is the data-parallel baseline of$2\Psi$. Stage two lands exactly on it; stage three overshoots by half:

There is a quieter virtue in all of this, and it explains why ZeRO spread so widely and so fast: none of it changes the training math. Every GPU computes exactly the gradients standard data parallelism would, takes exactly the optimizer step it would, and lands on exactly the same weights. ZeRO only rearranges where the bytes are stored and when they move. A model trained with ZeRO is bit-for-bit the model you would have trained without it, only larger, which is a stronger guarantee than pipeline schemes that keep stale weights or memory-thrifty optimizers that coarsen their statistics and can shift where training converges.

The memory that is left over

Partitioning the model states can uncover a second bottleneck. Once the optimizer, gradients, and parameters are no longer the dominant cost, the residual states from earlier, the activations, temporary buffers, and fragmentation, can become the thing that overflows. The paper's second half, ZeRO-R, is a set of three smaller optimizations aimed at these, and the first one reuses the same idea as ZeRO-DP.

Start with activations. The standard tool for shrinking them is activation checkpointing, which stores only a few activations during the forward pass and recomputes the rest during the backward pass, trading roughly a 30% compute overhead for a large memory saving (it holds about the square root of the activations it would otherwise keep). It is a prior technique, not a ZeRO invention, and it is not enough on its own: a 100B model still needs about 60 GB of activation memory even with checkpointing. Worse, when you combine checkpointing with model parallelism, the surviving checkpoints get replicated across the model-parallel GPUs, the same redundancy problem one level down.

ZeRO-R's answer, partitioned activation checkpointing ($P_a$), applies the same partition-don't-replicate idea to activations: split the checkpoints across the model-parallel GPUs instead of replicating them, and all-gather a full copy of a layer's activation only when the backward pass reaches it. Activation memory falls in proportion to the model-parallel degree. For a 100B model with 16-way model parallelism that is a 16-fold cut, from about 33 GB of checkpoints per GPU to about 2 GB, and for the very largest models those partitioned checkpoints can be pushed to CPU memory, taking activation memory to nearly zero.

That last move sounds reckless, offloading to CPU over the slow PCIe link, and yet it is safe. The safety comes down to arithmetic intensity, the ratio of compute done to bytes moved. For a transformer the compute in a layer grows with the square of the hidden dimension while an activation checkpoint grows only linearly with it, so the intensity grows with model width and, for GPT-2-scale models and up, exceeds ten thousand FLOPs per byte moved. When each byte fetched feeds that much arithmetic, there is enough compute time to hide the transfer behind it: the bytes still move, but the GPU is never left waiting on them. The extra traffic from partitioning checkpoints stays under 10% of the model-parallel communication, and CPU offload becomes a way to buy back memory without stalling.

The other two residual optimizations are more mundane and just as necessary at scale. Constant-size buffers ($CB$) cap the fused temporary buffers that libraries build for efficient collectives, which otherwise grow with the model until a single 3B model's buffer wants 12 GB. Memory defragmentation ($MD$) pre-allocates contiguous regions for the long-lived checkpoints and gradients, so the interleaving of short-lived and long-lived tensors does not leave memory so fragmented that an allocation fails with gigabytes nominally free, which the paper saw happen with over 30% of memory still available.

What it actually trains

The paper ships a subset of all this, called ZeRO-100B, which is stage two of ZeRO-DP ($P_{os+g}$) plus ZeRO-R, and measures it on 400 V100 GPUs. The point of stopping at stage two is that stage three's parameter partitioning was left for a later release; even so, the results are large.

On model size, ZeRO-100B trains models up to 170 billion parameters, about eight times the roughly 20B that the state-of-the-art Megatron setup could train with acceptable throughput on the same hardware. On speed, it runs 100B models at over 38 TFlops per GPU, more than 15 PetaFlops in aggregate and above 30% of the hardware peak, roughly a 10-fold speedup over the baseline at that size. And it does both at once, because one move drives both: less per-GPU memory leaves room for a larger batch, and a larger batch means more arithmetic per unit of communication.

That coupling produces the paper's strangest-looking result, super-linear scaling. Across the 64-to-400-GPU range on a 60B model, each doubling of the GPU count more than doubles throughput. Ordinarily doubling GPUs at best doubles throughput; here it does better, because stage two shrinks each GPU's model-state footprint as you add GPUs, which frees memory for bigger per-GPU batches, which raises efficiency. More GPUs make each GPU faster, not just more numerous.

There is a usability result that mattered as much as the raw numbers. Because ZeRO runs as ordinary data parallelism, a practitioner can train a 13B-parameter model, larger than T5's 11B, the biggest published model at the time, without any model or pipeline parallelism and without refactoring their code. Standard PyTorch distributed data parallelism runs out of memory at 1.4B. ZeRO moved the ceiling for an ordinary user by an order of magnitude, with a one-line change. The same system trained Turing-NLG, a 17B-parameter model that set a new record on WebText-103 with a perplexity of 10.21 (perplexity measures how surprised a model is by held-out text, so lower means it predicts real language more confidently).

The title's trillion is a projection, and the paper is careful to separate what fits from what is practical. On memory, the arithmetic is clean: a trillion parameters at sixteen bytes each is 16 terabytes of model states, and spread across 1024 GPUs with stage three that is 15.6 GB per GPU, which fits. The limit is compute, not memory. A trillion-parameter model carries roughly 3000 times the computation per sample of BERT-Large, which itself trains in about an hour on a 1024-GPU cluster. Three thousand hours works out to roughly 140 days at the same efficiency, and likely more than a year in practice once the data and sequence lengths grow too. ZeRO makes the trillion-parameter model fit. Making it train in reasonable time waits on faster hardware.

What ZeRO changed is less a single algorithm than a default. Fitting a model no longer meant carving up its computation and rewriting it for a particular cluster. It meant adding enough GPUs and letting the states shard themselves, with the training math untouched. The idea outlived the paper: PyTorch's Fully Sharded Data Parallel is stage three under another name, and it is how a large fraction of today's models are trained. It all began with a plain observation, that data parallelism was storing the same sixteen bytes on every GPU, and the decision to stop doing that.