Pytorch Fsdp2 Agent 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.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

• Launch with torchrun and set the CUDA device per process (usually via LOCAL_RANK).
• Apply fully_shard() bottom-up, i.e., shard submodules (e.g., Transformer blocks) before the root module.
• Call model(input), not model.forward(input), so the FSDP2 hooks run (unless you explicitly unshard() or register the forward method).
• Create the optimizer after sharding and make sure it is built on the DTensor parameters (post-fully_shard).
• 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 --nproc_per_node <gpus_per_node> ... and ensure RANK, WORLD_SIZE, LOCAL_RANK are visible.

Reference: references/pytorch_fsdp2_tutorial.md (launch commands and setup), references/pytorch_fully_shard_api.md (user contract).

1) Initialize distributed and set device

Minimal, correct pattern:

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

Reference: references/pytorch_device_mesh_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 fully_shard(...) on submodules, then fully_shard(model)
• model.to_empty(device="cuda")
• model.reset_parameters() (or your init routine)

Reference: references/pytorch_fsdp2_tutorial.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/pytorch_fully_shard_api.md (bottom-up requirement and why).

4) Configure reshard_after_forward 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/pytorch_fully_shard_api.md (full semantics).

5) Mixed precision & offload (optional but common)

FSDP2 uses:

• mp_policy=MixedPrecisionPolicy(param_dtype=..., reduce_dtype=..., output_dtype=..., cast_forward_inputs=...)
• 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/pytorch_fully_shard_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 (set_requires_gradient_sync) instead of FSDP1’s no_sync().

Gradient clipping:

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

Reference: references/pytorch_fsdp2_tutorial.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

• get_model_state_dict / set_model_state_dict with StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True, ...)
• For optimizer: get_optimizer_state_dict / set_optimizer_state_dict

Avoid:

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

References:

• references/pytorch_dcp_overview.md (DCP behavior and caveats)
• references/pytorch_dcp_recipe.md and references/pytorch_dcp_async_recipe.md (end-to-end usage)
• references/pytorch_fsdp2_tutorial.md (DTensor vs DCP state-dict flows)
• references/pytorch_examples_fsdp2.md (working checkpoint scripts)

Workflow checklists (copy-paste friendly)

Workflow A: Retrofit FSDP2 into an existing training script

• [ ] Launch with torchrun and initialize the process group.
• [ ] Set the CUDA device from LOCAL_RANK; create a DeviceMesh if you need multi-dim parallelism.
• [ ] Build the model (use meta if needed), apply fully_shard bottom-up, then fully_shard(model).
• [ ] Create the optimizer after sharding so it captures DTensor parameters.
• [ ] Use model(inputs) so hooks run; use set_requires_gradient_sync for accumulation.
• [ ] Add DCP save/load via torch.distributed.checkpoint helpers.

Reference: references/pytorch_fsdp2_tutorial.md, references/pytorch_fully_shard_api.md, references/pytorch_device_mesh_tutorial.md, references/pytorch_dcp_recipe.md.

Workflow B: Add DCP save/load (minimal pattern)

• [ ] Wrap state in Stateful or assemble state via get_state_dict.
• [ ] Call dcp.save(...) from all ranks to a shared path.
• [ ] Call dcp.load(...) and restore with set_state_dict.
• [ ] Validate any resharding assumptions when loading into a different mesh.

Reference: references/pytorch_dcp_recipe.md.

Debug checklist (what the agent should check first)

• All ranks on distinct GPUs?

If not, verify torch.cuda.set_device(LOCAL_RANK) and your torchrun flags.

• Did you accidentally call forward() directly?

Use model(input) or explicitly unshard() / register forward.

• Is fully_shard() applied bottom-up?

If only root is sharded, expect worse memory/perf and possible confusion.

• Optimizer created at the right time?

Must be built on DTensor parameters *after* sharding.

• Checkpointing path consistent?
• If using DCP, don’t mix with ad-hoc torch.save unless you understand conversions.
• Be mindful of PyTorch-version compatibility warnings for DCP.

Common issues and fixes

• Forward hooks not running → Call model(inputs) (or unshard() explicitly) instead of model.forward(...).
• Optimizer sees non-DTensor params → Create optimizer after all fully_shard calls.
• Only root module sharded → Apply fully_shard bottom-up on submodules before the root.
• Memory spikes after forward → Set reshard_after_forward=True for more modules.
• Gradient accumulation desync → Use set_requires_gradient_sync instead of FSDP1’s no_sync().

Reference: references/pytorch_fully_shard_api.md, references/pytorch_fsdp2_tutorial.md.

Minimal reference implementation outline (agent-friendly)

The coding agent should implement a script with these labeled blocks:

• init_distributed(): init process group, set device
• build_model_meta(): model on meta, apply fully_shard, materialize weights
• build_optimizer(): optimizer created after sharding
• train_step(): forward/backward/step with model(inputs) and DTensor-aware patterns
• checkpoint_save/load(): DCP or distributed state dict helpers

Concrete examples live in references/pytorch_examples_fsdp2.md and the official tutorial reference.

References

• references/pytorch_fsdp2_tutorial.md
• references/pytorch_fully_shard_api.md
• references/pytorch_ddp_notes.md
• references/pytorch_fsdp1_api.md
• references/pytorch_device_mesh_tutorial.md
• references/pytorch_tp_tutorial.md
• references/pytorch_dcp_overview.md
• references/pytorch_dcp_recipe.md
• references/pytorch_dcp_async_recipe.md
• references/pytorch_examples_fsdp2.md
• references/torchtitan_fsdp_notes.md (optional, production notes)
• references/ray_train_fsdp2_example.md (optional, integration example)