Also known as: Fully Sharded Data Parallel, ZeRO-3, FSDP2
FSDP shards parameters, gradients, and optimizer state across data-parallel ranks instead of replicating them. Each rank holds only 1/N of the weights at rest and gathers full layers on the fly during forward and backward.
FSDP — Fully Sharded Data Parallel — is the natural extension of data parallelism for models that no longer fit on one GPU. Instead of replicating the full model on every rank, FSDP shards the parameters, gradients, and optimizer state across all data-parallel ranks. Each rank holds only
FSDP is PyTorch’s native answer to DeepSpeed ZeRO and is the default for training any model that is “bigger than one GPU but smaller than one node.” For larger jobs, it composes with tensor and pipeline parallelism rather than replacing them.
Conceptually FSDP wraps each transformer block (or any user-chosen unit) and runs three collectives around it:
The reduce-scatter is the FSDP analog of DDP’s gradient allreduce, just split across the sharding axis. The two extra all-gathers are the price of not replicating parameters.
Take a 70B-parameter model in mixed precision with Adam. Resident state per replica:
Total: ~1.12 TB per replica. No single GPU has that. With plain DDP across 16 H100s (80 GB each), every rank still needs 1.12 TB locally — DDP does not help.
With FSDP across the same 16 ranks, each rank holds 1/16 of every tensor: ~70 GB resident. Add transient gathered-parameter buffers (one layer’s worth in bf16, ~3-5 GB for typical transformer blocks) and the activation memory budget, and a 70B model trains comfortably on a single 16-GPU node. This is the entire point of FSDP.
PyTorch FSDP exposes a small set of sharding choices:
The sharding unit also matters. FSDP wraps modules — typically each transformer block is its own FSDP unit. Wrapping too coarsely (the whole model as one unit) means you all-gather every parameter at once and lose the memory savings during forward / backward. Wrapping too finely (every linear layer separately) burns time on tiny collectives.
MixedPrecision policy.The other operational quirk is checkpointing. Naive state_dict() collects the full unsharded model on rank 0 — fine for a 7B model, fatal for a 70B one. Production FSDP code uses sharded checkpointing (each rank writes only its slice) and reconstructs the full state at load time, with torch.distributed.checkpoint as the canonical path.
FSDP is the boring-but-correct answer for the meaty middle of the model-size distribution. The day a model stops fitting on one GPU is the day you reach for it; the day activations stop fitting is the day you start composing it with everything else.
ZeRO-1 shards optimizer state only. ZeRO-2 shards optimizer state and gradients. ZeRO-3 shards parameters, gradients, and optimizer state — which is what PyTorch FSDP implements. The naming differs, the technique is the same; ZeRO came first (DeepSpeed, Microsoft) and FSDP is PyTorch's native re-implementation.
Roughly 1.3-1.5x more communication per step. DDP does one gradient allreduce; FSDP adds an all-gather for parameters before forward, an all-gather before backward, and a reduce-scatter for gradients. With overlap and bucketing, end-to-end throughput drop is usually 5-15% on intra-node NVLink. Cross-node, the gap widens and hybrid-shard becomes important.
When activations dominate memory rather than parameters — long-context training, MoE models with huge embedding tables, very wide FFN layers. FSDP shards model state but leaves activations replicated within a rank. Tensor parallelism shards activations along the feature axis. Production frontier-scale training combines both: TP within a node, FSDP across nodes.
