ZeRO & FSDP
Instead of having every GPU hold a copy of the entire model, we can slice the model like a pie. Every GPU only holds one slice, and they pass the slices around right when they are needed for the math.
The Memory Waste of Data Parallelism
In standard data-parallelism, you duplicate the entire model across multiple GPUs. If you have an 8-GPU cluster, you have 8 identical copies of the model.
This works brilliantly for small models, but it is a massive waste of memory. If a model takes up 40GB of memory, and you have 8 GPUs, you are using 320GB of VRAM just to store redundant copies of the exact same numbers.
Worse, it isn't just the model parameters. During training, you must also store:
- Gradients (same size as the parameters)
- Optimizer States (e.g., Adam uses momentum and variance, taking up 2x the size of the parameters).
This redundancy makes it impossible to train Large Language Models (LLMs) on standard Data Parallelism.
Enter ZeRO
ZeRO (Zero Redundancy Optimizer) is a memory optimization technology developed by Microsoft (and popularized by the DeepSpeed library) that solves this problem. ZeRO shards the memory across the GPUs, completely eliminating the redundancy while keeping the simplicity of Data Parallelism.
ZeRO does this in three stages.
ZeRO Stage 1: Optimizer State Partitioning
Instead of every GPU storing the full optimizer state (which is massive for Adam), the state is sharded. GPU 0 updates the first 25% of the weights, GPU 1 updates the next 25%, etc. At the end of the step, they broadcast the updated weights to each other.
- Memory Saved: ~4x reduction in optimizer memory.
ZeRO Stage 2: Gradient Partitioning
Building on Stage 1, we also shard the gradients. As soon as a gradient is computed during backpropagation, it is reduced and sent to the specific GPU responsible for updating that parameter. The local copy is then deleted.
- Memory Saved: ~8x reduction in gradient and optimizer memory.
ZeRO Stage 3: Parameter Partitioning
Building on Stages 1 and 2, we now shard the model weights themselves. The GPUs no longer hold a full copy of the model. Instead, right before a layer needs to do its matrix multiplication, it asks the other GPUs to send it the missing weights. It does the math, and then immediately deletes the weights it just downloaded.
- Memory Saved: Memory scales linearly with the number of GPUs. If you have 100 GPUs, your memory footprint per GPU drops by ~100x.
Fully Sharded Data Parallel (FSDP)
FSDP is PyTorch's native implementation of the concepts pioneered by ZeRO Stage 3. It is built directly into PyTorch's distributed library, making it much easier to use than installing third-party frameworks like DeepSpeed.
FSDP "shards" the model at the layer level. As the forward pass progresses through the network, FSDP dynamically gathers the parameters for the current layer using an All-Gather operation, computes the outputs, and then discards the parameters.
The Communication Tradeoff
ZeRO Stage 3 and FSDP sound like magic—you can train a trillion-parameter model on consumer GPUs just by adding more of them!
However, there is a catch: Network Bandwidth. Because the GPUs are constantly passing model weights back and forth, the network interconnect becomes the bottleneck. If your GPUs are connected by standard Ethernet instead of high-speed NVLink or InfiniBand, the GPUs will spend 90% of their time waiting for weights to arrive over the network, bringing your training speed to a crawl.
The Quick Version
- Data Parallelism wastes memory by storing redundant copies of the model on every GPU.
- ZeRO eliminates this redundancy by sharding the optimizer states, gradients, and parameters across the GPUs.
- FSDP is PyTorch's native implementation of ZeRO Stage 3.
- These techniques allow you to train massive models by combining the memory of multiple GPUs, but they require fast network interconnects to prevent communication bottlenecks.