FSDP Deep Dive

Overview

Fully Sharded Data Parallel (FSDP) shards model parameters, gradients, and optimizer state across GPUs — the strategy of choice for large models (multi-billion params) that don't fit on a single GPU. It trades extra communication for dramatically lower memory.

• Shards the model itself (parameters split across ranks), not just gradients.
• Automatically all-gathers params before forward/backward, reduce-scatters grads.
• ShardingStrategy: FULL_SHARD (params+grads+optimizer sharded), SHARD_GRAD_OP (params unsharded), NO_SHARD.
• Wraps submodules (auto_wrap_policy) so shards match transformer block boundaries.
• use_orig_params=True for param-compatible APIs (param.grad, LoRA).

💡 FSDP memory ≈ parameters/world_size + optimizer shard + activations. Communication ≈ DDP-level or slightly higher, but overlap hides it.

---

Minimal FSDP Setup

import torch, torch.nn as nn, torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

def main(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    model = nn.Sequential(*[nn.TransformerEncoderLayer(128, 4) for _ in range(4)])
    model = FSDP(model, device_id=rank)

    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    x = torch.randn(8, 16, 128).cuda()

    for step in range(5):
        opt.zero_grad()
        loss = model(x).abs().mean()
        loss.backward()          # reduce-scatter grads internally
        opt.step()               # optimizer sees only its shard's params
    dist.destroy_process_group()

if __name__ == "__main__":
    import torch.multiprocessing as mp
    mp.spawn(main, args=(2,), nprocs=2, join=True)

---

Which Strategy When

Strategy	Params	Grads	Optimizer	Memory	Comm
DDP	full	full	full	high	grads only
FSDP FULL_SHARD	sharded	sharded	sharded	low	params+g+o
FSDP SHARD_GRAD_OP	full	sharded	sharded	medium	grads
FSDP NO_SHARD	full	full	full	high	grads (≈DDP)

💡 For most LLM training: FULL_SHARD + auto_wrap on the transformer block.

---

Auto-Wrap Policy — Why It Matters

Wrapping per-block enables: all-gather only the shards used by the current block (not whole model), and overlap comm/compute block-by-block.

from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy, size_based_auto_wrap_policy
import torch.nn as nn

transformer_block = nn.TransformerEncoderLayer
policy = transformer_auto_wrap_policy(
    transformer_layer_cls={transformer_block},
)

model = nn.Sequential(*[transformer_block(128, 4) for _ in range(4)])
fsdp = FSDP(model, auto_wrap_policy=policy, device_id=rank)

Or size-based (wrap modules above threshold):

size_policy = size_based_auto_wrap_policy(min_num_params=5_000_000)

---

use_orig_params — The Pragmatic Choice

fsdp = FSDP(
    model,
    auto_wrap_policy=policy,
    use_orig_params=True,     # param objects behave like normal params
    device_id=rank,
)
# benefits: param.grad works, torch.compile-friendly, LoRA-friendly

⚠️ With use_orig_params=False (default), FSDP-created "flat" params differ: model-specific APIs (param.grad, .parameters()) behave differently. Prefer True in modern code.

---

Memory Optimization Tips

1. Activation checkpointing inside FSDP (Ch 08-03) — recompute attention activations, cheap blocks.
2. Gradient accumulation across micro-batches (Ch 08-01) to shrink activation peak.
3. Mixed precision via FSDP's own mixed_precision setting (bf16 params+grads).

from torch.distributed.fsdp import ShardingStrategy, MixedPrecision
bf16 = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)
fsdp = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD,
            mixed_precision=bf16, device_id=rank)

1. CPU offload (cpu_offload=CPUOffload(offload_params=True)) when massaging the last few GB, at big comm cost.

---

Checkpointing FSDP

# Save: only rank 0 writes, others skip (optimizer shards live on each rank)
state = {
    "model": fsdp.state_dict(),          # sharded view
    "optim": fsdp_optim.state_dict(),
    "step": step,
}
if rank == 0:
    torch.save(state, "ckpt.pt")

# load: all ranks load, state_dict() reshards automatically (ranks-independent)
fsdp.load_state_dict(torch.load("ckpt.pt")["model"], strict=True)

⚠️ Older fsdp used ShardedStateDict/reshard; modern state_dict() handles the sharded contract for you. Save/load on ALL ranks (or rank 0 + broadcast).

---

Debugging FSDP

• TORCH_DISTRIBUTED_DEBUG=DETAIL — dumps comm/params diagnostics.
• fsdp.summarize().print_summary() — prints sharding layout.
• Watch for the classic device mismatch: params on meta/CPU vs inputs on GPU during wrap.
• Profile: all-gather shows up as ncclKernel_*; if comm % high, increase block size / check cross-node bandwidth.

---

Key Takeaways

• FSDP = shard everything + all-gather/reduce-scatter around blocks; memory ∝ 1/world_size.
• Wrap per-transformer-block (transformer_auto_wrap_policy) — the single most important config.
• Use use_orig_params=True and FULL_SHARD as your default in modern PyTorch.
• Add activation checkpointing + gradient accumulation + bf16 mixed precision before reaching for CPU offload.
• Save/load on one rank (or all) with plain state_dict() — sharded contract is handled.

---

Related Topics

• Ddp
• Checkpointing
• Activation Checkpointing
• Mixed Precision