Agent skill · software engineering · zechenzhangagi
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.
Why this skill is useful
Adds executable scripts for initializing and configuring PyTorch FSDP2 with mixed precision and distributed checkpointing that aren't in public docs.
What it needs
Requires torch installed locally. About 5k tokens when loaded. Last updated 2026-06-16. 11,472 stars on the source repository.
What this skill does
Skill: Use PyTorch FSDP2 (fullyshard) 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.fullyshard and the FSDPModule methods it adds in-place to modules. See: references/pytorchfullyshardapi.md, references/pytorchfsdp2tutorial.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/pytorchddpnotes.md, references/pytorchfsdp1api.md. --- Contract the agent must follow 1. Launch with torchrun and set the CUDA device per process (usually via LOCALRANK). 2. Apply fullyshard() bottom-up, i.e., shard submodules (e.g., Transformer blocks) before the root module. 3. Call model(input), not model.forward(input), so the FSDP2 hooks run (unless you explicitly unshard() or register the forward method). 4. Create the optimizer after sharding and make sure it is built on the DTensor parameters (post-fullyshard). 5. Checkpoint using Distributed Checkpoint (DCP) or the distributed-state-dict helpers, not naïve torch.save(model.statedict()) unless you deliberately gather to full tensors. …
How to use it
Reference it in AdaL, Claude Code, Cursor or any coding agent — nothing to install:
@skills zechenzhangagi/pytorch-fsdp2