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

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¶

Activation checkpointing inside FSDP (Ch 08-03)— recompute attention activations, cheap blocks.
Gradient accumulation across micro-batches (Ch 08-01) to shrink activation peak.
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)

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

-

Checkpointing FSDP¶

# Save
state = {
 "model": fsdp.state_dict(), # sharded view
 "optim": fsdp_optim.state_dict(),
 "step": step,
}
if rank == 0:
 torch.save(state, "ckpt.pt")

# load
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.

-

Ddp
[Checkpointing](/06-pytorch/05-distributed-training/(04-checkpointing-fault-tolerance/)
[Activation Checkpointing](/06-pytorch/08-advanced-training-techniques/(03-activation-checkpointing-memory-restructuring/)
Mixed Precision