pytorch-fsdp2

Skill: Use PyTorch FSDP2 (fully_shard) correctly in a training script

This skill teaches a coding agent how to add PyTorch FSDP2 to a training loop with correct initialization, sharding, mixed precision/offload configuration, and checkpointing.

FSDP2 in PyTorch is exposed primarily via torch.distributed.fsdp.fully_shard and the FSDPModule methods it adds in-place to modules. See: references/pytorch_fully_shard_api.md, references/pytorch_fsdp2_tutorial.md.

---

When to use this skill

Use FSDP2 when:

• Your model doesn’t fit on one GPU (parameters + gradients + optimizer state).
• You want an eager-mode sharding approach that is DTensor-based per-parameter sharding (more inspectable, simpler sharded state dicts) than FSDP1.
• You may later compose DP with Tensor Parallel using DeviceMesh.

Avoid (or be careful) if:

• You need strict backwards-compatible checkpoints across PyTorch versions (DCP warns against this).
• You’re forced onto older PyTorch versions without the FSDP2 stack.

Alternatives (when FSDP2 is not the best fit)

• DistributedDataParallel (DDP): Use the standard data-parallel wrapper when you want classic distributed data parallel training.
• FullyShardedDataParallel (FSDP1): Use the original FSDP wrapper for parameter sharding across data-parallel workers.

Reference: references/pytorch_ddp_notes.md, references/pytorch_fsdp1_api.md.

---

Contract the agent must follow

1. Launch with torchrun and set the CUDA device per process (usually via LOCAL_RANK).
2. Apply fully_shard() bottom-up, i.e., shard submodules (e.g., Transformer blocks) before the root module.
3. Call model(input), not model.forward(input), so the FSDP2 hooks run (unless you explicitly unshard() or register the forward method).
4. Create the optimizer after sharding and make sure it is built on the DTensor parameters (post-fully_shard).
5. Checkpoint using Distributed Checkpoint (DCP) or the distributed-state-dict helpers, not naïve torch.save(model.state_dict()) unless you deliberately gather to full tensors.