Scale and post-training
ZeRO and FSDP
ZeRO and FSDP remove the copies that data parallelism makes of optimizer state, gradients and weights. Each GPU keeps only a shard and gathers the rest when a layer needs it.
Data parallelism stores the same optimizer state, gradients and weights on every GPU. That is wasteful. ZeRO (Zero Redundancy Optimizer) observes that these copies are redundant, and shards each one across the data-parallel group, so every GPU stores only a slice. FSDP (Fully Sharded Data Parallel) is PyTorch's native version of the most aggressive stage.
Concept
Before you start: sharding means splitting a tensor into pieces, one per GPU. A GPU keeps its own piece at rest, and it borrows the others temporarily when it needs the full tensor.
The three stages
Each stage shards more of the state. The more you shard, the more communication you add.
| Stage | What is sharded | What it saves |
|---|---|---|
| ZeRO-1 | Optimizer state | Most of the 12 bytes per parameter |
| ZeRO-2 | Optimizer state and gradients | Gradient memory too |
| ZeRO-3 / FSDP | Optimizer state, gradients and weights | Weights are gathered only when needed |
How FSDP moves data
Each rank owns a shard of every parameter. For each unit of the model, such as one transformer block:
- All-gather the full weights for that unit, so the forward pass can run.
- Free the gathered weights when the unit is done, keeping only the shard.
- In the backward pass, all-gather the weights again, compute gradients.
- Reduce-scatter the gradients, so each rank keeps only its shard of the summed gradient.
Communication rises compared with plain data parallelism, because weights are gathered repeatedly. Memory falls enough to train models that would not fit otherwise.
Example
Take the 112 GB of training state for a 7B model from the memory math page. Across 8 GPUs, full sharding leaves about 112 / 8 = 14 GB of state per GPU at rest. This is the ideal figure. Gathered layers and activations need extra room on top.
Choosing between them
Tip
Start with FSDP or ZeRO-2 or ZeRO-3 when a single replica does not fit. Stay on plain DDP when the model fits comfortably, because DDP has the simplest communication and often gets the best utilization on smaller models.
DeepSpeed made ZeRO popular. Hugging Face Accelerate, Lightning and torch-distributed all expose versions of these stages. Pick one stack and learn its options before comparing them.
Common mistakes
- Sharding a model that already fits. You add gathers and gain nothing. Use DDP.
- Ignoring gather traffic on slow links. FSDP gathers weights many times per step. Check that the network can keep up.
- Expecting ZeRO to shrink activations. It shards weights, gradients and optimizer state. Activation memory needs checkpointing.
Check yourself
- Which state does ZeRO-1 shard? Only the optimizer state.
- Why does FSDP use more communication than DDP? It gathers the full weights for each unit during forward and backward, instead of keeping a full copy on every GPU.
Key takeaways
- ZeRO and FSDP shard the state that data parallelism copies everywhere.
- Each stage shards more state and adds more communication.
- FSDP gathers weights per layer, runs the layer, and frees them again.
- Use sharding when a replica does not fit. Otherwise plain data parallelism is simpler.
Further learning
When sharding is not enough, you split the layers and the matrices themselves: tensor, pipeline and context parallelism.
Related
Distributed training
How to train models too big for one GPU, using data, sharding, tensor, pipeline and context parallelism. Covers the memory math that forces the choice and how to judge the result.
Data parallelism and all-reduce
Data parallelism gives each GPU a full copy of the model and a different slice of the batch. The GPUs then sum their gradients with an all-reduce, so every copy applies the same update.
Memory math for training
Training memory is much more than the weights. Mixed-precision Adam needs about 16 bytes per parameter before any activations, which is why a 7B model needs over 100 GB.