pytorch-fsdp2
Adds PyTorch FSDP2 (fully_shard) to training scripts with correct init, sharding, mixed precision/offload config, and distributed checkpointing. Use when models exceed single-GPU memory or when you need DTensor-based sharding with DeviceMesh.
npx skills add Orchestra-Research/AI-Research-SKILLs --skill pytorch-fsdp2 --agent claude-code
Same command for any agent — swap --agent for codex, cursor, copilot.
Weekly change comes from our own snapshots, not the repository page — it measures attention, not adoption.
# 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 parall
- When to use this skill
- Alternatives (when FSDP2 is not the best fit)
- Contract the agent must follow
- Step-by-step procedure
- 0) Version & environment sanity
- 1) Initialize distributed and set device
- 2) Build model on meta device (recommended for very large models)
- 3) Apply fullyshard() bottom-up (wrapping policy = “apply where needed”)
- 4) Configure reshardafterforward for memory/perf trade-offs
- 5) Mixed precision & offload (optional but common)
- 6) Optimizer, gradient clipping, accumulation
- 7) Checkpointing: prefer DCP or distributed state dict helpers
- Workflow checklists (copy-paste friendly)
- Workflow A: Retrofit FSDP2 into an existing training script
What does the pytorch-fsdp2 skill do?
Adds PyTorch FSDP2 (fully_shard) to training scripts with correct init, sharding, mixed precision/offload config, and distributed checkpointing. Use when models exceed single-GPU memory or when you need DTensor-based sharding with DeviceMesh.
How do I install it?
Run `npx skills add Orchestra-Research/AI-Research-SKILLs --skill pytorch-fsdp2 --agent claude-code` — it drops the skill into your project so the agent can pick it up. Swap the --agent value for codex, cursor or copilot if you use one of those.
Where does this skill come from?
From Orchestra-Research/AI-Research-SKILLs, a repository with 11,391 stars. We read it straight from the repository tree rather than a submitted listing, so what you see here is what is actually published.
Is a popular skill a good skill?
Not necessarily. Stars measure attention, not adoption — a repository can trend for a week and be abandoned. That is why we show the weekly change from our own snapshots next to the total, instead of a single flattering number.
