FSDP › Summary

FSDP (Fully Sharded Data Parallel) allows training of models too large to fit on a single GPU by sharding parameters, gradients, and optimizer states across all GPUs in a job.

The key ideas are:

• Sharding all three components: Parameters, gradients, and optimizer states are each split into $N$ equal fragments, so each GPU's steady-state memory is $1/N$ of the total. The full $16P$ bytes (float32 parameters, gradients, and Adam states) becomes $16P/N$ bytes per GPU.
• All-gather and reduce-scatter: FSDP reconstructs full parameter tensors just in time for computation via all-gather, then discards them immediately after use. Gradients are aggregated and re-sharded via reduce-scatter instead of all-reduce, making sure no GPU ever holds a complete gradient tensor.
• Sharding strategies: FULL_SHARD (maximum memory savings, corresponds to ZeRO Stage 3), SHARD_GRAD_OP (moderate savings, less communication, corresponds to ZeRO Stage 2), NO_SHARD (DDP equivalent for debugging), and HYBRID_SHARD (multi-node optimization that confines expensive all-gathers to the fast intra-node fabric).
• Relationship to ZeRO: FSDP with FULL_SHARD is algorithmically equivalent to DeepSpeed ZeRO Stage 3. The practical difference is ecosystem: FSDP is native PyTorch, ZeRO offers additional features like CPU offloading via ZeRO-Infinity.
• Wrapping policy: The granularity of FSDP units, controlled by the auto-wrap policy, determines the trade-off between all-gather frequency and buffer sizes. Wrapping at the transformer layer level is the standard recommendation.
• Mixed precision: The MixedPrecision policy allows parameters and activations in bfloat16 while keeping gradients and optimizer states in float32, balancing compute speed with numerical stability.
• Checkpointing: Full state dict mode produces portable checkpoints at the cost of a memory spike on rank 0. Local state dict mode is faster but tied to the specific FSDP configuration.

FSDP is the standard approach for large-scale model training in the PyTorch ecosystem and is directly supported by tools like Hugging Face Accelerate, PyTorch Lightning, and TorchTitan. Understanding its memory model and communication patterns equips you to reason about the practical constraints of large model training and to choose the right strategy for your hardware and model size.