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

1. Launch with torchrun and set the CUDA device per process (usually via LOCAL_RANK).
2. Apply fully_shard() 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 --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)

1. All ranks on distinct GPUs?
   If not, verify torch.cuda.set_device(LOCAL_RANK) and your torchrun flags.
2. Did you accidentally call forward() directly?
   Use model(input) or explicitly unshard() / register forward.
3. Is fully_shard() applied bottom-up?
   If only root is sharded, expect worse memory/perf and possible confusion.
4. Optimizer created at the right time?
   Must be built on DTensor parameters after sharding.
5. 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)