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)