Loading
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 ne
The skill file stays in the publisher's repository.
Paste the copied command in a terminal. It installs this SKILL.md into that agent.
npx skills add Orchestra-Research/AI-Research-SKILLs --skill pytorch-fsdp2 -a cursor -y
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.
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.
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)