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.
npxskills add orchestra-research/ai-research-skills--skill pytorch-fsdp2Loading…
fully_shard) correctly in a training scriptThis 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_shardand theFSDPModulemethods it adds in-place to modules. See:references/pytorch_fully_shard_api.md,references/pytorch_fsdp2_tutorial.md.
Use FSDP2 when:
Avoid (or be careful) if:
Reference: references/pytorch_ddp_notes.md, references/pytorch_fsdp1_api.md.
torchrun and set the CUDA device per process (usually via LOCAL_RANK).fully_shard() bottom-up, i.e., shard submodules (e.g., Transformer blocks) before the root module.model(input), not model.forward(input), so the FSDP2 hooks run (unless you explicitly unshard() or register the forward method).fully_shard).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.)
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).
Minimal, correct pattern:
dist.init_process_group(backend="nccl")torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))DeviceMesh to describe the data-parallel group(s)Reference: references/pytorch_device_mesh_tutorial.md (why DeviceMesh exists & how it manages process groups).
For big models, initialize on meta, apply sharding, then materialize weights on GPU:
with torch.device("meta"): model = ...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).
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:
if isinstance(m, TransformerBlock): fully_shard(m, ...)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).
reshard_after_forward for memory/perf trade-offsDefault behavior:
None means True for non-root modules and False for root modules (good default).Heuristics:
True on many blocks.False).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).
FSDP2 uses:
mp_policy=MixedPrecisionPolicy(param_dtype=..., reduce_dtype=..., output_dtype=..., cast_forward_inputs=...)offload_policy=CPUOffloadPolicy() if you want CPU offloadRules of thumb:
reduce_dtype aligned with your gradient reduction expectations.Reference: references/pytorch_fully_shard_api.md (MixedPrecisionPolicy / OffloadPolicy classes).
set_requires_gradient_sync) instead of FSDP1’s no_sync().Gradient clipping:
Reference: references/pytorch_fsdp2_tutorial.md.
Two recommended approaches:
A) Distributed Checkpoint (DCP) — best default
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, ...)get_optimizer_state_dict / set_optimizer_state_dictAvoid:
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)torchrun and initialize the process group.LOCAL_RANK; create a DeviceMesh if you need multi-dim parallelism.meta if needed), apply fully_shard bottom-up, then fully_shard(model).model(inputs) so hooks run; use set_requires_gradient_sync for accumulation.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.
Stateful or assemble state via get_state_dict.dcp.save(...) from all ranks to a shared path.dcp.load(...) and restore with set_state_dict.Reference: references/pytorch_dcp_recipe.md.
torch.cuda.set_device(LOCAL_RANK) and your torchrun flags.forward() directly?model(input) or explicitly unshard() / register forward.fully_shard() applied bottom-up?torch.save unless you understand conversions.model(inputs) (or unshard() explicitly) instead of model.forward(...).fully_shard calls.fully_shard bottom-up on submodules before the root.reshard_after_forward=True for more modules.set_requires_gradient_sync instead of FSDP1’s no_sync().Reference: references/pytorch_fully_shard_api.md, references/pytorch_fsdp2_tutorial.md.
The coding agent should implement a script with these labeled blocks:
init_distributed(): init process group, set devicebuild_model_meta(): model on meta, apply fully_shard, materialize weightsbuild_optimizer(): optimizer created after shardingtrain_step(): forward/backward/step with model(inputs) and DTensor-aware patternscheckpoint_save/load(): DCP or distributed state dict helpersConcrete examples live in references/pytorch_examples_fsdp2.md and the official tutorial reference.
references/pytorch_fsdp2_tutorial.mdreferences/pytorch_fully_shard_api.mdreferences/pytorch_ddp_notes.mdreferences/pytorch_fsdp1_api.mdreferences/pytorch_device_mesh_tutorial.mdreferences/pytorch_tp_tutorial.mdreferences/pytorch_dcp_overview.mdreferences/pytorch_dcp_recipe.mdreferences/pytorch_dcp_async_recipe.mdreferences/pytorch_examples_fsdp2.mdreferences/torchtitan_fsdp_notes.md (optional, production notes)references/ray_train_fsdp2_example.md (optional, integration example)
Create or update AgentSkills. Use when designing, structuring, or packaging skills with scripts, references, and assets.
Create or update AgentSkills. Use when designing, structuring, or packaging skills with scripts, references, and assets.
Set up and use 1Password CLI (op). Use when installing the CLI, enabling desktop app integration, signing in (single or multi-account), or reading/injecting/running secrets via op.
CLI to manage emails via IMAP/SMTP. Use `himalaya` to list, read, write, reply, forward, search, and organize emails from the terminal. Supports multiple accounts and message composition with MML (MIME Meta Language).
Create or update AgentSkills. Use when designing, structuring, or packaging skills with scripts, references, and assets.
Create or update AgentSkills. Use when designing, structuring, or packaging skills with scripts, references, and assets.
Set up and use 1Password CLI (op). Use when installing the CLI, enabling desktop app integration, signing in (single or multi-account), or reading/injecting/running secrets via op.
CLI to manage emails via IMAP/SMTP. Use `himalaya` to list, read, write, reply, forward, search, and organize emails from the terminal. Supports multiple accounts and message composition with MML (MIME Meta Language).