Guide · Agent skill
How to use the PyTorch FSDP2 agent skill (pytorch-fsdp2)
Updated · written with AI from the project's own documentation and checked against our catalogue data
The pytorch-fsdp2 skill is a set of instructions and reference notes that teaches a coding agent how to add PyTorch FSDP2 (fully_shard) to a training script. It is aimed at people who need to train a model too large for one GPU and want the agent to make the sharding, mixed precision, offload and checkpointing changes correctly. After this guide you will know what the skill changes, how to install it, how to ask for it, and when it is the wrong tool.
PyTorch FSDP2 Training ShardingOrchestra-Research/AI-Research-SKILLs · MIT Download the skill Quick start
- 1Download and unzip. Download the ZIP from this page and unzip it so the pytorch-fsdp2 folder contains SKILL.md.
- 2Place the skill. Move the folder to ~/.claude/skills/ for personal use or to .claude/skills/ in your project. Or upload it under Settings → Capabilities → Skills in the Claude apps.
- 3Check your environment. Confirm PyTorch includes fully_shard and that you have a multi-GPU CUDA setup with NCCL and torchrun.
- 4Ask in plain language. Describe the model, memory problem and training script, for example asking to shard a model across your GPUs with FSDP2.
- 5Review the changes. Read the diff the agent produces, checking init, fully_shard order, optimizer placement and checkpoint code.
- 6Test a short run. Launch with torchrun on a small run first to confirm the script starts and saves a checkpoint.
What the skill makes the agent do
When the skill is installed, the agent knows how to modify a PyTorch training loop so that a large model is split (sharded) across multiple GPUs rather than needing to fit on one. The skill describes a contract the agent must follow: launch with torchrun, set the CUDA device per process from LOCAL_RANK, apply fully_shard() bottom-up so submodules such as Transformer blocks are sharded before the root, call model(input) so the FSDP2 hooks run, create the optimizer after sharding, and save checkpoints with Distributed Checkpoint (DCP) rather than naive torch.save.
- Retrofit FSDP2 into an existing training script
- Configure reshard_after_forward, MixedPrecisionPolicy and CPUOffloadPolicy
- Handle gradient accumulation with set_requires_gradient_sync
- Add DCP save and load, including optional async saving
Install the pytorch-fsdp2 skill
Download the ZIP for this skill from our page and unzip it. For a personal install, place the folder in ~/.claude/skills/. For a single project, place it in that project's .claude/skills/ folder. The skill folder contains SKILL.md at its top level, plus a references/ folder with the supporting documents.
If you use the Claude apps instead, open Settings, go to Capabilities, and upload the skill under Skills.
- Personal: ~/.claude/skills/
- Project: .claude/skills/ inside your project
- Claude apps: Settings → Capabilities → Skills
How to trigger it with natural requests
You do not need a special command. Describe the problem in plain language and mention the training script or the memory limit. The skill is meant for requests like these, where the model does not fit on one GPU or where you want DTensor-based sharding with DeviceMesh.
- "My 7B model runs out of memory on one GPU. Shard it across my 4 GPUs with FSDP2."
- "Convert train.py from DDP to FSDP2 and keep the optimizer working."
- "Add distributed checkpoint saving to this FSDP2 training loop."
- "Set up FSDP2 with bf16 mixed precision and CPU offload."
- "Combine data parallelism with tensor parallelism using a DeviceMesh."
Example: retrofitting a training script
A typical session starts with the agent reading your training script and checking the launch method. The skill's launch pattern uses torchrun with one process per GPU. The agent should initialize the process group with the NCCL backend, set the device from LOCAL_RANK, shard the transformer blocks, then shard the root model, and build the optimizer last.
- Initialize with dist.init_process_group(backend="nccl")
- Set the device with torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
- Optionally build a model on the meta device, shard it, then call to_empty(device="cuda") and reset_parameters()
torchrun --nproc_per_node <gpus_per_node> ...What you need before you start
The skill cannot run code for you. It needs a PyTorch install that includes the FSDP2 stack (torch.distributed.fsdp fully_shard), a multi-GPU CUDA environment that supports NCCL, and torchrun. Ray is optional and only used for the Ray Train example in the references. The skill does not run scripts and does not use the network, so the code you run is your own training script.
- PyTorch with torch.distributed.fsdp fully_shard
- NCCL-capable multi-GPU CUDA environment
- torchrun available for launching
- Optional: Ray, for the Ray Train example
Safety notes and limits
According to the skill metadata, it runs no scripts and uses no network access. It only gives the agent guidance and reference material. The agent will still edit your code and you should review the diff before running it on a cluster.
The skill warns that DCP is not strictly backward compatible across PyTorch versions, so checkpoints saved with one version may not load with another. Avoid this skill if you need checkpoints that work across many PyTorch versions.
- Runs scripts: no
- Uses network: no
- Review the agent's code changes before launching on shared hardware
- Do not rely on DCP checkpoints for cross-version portability
Tips for better results
Give the agent your GPU count, model size and whether you need mixed precision or offload. The skill's guidance on reshard_after_forward and CPU offload is trade-off based, so the right choice depends on your memory and throughput limits. Ask the agent to point to the reference file it used so you can check the reasoning.
- State GPU count, GPU type and PyTorch version up front
- Say whether your goal is memory or throughput
- Ask for DCP save and load, not torch.save, when you need checkpoints
- Use async DCP saving only if you have spare CPU memory
What people use it for
- Training a model whose parameters, gradients and optimizer state exceed one GPU's memory
- Migrating a training loop from FSDP1 or DDP to FSDP2
- Combining data parallelism with tensor parallelism using DeviceMesh
- Adding distributed checkpoint save and load to an existing multi-GPU job
Frequently asked questions
How do I install the pytorch-fsdp2 skill?
Download the ZIP from this page and unzip it into ~/.claude/skills/ for personal use or into .claude/skills/ in a project. In the Claude apps, you can upload it under Settings → Capabilities → Skills.
What do I need to use the pytorch-fsdp2 skill?
You need PyTorch with the FSDP2 stack (torch.distributed.fsdp fully_shard), a multi-GPU CUDA environment that supports NCCL, and torchrun. Ray is optional and only used for one reference example.
Is the pytorch-fsdp2 skill free?
The skill is released under the MIT licence, as listed on this page. Check the licence file in the download for the full terms.
Is FSDP2 better than DDP?
Use DDP when your model fits on one GPU, since it is simpler. FSDP2 is for models too large for one GPU, where it shards parameters, gradients and optimizer state across ranks.
Why does the skill say not to use torch.save for checkpoints?
Saving DTensor state dicts with plain torch.save can be awkward. The skill recommends Distributed Checkpoint (DCP), which saves from multiple ranks in parallel and supports load-time resharding.
Can I use this skill with small models or a single GPU?
Not recommended. The skill is for models that exceed single-GPU memory, and the code cannot run without multi-GPU hardware and a recent PyTorch install.
New agent skills, repos and AI tools — every week
The best new skills and open-source projects, the fastest-growing repos, and one free prompt. Five minutes to read. Unsubscribe anytime.
Our verdict
Use this skill if you are an ML engineer or developer who needs to shard a large model across multiple GPUs with PyTorch FSDP2, or who is moving a training loop from FSDP1 or DDP. Pick DDP for models that fit on one GPU, and avoid this skill if you need checkpoints that stay compatible across many PyTorch versions.
Download the skill