FSDP › FSDP vs. ZeRO

FSDP and Microsoft's ZeRO optimizer (used in DeepSpeed) address the same fundamental problem and use nearly identical techniques. Understanding their relationship clarifies why FSDP exists and how to choose between them.

The ZeRO paper, "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models," was published by Rajbhandari et al. at Microsoft Research in 2020. At the time, the dominant approach for large model training was model parallelism, which splits a model across GPUs by assigning different layers to different devices. Model parallelism requires careful pipeline management to keep all GPUs busy and introduces complex inter-layer communication patterns. ZeRO proposed a different insight: instead of splitting the model structurally, shard the redundant state that data parallelism was already replicating across GPUs.

The insight was that in standard DDP with $N$ GPUs, the optimizer states, gradients, and even parameters are all replicated $N$ times. Each replica is identical; the replication exists purely for convenience, not necessity. By eliminating this redundancy through sharding, you can reduce per-GPU memory by a factor of $N$ while preserving the data-parallel communication structure.

ZeRO introduced three progressive stages of sharding, which became the conceptual framework that FSDP adopted.

ZeRO (Zero Redundancy Optimizer) defines three progressive stages of memory optimization:

• ZeRO Stage 1: Shards only optimizer states across GPUs. Each GPU holds $1/N$ of the optimizer states but the full parameters and gradients. This is the least disruptive change, requiring only that the optimizer step communicate parameter shards back to all GPUs after updating.

• ZeRO Stage 2: Shards optimizer states and gradients. Each GPU holds $1/N$ of optimizer states and $1/N$ of gradients but the full parameters. A reduce-scatter replaces the all-reduce for gradient synchronization.

• ZeRO Stage 3: Shards all three components. Each GPU holds $1/N$ of parameters, gradients, and optimizer states. All-gather operations are required before each layer's computation. This is the mode with maximum memory savings and highest communication overhead.

ZeRO Stage 3 is conceptually identical to what FSDP with FULL_SHARD implements. Both shard all three memory components and use all-gather plus reduce-scatter communication patterns. The algorithms are the same at the mathematical level.

The practical differences between FSDP and ZeRO Stage 3 are largely about ecosystem and integration rather than algorithmic substance.

FSDP is a native PyTorch feature. It integrates directly with the PyTorch autograd engine, requires no additional dependencies, and benefits from PyTorch's ongoing optimization work. Because it is part of the standard library, it works out of the box with tools like Hugging Face Accelerate, PyTorch Lightning, and most training frameworks that target PyTorch. The API is designed around PyTorch idioms: you wrap a model with FSDP() the same way you wrap with DDP(), and the training loop is otherwise unchanged.

DeepSpeed's ZeRO Stage 3 has a richer feature set around communication optimization, CPU offloading, and mixed precision handling. It also offers ZeRO-Infinity, which extends sharding to CPU and NVMe storage, enabling model sizes far beyond what fits in combined GPU memory. For researchers who need to push the absolute limit of model size on existing hardware, ZeRO-Infinity provides capabilities that FSDP does not yet match.

For most practical use cases, FSDP is the simpler choice if you are already in the PyTorch ecosystem. DeepSpeed is preferable when you need its extended optimization features or CPU offloading capabilities.