DEV Community

Shrijith Venkatramana
Shrijith Venkatramana

Posted on

ZeRO and FSDP: How LLM Training Escapes the GPU Memory Wall

Hello, I'm Shrijith Venkatramana. I'm building git-lrc, an AI code reviewer that runs on every commit. Star Us to help devs discover the project. Do give it a try and share your feedback for improving the product.


A 70-billion-parameter model contains 70 billion numbers.

That sounds like the problem.

It isn't.

The real problem is that training those 70 billion numbers requires you to keep several different copies of information around: parameters, gradients, optimizer state, and sometimes higher-precision copies of the weights.

Put those together and a 70B model can require well over a terabyte of memory.

Now imagine trying to put that on GPUs with 80 GB each.

This is where one of the most important ideas in large-scale deep learning enters the story: don't make every GPU own everything.

That deceptively simple observation led to Microsoft's ZeRO (Zero Redundancy Optimizer) and, independently in the PyTorch ecosystem, Fully Sharded Data Parallel (FSDP).

Today, these techniques are among the fundamental building blocks behind large-scale LLM pretraining.

The interesting part is not merely that they save memory. It is that they turn a seemingly impossible memory problem into a distributed systems problem involving communication, computation, memory hierarchy, and economics.

1. First, remember what ordinary data parallelism actually does

Suppose we have 8 GPUs training the same Transformer.

With ordinary Distributed Data Parallel (DDP), every GPU gets:

  • a complete copy of the model parameters
  • a complete copy of the gradients
  • a complete copy of the optimizer state

The GPUs see different batches, compute gradients independently, and then synchronize those gradients.

Conceptually:

GPU 0: [entire model] + [entire optimizer] + batch 0
GPU 1: [entire model] + [entire optimizer] + batch 1
GPU 2: [entire model] + [entire optimizer] + batch 2
...
GPU 7: [entire model] + [entire optimizer] + batch 7
Enter fullscreen mode Exit fullscreen mode

This is beautifully simple.

It is also extremely wasteful.

The eight GPUs have eight copies of essentially the same training state.

The obvious question is:

Why does every GPU need to own the entire model if the GPUs are cooperating?

That question is the starting point for ZeRO.

The historical context

The late 2010s saw an extraordinary escalation in neural-network model sizes.

BERT-Large had roughly 340M parameters.

GPT-2 had 1.5B.

GPT-3, announced in 2020, had 175B.

At that scale, simply adding more GPUs did not solve the problem. Data parallelism gave you more aggregate memory, but the model still had to fit on every individual GPU.

Microsoft researchers Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase and Yuxiong He attacked precisely this redundancy problem.

Their 2019 paper introduced ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. The paper demonstrated training models over 100B parameters on 400 GPUs and argued that the approach could extend toward trillion-parameter models.

The key insight was wonderfully economical:

If the GPUs are already communicating, why replicate state that could instead be partitioned among them?

2. The arithmetic that makes the problem obvious

Let's do a rough calculation.

Suppose our model has:

P = 70 billion parameters
Enter fullscreen mode Exit fullscreen mode

During mixed-precision training with Adam, a useful back-of-the-envelope estimate is roughly:

parameters             ~2 bytes/parameter
gradients              ~2 bytes/parameter
FP32 master weights    ~4 bytes/parameter
Adam first moment      ~4 bytes/parameter
Adam second moment     ~4 bytes/parameter
                         ----------------
                         ~16 bytes/parameter
Enter fullscreen mode Exit fullscreen mode

So:

70B * 16 bytes
= 1.12 TB
Enter fullscreen mode Exit fullscreen mode

That's per replica.

Not per cluster.

Per GPU, if you're doing ordinary data parallelism.

An 80 GB GPU cannot possibly hold that.

And this calculation hasn't even included:

  • activations
  • temporary buffers
  • CUDA workspace
  • communication buffers
  • fragmentation
  • attention-related memory
  • checkpointing overhead

So simply saying:

"I'll train my 70B model on 16 x 80GB GPUs"

doesn't work if every GPU needs a complete training replica.

The aggregate GPU memory is:

16 * 80 GB = 1.28 TB
Enter fullscreen mode Exit fullscreen mode

which sounds sufficient.

But distributed memory doesn't magically become one giant 1.28 TB RAM pool.

DDP effectively asks each GPU:

"Can YOU hold 1.12 TB?"
Enter fullscreen mode Exit fullscreen mode

The answer is no.

This distinction is fundamental:

Aggregate memory is not the same thing as usable memory.

ZeRO/FSDP changes the answer by allowing the aggregate memory of the cluster to actually be used.

3. ZeRO's trick: eliminate redundancy one state at a time

ZeRO doesn't begin by radically changing data parallelism.

Instead, it asks:

What exactly are we redundantly storing?

There are three major pieces of replicated training state:

  1. optimizer states
  2. gradients
  3. parameters

ZeRO removes the redundancy progressively.

ZeRO Stage 1

Shard the optimizer states.

Instead of:

GPU 0: Adam states for all parameters
GPU 1: Adam states for all parameters
GPU 2: Adam states for all parameters
...
Enter fullscreen mode Exit fullscreen mode

we do:

GPU 0: Adam states for 1/8 of parameters
GPU 1: Adam states for 1/8
...
GPU 7: Adam states for 1/8
Enter fullscreen mode Exit fullscreen mode

The model parameters themselves are still replicated.

ZeRO Stage 2

Shard:

optimizer states
+
gradients
Enter fullscreen mode Exit fullscreen mode

Now each GPU only owns its fraction of the gradients and optimizer state.

ZeRO Stage 3

Shard everything:

parameters
gradients
optimizer states
Enter fullscreen mode Exit fullscreen mode

Now the persistent training state is distributed across the GPUs.

For N GPUs, the rough persistent-memory cost becomes:

~16P / N bytes
Enter fullscreen mode Exit fullscreen mode

instead of:

~16P bytes
Enter fullscreen mode Exit fullscreen mode

per GPU.

For our 70B example on 64 GPUs:

70B * 16 / 64
= 17.5 GB
Enter fullscreen mode Exit fullscreen mode

per GPU.

Suddenly an 80 GB GPU has a plausible starting point.

There is an important catch, however.

The model parameters needed for computation must temporarily exist in full.

And this is where the next idea becomes crucial.

4. The clever part: don't keep the full model around

Suppose a Transformer has 80 layers.

You don't need all 80 layers' full parameters sitting on a GPU while you're computing layer 17.

You need layer 17.

So imagine that each GPU permanently stores only its shard:

GPU 0: 1/64 of parameters
GPU 1: 1/64
...
GPU 63: 1/64
Enter fullscreen mode Exit fullscreen mode

When layer 17 is about to execute, the GPUs perform an all-gather.

Each GPU contributes its shard, and everybody temporarily obtains the complete parameters for that layer.

Conceptually:

          shard 0
          shard 1
          shard 2
             ...
          shard 63
             |
             v
        ALL-GATHER
             |
             v
     [complete layer 17]
             |
          compute
             |
             v
        free full copy
Enter fullscreen mode Exit fullscreen mode

Then you move on.

The full layer doesn't need to remain resident after computation.

This is the central mental model for FSDP/ZeRO-3:

Keep the model sharded most of the time; temporarily materialize the piece you are computing.

PyTorch's FSDP implementation describes essentially this mechanism: parameters are sharded across workers, all-gathered before computation, and the unsharded parameters can then be freed to recover memory.

This is why the word fully in Fully Sharded Data Parallel matters.

It isn't merely sharding the batch.

It is sharding the model's training state.

5. FSDP: the PyTorch incarnation of the same idea

If you have encountered both terms, the relationship can initially be confusing.

ZeRO is the Microsoft/DeepSpeed family of techniques.

FSDP is PyTorch's native fully-sharded data-parallel implementation.

The conceptual overlap is substantial. In fact, PyTorch explicitly describes FSDP as being inspired by ZeRO Stage 3.

A useful simplified mapping is:

DDP
 |
 +-- ZeRO-1 / partial sharding
 |
 +-- ZeRO-2
 |
 +-- ZeRO-3
       |
       +-- FSDP-style full sharding
Enter fullscreen mode Exit fullscreen mode

FSDP preserves much of the programming model of ordinary data parallelism:

model = FSDP(model)
optimizer = AdamW(model.parameters(), ...)
Enter fullscreen mode Exit fullscreen mode

rather than requiring the developer to manually split every Transformer layer across devices.

That's an important engineering achievement.

Because distributed training has two very different problems:

Can the mathematics be distributed?
        +
Can humans actually program it?
Enter fullscreen mode Exit fullscreen mode

The second problem is often underestimated.

Meta's FSDP work emerged from precisely this ecosystem. The original FSDP implementation drew on earlier work around sharded optimizers and model parallelism, and was subsequently integrated into PyTorch. Meta reported experiments reaching large-model training at scale, including a 1T-parameter GPT configuration in its early FSDP work.

The later PyTorch FSDP scaling work reported near-linear scaling in TFLOPS while supporting substantially larger models than ordinary DDP.

6. But there is no free lunch: memory becomes communication

Here is the really interesting systems tradeoff.

DDP has:

replicated parameters
+
replicated optimizer
+
replicated gradients

but relatively simple communication
Enter fullscreen mode Exit fullscreen mode

FSDP/ZeRO-3 says:

less memory
+
more communication
Enter fullscreen mode Exit fullscreen mode

Every time you need a layer, you have to obtain its shards.

During forward:

sharded parameters
        |
        v
    all-gather
        |
        v
full parameters
        |
        v
    forward
        |
        v
free / reshard
Enter fullscreen mode Exit fullscreen mode

During backward, you similarly need the relevant parameters and eventually synchronize gradients with a reduce-scatter.

So the basic communication pattern becomes:

FORWARD

all-gather parameters
        |
     compute
        |
     reshard


BACKWARD

all-gather parameters
        |
     compute
        |
  reduce-scatter gradients
        |
     reshard
Enter fullscreen mode Exit fullscreen mode

This is why high-bandwidth GPU interconnects matter so much.

An A100 or H100 is enormously fast at computation.

But if your training system repeatedly waits for network communication, your expensive GPUs sit idle.

This is the fundamental distributed-training equation:

training efficiency
    ~ computation efficiency
      / communication + synchronization overhead
Enter fullscreen mode Exit fullscreen mode

Not literally a mathematical identity, but an excellent engineering mental model.

The whole game is to make communication happen while computation is happening.

For example:

compute layer 10
       ||
       || communication for layer 11
       ||
compute layer 11
       ||
       || communication for layer 12
       ||
compute layer 12
Enter fullscreen mode Exit fullscreen mode

rather than:

communicate
wait
compute
wait
communicate
wait
compute
Enter fullscreen mode Exit fullscreen mode

This is why FSDP implementations contain machinery such as prefetching, communication buckets, asynchronous collectives, and computation/communication overlap.

The PyTorch FSDP scaling work specifically emphasizes these engineering techniques as part of making fully sharded training practical at scale.

7. The economics: you're trading GPU memory for network bandwidth

This is where ZeRO/FSDP stops looking like a clever PyTorch feature and starts looking like infrastructure economics.

Consider two hypothetical architectures.

Architecture A: enormous GPUs

Suppose a hypothetical GPU has enough memory to hold your entire model and optimizer state.

You get:

simple programming
low communication
high memory requirement
expensive GPUs
Enter fullscreen mode Exit fullscreen mode

Architecture B: many smaller-memory GPUs

Instead:

GPU 0 owns shard 0
GPU 1 owns shard 1
...
GPU N owns shard N
Enter fullscreen mode Exit fullscreen mode

You get:

lower memory requirement per GPU
more GPUs
more communication
more networking requirements
more distributed-systems complexity
Enter fullscreen mode Exit fullscreen mode

The optimal solution depends on the relative price of:

GPU memory
GPU compute
network bandwidth
network latency
power
rack capacity
engineering time
Enter fullscreen mode Exit fullscreen mode

This is why large-scale AI infrastructure is not simply:

"Buy the fastest GPU."

It is closer to:

"Find the cheapest system that keeps enough expensive compute occupied enough of the time."

If communication reduces your GPU utilization from 50% to 20%, buying more GPUs may make the system worse economically.

But if communication can be overlapped with computation and you achieve something close to linear scaling, sharding becomes enormously attractive.

A concrete 70B thought experiment

Suppose you have:

70B parameters
80 GB GPU
64 GPUs
Enter fullscreen mode Exit fullscreen mode

Using the rough 16 bytes/parameter estimate:

unsharded training state:

70B * 16
= 1.12 TB per GPU
Enter fullscreen mode Exit fullscreen mode

Impossible.

With full sharding:

1.12 TB / 64
= 17.5 GB per GPU
Enter fullscreen mode Exit fullscreen mode

Now there is substantial room for:

temporary unsharded layer parameters
+
activations
+
CUDA workspaces
+
communication buffers
Enter fullscreen mode Exit fullscreen mode

But notice what happened.

We did not magically compress 1.12 TB into 17.5 GB.

We distributed the 1.12 TB across 64 machines.

And when a GPU needs something it doesn't own, it asks the other GPUs for it.

That is the entire conceptual leap.

8. ZeRO/FSDP is not the whole recipe for frontier-model training

One subtle point is worth emphasizing.

FSDP solves a particular problem:

How do we make data-parallel training state fit across many devices?

It does not mean that arbitrary 1T+ models can simply be wrapped in FSDP and trained efficiently.

At sufficiently large scales, modern systems often combine several forms of parallelism.

For example:

                 Training
                    |
       +------------+------------+
       |            |            |
   Data/FSDP    Tensor Parallel  Pipeline
   parallelism    parallelism    parallelism
Enter fullscreen mode Exit fullscreen mode

You might have:

FSDP / ZeRO
    +
tensor parallelism
    +
pipeline parallelism
    +
activation checkpointing
    +
mixed precision
Enter fullscreen mode Exit fullscreen mode

Each attacks a different bottleneck.

Tensor parallelism splits individual matrix operations across GPUs.

Pipeline parallelism places different layers on different GPUs.

FSDP/ZeRO shards the training state across data-parallel workers.

This distinction matters because developers sometimes encounter a 3D-parallel training system and think:

"Why not just use FSDP?"

Because once the model becomes sufficiently large, communication inside a single enormous layer can itself become the bottleneck.

FSDP is best understood as one dimension in a larger distributed-training architecture.

9. The surprising historical trajectory: from optimization trick to default infrastructure

One of the interesting things about ZeRO is how quickly the idea escaped its original paper.

The 2019 ZeRO paper was explicitly motivated by trillion-parameter training. It showed that eliminating redundancy in data-parallel training could dramatically increase the feasible model size without requiring developers to manually adopt complicated model-parallel programming.

Then came increasingly aggressive extensions.

ZeRO-Offload pushed optimizer state and other training state into CPU memory, exploiting the fact that CPU memory is vastly cheaper and more plentiful than GPU memory. The published results demonstrated, for example, 10B-parameter training on a single V100 with substantially higher attainable model size than conventional PyTorch training.

Then ZeRO-Infinity extended the basic idea further into a heterogeneous memory hierarchy involving GPU memory, CPU memory, and storage.

Meanwhile, the PyTorch ecosystem developed FSDP into an increasingly native distributed-training abstraction.

The trajectory is revealing:

DDP
 |
 | "Why replicate everything?"
 v
ZeRO-1
 |
 | "Why replicate gradients?"
 v
ZeRO-2
 |
 | "Why replicate parameters?"
 v
ZeRO-3
 |
 | "Why keep everything on GPU?"
 v
Offload / heterogeneous memory
 |
 | "Can the framework make this transparent?"
 v
FSDP
Enter fullscreen mode Exit fullscreen mode

The deeper idea is not "use FSDP."

It is:

Treat memory as a distributed resource rather than a property of an individual accelerator.

That is a much more general systems principle.

10. What a developer should actually remember

If you are building or debugging an LLM pretraining system, I would keep five mental models in your head.

1. DDP replicates

Every GPU owns:

model + gradients + optimizer
Enter fullscreen mode Exit fullscreen mode

Simple, robust, but memory-inefficient.

2. ZeRO partitions

Progressively shard:

optimizer
    -> gradients
        -> parameters
Enter fullscreen mode Exit fullscreen mode

3. FSDP materializes temporarily

The GPU usually owns only a shard.

It temporarily reconstructs what it needs:

shard
  -> all-gather
  -> compute
  -> free
Enter fullscreen mode Exit fullscreen mode

4. Memory savings become communication costs

The important question becomes:

Can I move parameters between GPUs
fast enough that computation doesn't stall?
Enter fullscreen mode Exit fullscreen mode

5. Scaling is an optimization problem across the whole machine

The objective isn't:

minimize GPU memory
Enter fullscreen mode Exit fullscreen mode

It is closer to:

maximize useful training FLOPs / dollar
Enter fullscreen mode Exit fullscreen mode

subject to:

memory
network bandwidth
latency
compute
power
fault tolerance
checkpointing
software complexity
Enter fullscreen mode Exit fullscreen mode

That is why distributed LLM training is fundamentally a systems problem.

Conclusion: the trick was never really "zero redundancy"

The most important conceptual shift behind ZeRO and FSDP is surprisingly simple.

A model does not have to exist in its entirety on every GPU merely because every GPU participates in training it.

Instead, you can arrange the system so that:

persistent state
    = distributed

temporary computation state
    = local
Enter fullscreen mode Exit fullscreen mode

That lets a cluster behave, from the perspective of the training algorithm, almost like one enormous distributed memory system.

And this is a recurring pattern in computer science:

When one machine cannot hold the thing you want, don't immediately make the thing smaller. First ask whether the thing really needs to be in one machine.

ZeRO asked that question of optimizer states.

Then gradients.

Then parameters.

FSDP brought the resulting abstraction directly into PyTorch's distributed-training stack.

The result is one of the key pieces of infrastructure that made the modern LLM scaling story possible.


*AI agents write code fast. They also silently remove logic, change behavior, and introduce bugs -- without telling you. You often find out in production.

git-lrc fixes this. It hooks into git commit and reviews every diff before it lands. 60-second setup. Completely free.*

Any feedback or contributors are welcome! It's online, source-available, and ready for anyone to use.

GitHub logo HexmosTech / git-lrc

Free, Micro AI Code Reviews That Run on Git Commit




GenAI today is a race car without brakes. It accelerates fast -- you describe something, and large blocks of code appear instantly. But AI agents silently break things: they remove logic, relax constraints, introduce expensive cloud calls, leak credentials, and change behavior -- without telling you. You often find out in production.

git-lrc is your braking system. It hooks into git commit and runs an AI review on every diff before it lands. 60-second setup. Completely free.

In short, git-lrc helps Prevent Outages, Breaches, and Technical Debt Before They Happen

At a glance: 10 risk categories · 100+ failure patterns tracked · every commit…

Top comments (0)