FSDP

FSDP (Fully Sharded Data Parallel) is PyTorch’s native implementation of ZeRO-3. It shards parameters, gradients, and optimizer states across all data-parallel ranks, gathering them on-the-fly during forward and backward.

---

1. How FSDP Works

The Lifecycle of a Training Step

1. FORWARD PASS:
   For each FSDP unit (typically one Transformer layer):
     a. AllGather: Collect full parameters from all ranks
     b. Forward compute with full parameters
     c. Discard non-local parameter shards (free memory)

2. BACKWARD PASS:
   For each FSDP unit (reverse order):
     a. AllGather: Collect full parameters again (for gradient computation)
     b. Backward compute
     c. ReduceScatter: Average gradients, keep only local shard
     d. Discard non-local parameters

3. OPTIMIZER STEP:
   Each rank updates only its local parameter shard
   (using its local optimizer state shard)

Memory During Execution

Phase	Parameters in memory	Gradients in memory
Idle (between steps)	Ψ/N (local shard only)	0
Forward (one layer)	Ψ/N + Ψ_layer (gathered)	0
Backward (one layer)	Ψ/N + Ψ_layer (gathered)	Ψ_layer (local only after ReduceScatter)
Optimizer step	Ψ/N	Ψ/N

Peak memory: Ψ/N (shards) + Ψ_layer (one layer gathered) + activations.

---

2. FSDP2 (PyTorch 2.x)

PyTorch 2.x introduced FSDP2 with a cleaner API:

import torch
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy

# Define mixed precision
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
)

# Apply FSDP to each Transformer layer
for layer in model.transformer.layers:
    fully_shard(layer, mp_policy=mp_policy)

# Apply FSDP to the whole model (outermost wrapper)
fully_shard(model, mp_policy=mp_policy)

# Training loop — same as regular PyTorch
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Key Differences: FSDP1 vs FSDP2

	FSDP1 (legacy)	FSDP2
API	FullyShardedDataParallel(module)	fully_shard(module)
Composability	Wraps module	Composable with TP, CP, PP
Mixed precision	MixedPrecision class	MixedPrecisionPolicy
Activation checkpointing	Separate wrapper	torch.utils.checkpoint
torch.compile	Limited	Full support

---

3. Sharding Strategies

Strategy	What’s sharded	Per-GPU memory	Communication
NO_SHARD	Nothing (DDP equivalent)	16Ψ	AllReduce gradients
SHARD_GRAD_OP	Gradients + optimizer (ZeRO-2)	2Ψ + 14Ψ/N	≈ AllReduce
FULL_SHARD	Everything (ZeRO-3)	16Ψ/N	3× AllReduce

FULL_SHARD is the default and most commonly used — it’s the whole point of FSDP.

---

4. FSDP Wrapping Granularity

You control which modules form an “FSDP unit” — the granularity of gather/scatter:

# Fine-grained: each Transformer layer is an FSDP unit
for layer in model.layers:
    fully_shard(layer)  # parameters gathered/scattered per layer

# Coarse-grained: entire model is one FSDP unit
fully_shard(model)  # all parameters gathered at once — huge memory spike!

Per-layer wrapping is standard: only one layer’s worth of parameters is fully materialized at a time.

Memory Impact of Wrapping Granularity

Wrapping	Peak param memory	Notes
Per-layer (standard)	Ψ/N + Ψ/L (one layer)	Minimal peak
Per-block (2 layers)	Ψ/N + 2Ψ/L	Slightly higher, fewer AllGathers
Whole model	Ψ (full model!)	Defeats the purpose of FSDP

---

5. FSDP + Tensor Parallelism

FSDP and TP compose naturally in PyTorch 2.x:

from torch.distributed.tensor.parallel import parallelize_module

# Step 1: Apply TP within each node
for layer in model.layers:
    parallelize_module(layer, tp_mesh, tp_plan)

# Step 2: Apply FSDP across nodes
for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)
fully_shard(model, mp_policy=mp_policy)

The combination: TP splits computation within a node (NVLink), FSDP shards memory across nodes (InfiniBand).

---

6. Common FSDP Pitfalls

Pitfall	Symptom	Fix
Wrapping too coarsely	OOM on AllGather (full model materialized)	Wrap per-layer
Not wrapping shared parameters	Incorrect gradients, training diverges	Use ignored_modules or handle manually
Mixed precision mismatch	Silently slower (FP32 compute) or NaN	Explicitly set MixedPrecisionPolicy
Forgetting to shard outermost model	Parameters not sharded, full memory	Always fully_shard(model) at the end
Saving/loading checkpoints wrong	Can’t resume from different GPU count	Use torch.distributed.checkpoint

For activation checkpointing with FSDP, see Activation Checkpointing. For ZeRO stages in DeepSpeed (alternative to FSDP), see ZeRO Stage 1/2/3.

---

Self-Check

• Explain the AllGather → compute → discard lifecycle of one FSDP forward layer
• Compare FSDP wrapping per-layer vs whole-model and the memory impact
• Write FSDP2 setup code for a Transformer model with mixed precision
• Explain how FSDP composes with TP (which dimension each handles)
• Choose between FSDP FULL_SHARD vs SHARD_GRAD_OP for a given scenario

---

References

• PyTorch FSDP Tutorial
• Composable FSDP2
• ZeRO: Memory Optimizations (Rajbhandari et al., 2019)