Pytorch Fsdp2

This teaches Claude to correctly implement PyTorch's FSDP2 (the DTensor-based sharding API via fully_shard) in training scripts. It covers the full setup: launching with torchrun, applying bottom-up sharding to submodules before the root, creating optimizers after sharding so they see DTensor parameters, and checkpointing with Distributed Checkpoint or the state dict helpers. The skill is opinionated about the right patterns (use model() not model.forward(), never naively torch.save DTensor state dicts) and includes concrete migration steps from meta device initialization through mixed precision config. Use this when your model won't fit on one GPU or you want cleaner per-parameter sharding than FSDP1.

skills.sh installs: 1K
GitHub stars: 14.4K
last updated: Aug 21, 2026

View on GitHub

Install

$ npx -y skills add orchestra-research/ai-research-skills --skill pytorch-fsdp2 --agent claude-code

Installs into .claude/skills of the current project.

Files

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.