PyTorch FSDP2 Training Sharding
Guides an agent to add PyTorch FSDP2 sharding, mixed precision, offload and distributed checkpointing to training scripts.
- Repository
- Orchestra-Research/AI-Research-SKILLs
- Stars
- 14,195
- Licence
- MIT
- Files
- 13 files
What it does
When installed, the agent knows how to modify a PyTorch training script so a large model is split (sharded) across multiple GPUs instead of needing to fit on one. It follows a set of rules: launch with torchrun, shard submodules before the root, create the optimizer after sharding, and save checkpoints with PyTorch's Distributed Checkpoint tools rather than a plain torch.save. The skill also points to reference documents on mixed precision, CPU offload and device meshes.
Great for
- ML engineers training models too large for a single GPU
- Developers migrating a training loop from FSDP1 or DDP to FSDP2
- Teams combining data parallelism with tensor parallelism via DeviceMesh
Skip it if
- Small models that fit on one GPU, where DDP or plain training is simpler and avoids added complexity
- Projects needing checkpoints compatible across many PyTorch versions, since the skill itself warns DCP is not strictly backward compatible
- Users without multi-GPU hardware or a recent PyTorch install, as the code cannot run without them
Try asking
- “Add FSDP2 fully_shard to my transformer training script so it fits on 4 GPUs”
- “Set up bf16 mixed precision and CPU offload with FSDP2 for my model”
- “Switch my checkpoint saving to Distributed Checkpoint so each rank writes in parallel”
- “Migrate my FSDP1 training loop to FSDP2 using DTensor and DeviceMesh”
Install in 3 steps
- 1
Download the skill
The ZIP contains the 13 original files, the MIT licence and an attribution note.
pytorch-fsdp2.zip - 2
Add it to your agent
Claude Code (all projects)
mkdir -p ~/.claude/skills && unzip pytorch-fsdp2.zip -d ~/.claude/skills/Claude Code / Agent SDK (this project only)
mkdir -p .claude/skills && unzip pytorch-fsdp2.zip -d .claude/skills/Claude apps (web/desktop): Settings → Capabilities → Skills → Upload skill, and choose the ZIP.
- 3
Ask for the task
No special command needed — the agent reads the skill's description and loads it when your request matches. Try one of the prompts above.
Get better results
- Give the agent your model code and GPU count so it can pick the right sharding granularity and launch command.
- Verify the edited script on a small number of GPUs before launching a long run, and check that the optimizer uses DTensor parameters.
- Review the agent's changes to checkpoint code, since mistakes there can silently lose training state.
What's inside
- references/pytorch_dcp_async_recipe.md
- references/pytorch_dcp_overview.md
- references/pytorch_dcp_recipe.md
- references/pytorch_ddp_notes.md
- references/pytorch_device_mesh_tutorial.md
- references/pytorch_examples_fsdp2.md
- references/pytorch_fsdp1_api.md
- references/pytorch_fsdp2_tutorial.md
- references/pytorch_fully_shard_api.md
- references/pytorch_tp_tutorial.md
- references/ray_train_fsdp2_example.md
- references/torchtitan_fsdp_notes.md
- SKILL.md
Redistributed unmodified from Orchestra-Research/AI-Research-SKILLs under the MIT licence. The assessment was drafted with AI from the skill's files on 2026-10-29. Skills can instruct your agent to run code — read SKILL.md and any scripts before installing.