ZeRO — how training a huge model stops replicating everything on every GPU

Yesterday we saw that training memory is not just the model weights. A training GPU may also hold gradients, optimizer state, and activations. Activation checkpointing attacks the activation term by trading memory for recomputation.

Today we will attack a different source of waste:

The family of techniques known as ZeRO — Zero Redundancy Optimizer — observes that ordinary data-parallel training contains enormous redundancy.

Its central idea is:

Keep using many GPUs to process different data, but partition model state across them instead of replicating all of it everywhere.

This idea underlies much of modern large-model training.

Start with ordinary data parallelism¶

Suppose we train a model on four GPUs.

In classic data parallelism, every GPU has a complete copy of the model:

GPU 0: full model + batch A
GPU 1: full model + batch B
GPU 2: full model + batch C
GPU 3: full model + batch D

Each GPU independently performs forward and backward on its own microbatch.

After backward, GPU 0 has gradients:

GPU 1 has:

and so on.

To make all replicas perform the same update, the GPUs combine their gradients, conceptually:

An all-reduce commonly performs this synchronization.

Then every GPU applies the same optimizer update and remains an identical model replica.

This is beautifully simple.

It is also memory-hungry.

What is duplicated?¶

Suppose the model has \(P\) parameters.

With Adam-like training, a simplified GPU may need persistent storage for:

parameters
gradients
first-moment optimizer state
second-moment optimizer state

Often there are also higher-precision master weights or other implementation-specific state.

For intuition, imagine the persistent training state costs roughly:

The exact number depends on precision and optimizer implementation; the point is the scaling.

For a:

parameter model, that rough accounting gives:

or about:

And ordinary data parallelism says:

Put essentially that persistent state on every data-parallel worker.

Adding GPUs increases aggregate cluster memory, but replication prevents us from using that memory to fit a larger model.

The apartment-building mental model¶

Imagine four accountants working on the same enormous company ledger.

Classic data parallelism gives each accountant:

a complete copy of the ledger
a complete copy of all corrections
a complete copy of all bookkeeping history

They work on different receipts, but each desk stores the entire accounting system.

ZeRO asks:

What if the accountants still work in parallel, but divide the filing cabinets among themselves?

For example:

Accountant 0 owns files A–F
Accountant 1 owns files G–L
Accountant 2 owns files M–R
Accountant 3 owns files S–Z

When someone temporarily needs information owned by another accountant, they communicate.

We reduce storage duplication by accepting more coordination.

That is the core trade.

ZeRO Stage 1: shard optimizer state¶

The first redundancy to remove is optimizer state.

Suppose there are four GPUs.

Instead of every GPU storing Adam's optimizer state for every parameter:

GPU 0: optimizer state 0–100%
GPU 1: optimizer state 0–100%
GPU 2: optimizer state 0–100%
GPU 3: optimizer state 0–100%

we partition it:

GPU 0: optimizer state  0–25%
GPU 1: optimizer state 25–50%
GPU 2: optimizer state 50–75%
GPU 3: optimizer state 75–100%

Each GPU still has the full model parameters and, conceptually, full gradients.

But optimizer memory per GPU falls roughly with the data-parallel world size \(N\):

This is ZeRO Stage 1.

ZeRO Stage 2: shard gradients too¶

Gradients are also duplicated.

After backward, classic data parallelism ultimately leaves every worker with the synchronized gradient tensor.

But does every worker really need to retain every gradient?

Not if responsibility for the update is partitioned.

ZeRO Stage 2 shards:

optimizer state
+
gradients

across workers.

Conceptually:

GPU 0 owns gradients for parameter shard 0
GPU 1 owns gradients for parameter shard 1
GPU 2 owns gradients for parameter shard 2
GPU 3 owns gradients for parameter shard 3

Now per-GPU persistent memory falls further.

The full parameters are still replicated.

ZeRO Stage 3: shard the parameters themselves¶

This is the big step.

Why should every GPU permanently store every parameter if each GPU owns only a shard?

ZeRO Stage 3 partitions:

optimizer state
gradients
parameters

So with four GPUs:

GPU 0: parameter shard 0
GPU 1: parameter shard 1
GPU 2: parameter shard 2
GPU 3: parameter shard 3

At rest, no GPU needs a permanent full copy of the model.

Ignoring activations and temporary buffers, the persistent model-state memory per GPU can now approach:

This is the crucial scaling result.

If total persistent training state is 1.12 TB and we shard it perfectly across:

GPUs, the idealized average becomes:

That suddenly looks much more plausible on high-memory accelerators.

But we have created a new problem.

How can a GPU compute with weights it does not have?¶

Suppose GPU 0 owns only one quarter of the parameters.

A Transformer layer's matrix multiplication may require a weight matrix that is distributed across all four GPUs.

Before computing that layer, the workers can temporarily gather the needed parameter shards.

Conceptually:

at rest:

GPU 0: [W0]
GPU 1: [W1]
GPU 2: [W2]
GPU 3: [W3]

need full W for computation:

        communication
             ↓

GPU 0: [W0 W1 W2 W3] temporarily

After the layer has been computed, the gathered copies can be released when they are no longer needed.

Then the next layer's parameters are gathered.

So ZeRO-3 turns parameter memory into a streaming problem:

gather layer weights
      ↓
compute layer
      ↓
discard non-owned copies
      ↓
gather next layer weights

The model exists globally across the cluster even though it need not exist permanently on any one GPU.

The communication primitives matter¶

Three collective-communication ideas are useful here.

All-reduce¶

Everyone contributes data and everyone receives the combined result.

Conceptually:

GPU 0 ─┐
GPU 1 ─┼→ SUM → result on every GPU
GPU 2 ─┼
GPU 3 ─┘

Classic data-parallel gradient synchronization commonly uses all-reduce.

Reduce-scatter¶

Everyone contributes, but the reduced result is divided among workers:

combined gradient
      ↓
GPU 0 gets shard 0
GPU 1 gets shard 1
GPU 2 gets shard 2
GPU 3 gets shard 3

This naturally fits sharded gradients.

All-gather¶

Each worker owns a shard, and the shards are collected so workers can temporarily reconstruct a larger tensor:

GPU 0: W0 ─┐
GPU 1: W1 ─┼→ [W0 W1 W2 W3]
GPU 2: W2 ─┼
GPU 3: W3 ─┘

ZeRO-3 relies heavily on this kind of parameter materialization.

The memory savings therefore come with network traffic.

A simplified memory example¶

Imagine persistent state per parameter consists of:

So total persistent state is:

With ordinary data parallelism, per GPU:

With \(N\) workers, idealized ZeRO-1 gives approximately:

ZeRO-2:

ZeRO-3:

For:

the simplified comparison becomes:

These equations omit activations, temporary gathered parameters, fragmentation, communication buffers, and implementation details.

But they reveal the central scaling behavior.

ZeRO does not make the information disappear¶

This point is worth emphasizing.

If a model requires 1 TB of persistent training state, ZeRO does not somehow compress that state into 100 GB.

The cluster still stores roughly the same information.

It changes:

Instead of:

GPU 0: entire 1 TB
GPU 1: entire 1 TB
GPU 2: entire 1 TB
...

we aim for:

GPU 0: shard
GPU 1: shard
GPU 2: shard
...

The aggregate cluster memory was already there.

ZeRO makes that aggregate memory usable.

Why this resembles PagedAttention¶

There is an interesting connection to our inference lessons.

PagedAttention said:

A request's logical KV cache does not need to be physically contiguous.

ZeRO says something analogous:

A model's logical training state does not need to be physically resident on every worker.

Both techniques separate a logical object from its physical placement.

For inference:

For distributed training:

That abstraction is enormously powerful.

The cost is communication¶

Suppose a layer has a large weight tensor.

With ZeRO-3, workers may need to communicate parameter shards before computing the layer.

So we have traded:

for:

This should sound familiar from activation checkpointing.

Yesterday we traded:

for:

Today we trade:

for:

Large-model training is full of these exchanges.

Communication can sometimes be hidden¶

Suppose we are computing layer \(k\).

While the GPU performs arithmetic for that layer, the system may begin gathering parameters for layer:

Conceptually:

time ─────────────────────────────→

compute layer k:
[=====================]

gather weights for k+1:
          [==============]

compute layer k+1:
                         [=====================]

If communication overlaps successfully with computation, some network cost disappears from the critical path.

This is one reason distributed-training performance depends on scheduling, not merely bandwidth.

A fast network is valuable.

A system that knows when to move tensors can be equally important.

Parameter prefetching has a risk¶

Why not gather many future layers far in advance?

Because then their weights occupy GPU memory.

If we prefetch too aggressively:

current layer
next layer
next-next layer
next-next-next layer

we gradually recreate the memory problem ZeRO was designed to solve.

So prefetch distance creates another tradeoff:

The system must keep enough future work ready to hide communication without materializing too much of the model at once.

ZeRO and tensor parallelism solve different problems¶

These techniques are easy to conflate.

Tensor parallelism splits the computation of a layer across devices.

For example, a matrix multiplication might be partitioned:

with:

Different GPUs compute different pieces of the same layer concurrently.

ZeRO is primarily about sharding training state across data-parallel workers so that redundant copies do not consume memory.

A rough distinction is:

tensor parallelism:
How do multiple GPUs jointly compute this layer?

ZeRO:
How do multiple training replicas avoid storing redundant state?

Modern training systems often use both.

And pipeline parallelism solves yet another problem¶

Pipeline parallelism assigns different layers to different devices:

GPU group A:
layers 1–20

GPU group B:
layers 21–40

GPU group C:
layers 41–60

So we now have several axes of distribution:

data parallelism
→ split examples

tensor parallelism
→ split operations inside a layer

pipeline parallelism
→ split layers

ZeRO / sharded data parallelism
→ split persistent training state

Real large-model training may combine all of them.

This is why statements like "the model was trained on 10,000 GPUs" tell us surprisingly little about how the computation was actually arranged.

ZeRO and activation checkpointing compose¶

Yesterday we divided training memory into:

parameters
gradients
optimizer state
activations
temporary buffers

ZeRO attacks:

parameters
gradients
optimizer state

depending on stage.

Activation checkpointing attacks:

activations.

So the techniques are complementary.

A large training job might use:

ZeRO-3
+
activation checkpointing
+
mixed precision
+
FlashAttention
+
tensor parallelism
+
pipeline parallelism

Each technique attacks a different bottleneck.

The final system works because none of them has to solve the whole problem alone.

The deeper idea: replication is a luxury¶

Replication is wonderful when memory is plentiful.

It makes computation simple:

everything I need is local

and local memory is much faster and easier to reason about than remote state.

But once a model becomes large enough, full replication is no longer affordable.

Then the system must ask:

Which information truly needs to exist here right now?

For a particular layer during a particular instant, perhaps the GPU needs a full weight tensor.

But that does not imply it must keep that tensor permanently.

This is exactly the same reasoning we used for activations:

That distinction is one of the foundations of scalable AI systems.

The key mental model¶

Imagine the entire training state as a giant library.

Classic data parallelism builds an identical complete library beside every researcher:

Researcher 0 → full library
Researcher 1 → full library
Researcher 2 → full library
Researcher 3 → full library

This is wonderfully convenient and spectacularly wasteful.

ZeRO builds one distributed library:

Researcher 0 → shelves 0–24%
Researcher 1 → shelves 25–49%
Researcher 2 → shelves 50–74%
Researcher 3 → shelves 75–100%

When a researcher needs a book stored elsewhere, the system moves it temporarily.

That gives us today's main principle:

And that connects directly to yesterday.

Activation checkpointing says:

Don't store an activation if recomputing it later is cheaper.

ZeRO says:

Don't replicate training state if fetching or communicating it when needed is cheaper.

Both are examples of the same deeper strategy:

Once models outgrow a single accelerator, training performance is no longer determined only by how fast GPUs multiply matrices.

It depends on how intelligently the system decides what must be local, what can be sharded, and when data should move.