Fully Sharded Data Parallel (FSDP)

FSDP is PyTorch's native realization of the ZeRO stage-3 idea, made into a first-class module wrapper. You group the model into units — typically each transformer block — and FSDP shards every parameter, gradient, and optimizer-state tensor in that unit evenly across the data-parallel workers. At rest a worker holds only its 1/N slice of each weight, so the resident footprint is small even for very large models.

Execution proceeds just-in-time. Right before a unit's forward pass, FSDP issues an all-gather so every worker temporarily materializes the full weights, runs the layer, then frees the gathered copy. The backward pass repeats the all-gather to recompute gradients locally and then performs a reduce-scatter, leaving each worker with the averaged gradient shard it owns. Because gathers for the next unit can be prefetched while the current unit computes, much of the communication hides behind math. Mixed precision, activation checkpointing, and CPU offload all compose with the wrapping.

The practical art of FSDP is choosing the wrapping granularity. Wrap too coarsely and the transient all-gathered weights blow the memory budget you were trying to save; wrap too finely and you pay launch and synchronization overhead on tiny collectives. FSDP2 refines this with per-parameter sharding (DTensor) for cleaner composition with tensor parallelism and more predictable memory.