pytorch-fsdp2

1# Skill: Use PyTorch FSDP2 (`fully_shard`) correctly in a training script
2 
3This skill teaches a coding agent how to **add PyTorch FSDP2** to a training loop with correct initialization, sharding, mixed precision/offload configuration, and checkpointing.
4 
5> 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`.
6 
7---
8 
9## When to use this skill
10 
11Use FSDP2 when:
12- Your model **doesn’t fit** on one GPU (parameters + gradients + optimizer state).
13- You want an eager-mode sharding approach that is **DTensor-based per-parameter sharding** (more inspectable, simpler sharded state dicts) than FSDP1.  
14- You may later compose DP with **Tensor Parallel** using **DeviceMesh**.
15 
16Avoid (or be careful) if:
17- You need strict backwards-compatible checkpoints across PyTorch versions (DCP warns against this).
18- You’re forced onto older PyTorch versions without the FSDP2 stack.
19 
20## Alternatives (when FSDP2 is not the best fit)
21 
22- **DistributedDataParallel (DDP)**: Use the standard data-parallel wrapper when you want classic distributed data parallel training.
23- **FullyShardedDataParallel (FSDP1)**: Use the original FSDP wrapper for parameter sharding across data-parallel workers.
24 
25Reference: `references/pytorch_ddp_notes.md`, `references/pytorch_fsdp1_api.md`.
26 
27---
28 
29## Contract the agent must follow
30…