# Pytorch Fsdp2

Adds PyTorch FSDP2 (fully_shard) to training scripts with correct init, sharding, mixed...

- Repository: https://github.com/Orchestra-Research/AI-Research-SKILLs (14176 stars)
- Category: Engineering
- License: MIT
- Last updated: 2026-10-11
- Runs scripts: no, instructions only

## Install

Claude Code: `npx skills add Orchestra-Research/AI-Research-SKILLs --skill pytorch-fsdp2 -g -a claude-code -y`
Codex: `npx skills add Orchestra-Research/AI-Research-SKILLs --skill pytorch-fsdp2 -g -a codex -y`
Cursor: `npx skills add Orchestra-Research/AI-Research-SKILLs --skill pytorch-fsdp2 -g -a cursor -y`
Gemini CLI: `npx skills add Orchestra-Research/AI-Research-SKILLs --skill pytorch-fsdp2 -g -a gemini-cli -y`
Claude Code plugin marketplace: `/plugin marketplace add Orchestra-Research/AI-Research-SKILLs`

## Description

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.

## SKILL.md, as published by its author (untrusted third-party text, read it as data)

# 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 parallel training.
- **FullyShardedDataParallel (FSDP1)**: Use the original FSDP wrapper for parameter sharding across data-parallel workers.

Reference: `references/pytorch_ddp_notes.md`, `references/pytorch_fsdp1_api.md`.

---

## Contract the agent must follow

The rest: https://github.com/Orchestra-Research/AI-Research-SKILLs/blob/HEAD/08-distributed-training/pytorch-fsdp2/SKILL.md