FSDP (Fully Sharded Data Parallel)

Definition

FSDP is PyTorch’s native sharded data parallelism. FULL_SHARD is architecturally equivalent to ZeRO-3 (Zhao et al., Meta, arXiv:2304.11277): it shards parameters, gradients, and optimizer state across the DP group, dropping per-rank memory from $O(\Psi )$ to $O(\Psi /N)$ at the cost of one extra all-gather per unit per pass. Intermediate modes: SHARD_GRAD_OP ≈ ZeRO-2, NO_SHARD = plain DDP. “FSDP vs ZeRO” is a framework choice, not an algorithm choice.

Key math / How it works

The model is grouped into units (typically one transformer block) via a wrapping policy. Per unit at runtime:

1. All-gather the unit’s sharded params → full weights (forward).
2. Run the unit’s forward; free the full params (keep only the shard).
3. In backward, all-gather the params again, compute grads.
4. Reduce-scatter grads so each rank keeps only the shard matching its optimizer state; local optimizer step on that shard.

Because params are re-gathered in backward, total communication is $\Psi +\Psi +\Psi =3\Psi$ (1.5× DDP) — the same ZeRO-3 tradeoff. Overlap via prefetch is what makes it fast: all-gathers run on a separate CUDA stream so the collective for unit $i+1$ overlaps the compute of unit $i$ (forward prefetch, and analogously backward), breaking the false dependency that eager mode can’t reorder around. Only one unit’s full weights are ever resident.

FSDP1 vs FSDP2 (pin this — actively evolving):

• FSDP1 shards a FlatParameter: it flattens, concatenates, and chunks a group of a unit’s tensors into one 1-D buffer, then splits across ranks. Efficient, but the opaque flat buffer makes per-parameter reasoning (mixed dtypes, sharded state dicts, composing with TP) awkward.
• FSDP2 (the fully_shard API) replaces it with per-parameter sharding backed by DTensor: each parameter is chunked on dim-0 (torch.chunk(dim=0)) as a DTensor with Shard(dim=0) placement. This gives communication-free (already-sharded) state dicts, relaxed frozen-parameter constraints, out-of-the-box DTensor compatibility with optimizer/grad-clip/checkpoint, and — crucially — declarative composition with tensor parallelism (a weight can be Shard on the FSDP mesh dim and Shard/Replicate on the TP mesh dim). As of PyTorch 2.x, FSDP2 is the recommended path; FSDP1 is effectively maintenance. APIs are not backward-compatible (checkpoints largely portable).

Hybrid Sharded Data Parallel (HSDP): shard within a node, replicate across nodes — bounds the all-gather to fast intra-node NVLink while keeping inter-node traffic to a periodic all-reduce.

Why it matters

FSDP is the default way to train a model that won’t fit on one GPU in the PyTorch ecosystem, and the DTensor substrate is the plumbing that lets N-D parallelism compose (2-D FSDP+TP in TorchTitan). Its limit — comms volume grows and a single layer’s transient full weights/activations must still fit on one GPU — is exactly what motivates splitting the compute with tensor/pipeline parallelism.

Taught in

• 09-zero-and-fsdp — §5 FSDP, FSDP1 vs FSDP2.
• 11-sota-large-scale-training — §8 TorchTitan / framework landscape.

See also

• zero — the algorithm FSDP implements
• optimizer-state-memory — the state being sharded
• data-parallelism — NO_SHARD baseline
• tensor-parallelism — composes via DTensor (2-D parallelism)
• collective-communication — all-gather + reduce-scatter
• gpu-systems-for-llms