pytorch-fsdp2 skill

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.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 LOCAL_RANK).
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-fully_shard).
5. Checkpoint using Distributed Checkpoint (DCP) or the distributed-state-dict helpers, not naïve torch.save(model.state_dict()) unless you deliberately gather to full tensors.

(Each of these rules is directly described in the official API docs/tutorial; see references.)

Step-by-step procedure

0) Version & environment sanity

• Prefer a recent stable PyTorch where the docs show FSDP2 and DCP updated recently.
• Use torchrun --nprocpernode ... and ensure RANK, WORLDSIZE, LOCALRANK are visible.

Reference: references/pytorchfsdp2tutorial.md (launch commands and setup), references/pytorchfullyshard_api.md (user contract).

1) Initialize distributed and set device

Minimal, correct pattern:

• dist.initprocessgroup(backend="nccl")
• torch.cuda.setdevice(int(os.environ["LOCALRANK"]))
• Optionally create a DeviceMesh to describe the data-parallel group(s)

Reference: references/pytorchdevicemesh_tutorial.md (why DeviceMesh exists & how it manages process groups).

2) Build model on meta device (recommended for very large models)

For big models, initialize on meta, apply sharding, then materialize weights on GPU:

• with torch.device("meta"): model = ...
• apply fullyshard(...) on submodules, then fullyshard(model)
• model.to_empty(device="cuda")
• model.reset_parameters() (or your init routine)

Reference: references/pytorchfsdp2tutorial.md (migration guide shows this flow explicitly).

3) Apply fully_shard() bottom-up (wrapping policy = “apply where needed”)

Do not only call fully_shard on the topmost module.

Recommended sharding pattern for transformer-like models:

• iterate modules, if isinstance(m, TransformerBlock): fully_shard(m, ...)
• then fully_shard(model, ...)

Why:

• fully_shard forms “parameter groups” for collective efficiency and excludes params already grouped by earlier calls. Bottom-up gives better overlap and lower peak memory.

Reference: references/pytorchfullyshard_api.md (bottom-up requirement and why).

4) Configure reshardafterforward for memory/perf trade-offs

Default behavior:

• None means True for non-root modules and False for root modules (good default).

Heuristics:

• If you’re memory-bound: keep defaults or force True on many blocks.
• If you’re throughput-bound and can afford memory: consider keeping unsharded params longer (root often False).
• Advanced: use an int to reshard to a smaller mesh after forward (e.g., intra-node) if it’s a meaningful divisor.

Reference: references/pytorchfullyshard_api.md (full semantics).

5) Mixed precision & offload (optional but common)

FSDP2 uses:

• mppolicy=MixedPrecisionPolicy(paramdtype=..., reducedtype=..., outputdtype=..., castforwardinputs=...)
• offload_policy=CPUOffloadPolicy() if you want CPU offload

Rules of thumb:

• Start with BF16 parameters/reductions on H100/A100-class GPUs (if numerically stable for your model).
• Keep reduce_dtype aligned with your gradient reduction expectations.
• If you use CPU offload, budget for PCIe/NVLink traffic and runtime overhead.

Reference: references/pytorchfullyshard_api.md (MixedPrecisionPolicy / OffloadPolicy classes).

6) Optimizer, gradient clipping, accumulation

• Create the optimizer after sharding so it holds DTensor params.
• If you need gradient accumulation / no_sync:
• use the FSDP2 mechanism (setrequiresgradientsync) instead of FSDP1’s nosync().

Gradient clipping:

• Use the approach shown in the FSDP2 tutorial (“Gradient Clipping and Optimizer with DTensor”), because parameters/gradients are DTensors.

Reference: references/pytorchfsdp2tutorial.md.

7) Checkpointing: prefer DCP or distributed state dict helpers

Two recommended approaches:

A) Distributed Checkpoint (DCP) — best default

• DCP saves/loads from multiple ranks in parallel and supports load-time resharding.
• DCP produces multiple files (often at least one per rank) and operates “in place”.

B) Distributed state dict helpers

• getmodelstatedict / setmodelstatedict with StateDictOptions(fullstatedict=True, cpuoffload=True, broadcastfrom_rank0=True, ...)
• For optimizer: getoptimizerstatedict / setoptimizerstatedict

Avoid:

• Saving DTensor state dicts with plain torch.save unless you intentionally convert with DTensor.full_tensor() and manage memory carefully.

References:

• references/pytorchdcpoverview.md (DCP behavior and caveats)
• references/pytorchdcprecipe.md and references/pytorchdcpasync_recipe.md (end-to-end usage)
• references/pytorchfsdp2tutorial.md (DTensor vs DCP state-dict flows)
• references/pytorchexamplesfsdp2.md (working checkpoint scripts)