FSDP and ZeRO — training a model without replicating every optimizer byte

Yesterday we explored pipeline parallelism: dividing a Transformer into consecutive stages so multiple GPUs can work on different microbatches. Before that, tensor parallelism split individual layers across GPUs.

Both techniques divide the model's computation. Today we'll ask a different question:

What if every GPU could execute the full model, but no GPU had to permanently store all the model's training state?

That is the idea behind ZeRO (Zero Redundancy Optimizer) and Fully Sharded Data Parallelism (FSDP). They make large-model training practical by eliminating redundant copies of parameters, gradients, and optimizer state across data-parallel workers.

Why training needs much more memory than inference¶

During inference, we mainly store model weights, activations needed for the current computation, and the KV cache. During training, we also need gradients and the optimizer's running statistics.

Consider one common mixed-precision Adam training configuration:

This is an illustrative accounting, not a universal constant: some optimizers, precision modes, and implementations store gradients or master weights differently.

For a 7-billion-parameter model, these persistent states alone require roughly

That is about 104 GiB, before counting activations, temporary buffers, and communication workspace. Merely fitting the BF16 inference weights (about 14 GB) tells us surprisingly little about whether we can train the model.

Ordinary data parallelism duplicates everything¶

In conventional data parallelism, each GPU has a complete model replica and processes different examples. Workers then combine their gradients so they can make a consistent optimizer update.

GPU 0: full weights + full gradients + full Adam state → data shard A
GPU 1: full weights + full gradients + full Adam state → data shard B
GPU 2: full weights + full gradients + full Adam state → data shard C
GPU 3: full weights + full gradients + full Adam state → data shard D
                         ↓
                  synchronize gradients

This scales compute throughput across examples, but it does not reduce the persistent training-state memory required per GPU. Four replicas of a 112 GB state still need roughly 112 GB each, not 28 GB each.

The key insight behind ZeRO is that these replicas contain enormous amounts of redundant state.

Three progressively deeper levels of sharding¶

ZeRO describes three stages of removing that redundancy:

• ZeRO-1: shard optimizer state across data-parallel workers.
• ZeRO-2: additionally shard gradients.
• ZeRO-3: additionally shard model parameters.

FSDP's fully sharded mode follows a closely related strategy to ZeRO-3: parameters, gradients, and optimizer state are distributed rather than persistently replicated. The frameworks differ in implementation details, configuration, and supported execution patterns.

For the illustrative 16-bytes-per-parameter configuration above, with \(p\) workers and ideal balanced sharding, the persistent memory per parameter is approximately

The 12 bytes in ZeRO-1 represent the FP32 master copy and Adam moments; ZeRO-2 also shards the two-byte gradients. Actual memory depends on how a framework organizes and buffers these states.

Worked example: training 7B parameters on eight GPUs¶

Set \(p=8\). The idealized persistent training-state footprint becomes:

This is a dramatic reduction: 112 GB to 14 GB per GPU, ideally an 8× reduction in persistent training-state memory.

But 14 GB is not the peak memory needed to train. Each worker still needs activations, temporary all-gathered weights, communication buffers, and kernel workspace. Longer sequences and larger microbatches can make activation memory substantial. Activation checkpointing addresses that separate problem by recomputing some activations during backward instead of retaining them.

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

This is the clever part. A fully sharded worker normally stores only its assigned parameter shards, but it can temporarily reconstruct a layer's full parameters just before using them.

A simplified FSDP execution looks like this:

Across 4 GPUs, each owns a shard of layer L:

GPU 0: [L0]   GPU 1: [L1]   GPU 2: [L2]   GPU 3: [L3]
                 |
          ALL-GATHER layer L
                 v
Each GPU temporarily has [L0 L1 L2 L3]
                 |
     Each GPU runs layer L on its own data
                 |
       Release unneeded full weights
                 |
      BACKWARD: gather again if needed
                 |
       Compute local gradient contributions
                 |
       REDUCE-SCATTER gradients
                 v
Each GPU retains and updates only its assigned shards

All-gather assembles a complete layer from pieces held by different GPUs. Reduce-scatter sums gradient contributions across workers and leaves each worker with only its designated gradient shard.

The important distinction from tensor parallelism is that each FSDP worker performs the full layer computation on its own data, once the layer's parameters have been gathered. Tensor-parallel GPUs instead cooperate on different pieces of the same layer computation.

The new bottleneck: communication¶

Memory savings do not come for free. Suppose the 7B model has 14 GB of BF16 weights and we use eight fully sharded workers.

For an ideal all-gather, each GPU receives approximately

If gathered weights are released after forward and needed again for backward, there may be two parameter all-gathers per training step. A gradient reduce-scatter adds another transfer of roughly \(12.25\) GB per worker when BF16 gradients are assumed.

That gives a useful, deliberately simplified traffic estimate:

This counts the main collective payloads, not every network byte or buffer copy. Real implementations communicate layer by layer, may retain some gathered parameters, overlap transfers with computation, and use different gradient precision or collective algorithms.

The tradeoff is now visible: we saved persistent GPU memory by repeatedly moving parameters over the interconnect. If the network is slow or layers are too small to overlap communication with useful computation, the sharded system can run slower than replicated training.

How this fits with tensor and pipeline parallelism¶

Data parallelism:     each GPU trains on different examples
Tensor parallelism:   GPUs split the math within a layer
Pipeline parallelism: GPUs own different sequences of layers
FSDP / ZeRO:          data-parallel GPUs share storage of training state

These are not mutually exclusive. A large training job might use tensor parallelism within a fast-connected GPU group, pipeline parallelism between groups of layers, and sharded data parallelism across replicas of that combined arrangement.

The optimal combination depends on model size, activation memory, microbatch size, GPU memory capacity, interconnect topology, and how much communication can overlap computation.

Mental model: a shared reference library¶

Imagine eight analysts doing similar work on different customer accounts. Ordinary data parallelism gives each analyst a complete, expensive reference library at their desk. It is convenient, but eight libraries occupy eight times the shelf space.

FSDP instead distributes the books among the analysts. When everyone needs Chapter 12, they temporarily assemble the chapter from the distributed collection, do their separate calculations, and then put the shared material away. Each analyst keeps only the records they are responsible for updating.

The office saves enormous shelf space—but now depends on a very fast delivery system. If requesting chapters takes longer than reading them, the new arrangement becomes inefficient.

The deeper lesson¶

Tensor and pipeline parallelism ask how to distribute computation. ZeRO and FSDP ask which training state actually needs to be resident on each GPU at any moment.

That is often a worthwhile exchange: a model that cannot fit with replicated optimizer state may become trainable on existing GPUs. But the right benchmark is not just bytes saved. It is end-to-end training throughput, peak memory, and network cost at the intended sequence length and batch size.

The practical question is: How much redundant training state can we eliminate before reconstructing it becomes the dominant cost?