{"id":"pytorch-fsdp2","name":"pytorch-fsdp2","summary":"正しいinit、シャーディング、混合精度/オフロード設定、分散チェックポイントを備えたトレーニングスクリプトにPyTorch FSDP2(fully_shard)を追加します。","body":"# Skill: Use PyTorch FSDP2 (`fully_shard`) correctly in a training script\n\nThis skill teaches a coding agent how to **add PyTorch FSDP2** to a training loop with correct initialization, sharding, mixed precision/offload configuration, and checkpointing.\n\n> 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`.\n\n---\n\n## When to use this skill\n\nUse FSDP2 when:\n- Your model **doesn’t fit** on one GPU (parameters + gradients + optimizer state).\n- You want an eager-mode sharding approach that is **DTensor-based per-parameter sharding** (more inspectable, simpler sharded state dicts) than FSDP1.  \n- You may later compose DP with **Tensor Parallel** using **DeviceMesh**.\n\nAvoid (or be careful) if:\n- You need strict backwards-compatible checkpoints across PyTorch versions (DCP warns against this).\n- You’re forced onto older PyTorch versions without the FSDP2 stack.\n\n## Alternatives (when FSDP2 is not the best fit)\n\n- **DistributedDataParallel (DDP)**: Use the standard data-parallel wrapper when you want classic distributed data parallel training.\n- **FullyShardedDataParallel (FSDP1)**: Use the original FSDP wrapper for parameter sharding across data-parallel workers.\n\nReference: `references/pytorch_ddp_notes.md`, `references/pytorch_fsdp1_api.md`.\n\n---\n\n## Contract the agent must follow\n\n1. **Launch with `torchrun`** and set the CUDA device per process (usually via `LOCAL_RANK`).  \n2. **Apply `fully_shard()` bottom-up**, i.e., shard submodules (e.g., Transformer blocks) before the root module.  \n3. **Call `model(input)`**, not `model.forward(input)`, so the FSDP2 hooks run (unless you explicitly `unshard()` or register the forward method).  \n4. **Create the optimizer after sharding** and make sure it is built on the **DTensor parameters** (post-`fully_shard`).  \n5. **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.\n\n(Each of these rules is directly described in the official API docs/tutorial; see references.)\n\n---\n\n## Step-by-step procedure\n\n### 0) Version & environment sanity\n- Prefer a recent stable PyTorch where the docs show FSDP2 and DCP updated recently.\n- Use `torchrun --nproc_per_node <gpus_per_node> ...` and ensure `RANK`, `WORLD_SIZE`, `LOCAL_RANK` are visible.\n\nReference: `references/pytorch_fsdp2_tutorial.md` (launch commands and setup), `references/pytorch_fully_shard_api.md` (user contract).\n\n---\n\n### 1) Initialize distributed and set device\nMinimal, correct pattern:\n- `dist.init_process_group(backend=\"nccl\")`\n- `torch.cuda.set_device(int(os.environ[\"LOCAL_RANK\"]))`\n- Optionally create a `DeviceMesh` to describe the data-parallel group(s)\n\nReference: `references/pytorch_device_mesh_tutorial.md` (why DeviceMesh exists & how it manages process groups).\n\n---\n\n### 2) Build model on meta device (recommended for very large models)\nFor big models, initialize on `meta`, apply sharding, then materialize weights on GPU:\n- `with torch.device(\"meta\"): model = ...`\n- apply `fully_shard(...)` on submodules, then `fully_shard(model)`\n- `model.to_empty(device=\"cuda\")`\n- `model.reset_parameters()` (or your init routine)\n\nReference: `references/pytorch_fsdp2_tutorial.md` (migration guide shows this flow explicitly).\n\n---\n\n### 3) Apply `fully_shard()` bottom-up (wrapping policy = “apply where needed”)\n**Do not** only call `fully_shard` on the topmost module.\n\nRecommended sharding pattern for transformer-like models:\n- iterate modules, `if isinstance(m, TransformerBlock): fully_shard(m, ...)`\n- then `fully_shard(model, ...)`\n\nWhy:\n- `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.\n\nReference: `references/pytorch_fully_shard_api.md` (bottom-up requirement and why).\n\n---\n\n### 4) Configure `reshard_after_forward` for memory/perf trade-offs\nDefault behavior:\n- `None` means `True` for non-root modules and `False` for root modules (good default).\n\nHeuristics:\n- If you’re memory-bound: keep defaults or force `True` on many blocks.\n- If you’re throughput-bound and can afford memory: consider keeping unsharded params longer (root often `False`).\n- Advanced: use an `int` to reshard to a smaller mesh after forward (e.g., intra-node) if it’s a meaningful divisor.\n\nReference: `references/pytorch_fully_shard_api.md` (full semantics).\n\n---\n\n### 5) Mixed precision & offload (optional but common)\nFSDP2 uses:\n- `mp_policy=MixedPrecisionPolicy(param_dtype=..., reduce_dtype=..., output_dtype=..., cast_forward_inputs=...)`\n- `offload_policy=CPUOffloadPolicy()` if you want CPU offload\n\nRules of thumb:\n- Start with BF16 parameters/reductions on H100/A100-class GPUs (if numerically stable for your model).\n- Keep `reduce_dtype` aligned with your gradient reduction expectations.\n- If you use CPU offload, budget for PCIe/NVLink traffic and runtime overhead.\n\nReference: `references/pytorch_fully_shard_api.md` (MixedPrecisionPolicy / OffloadPolicy classes).\n\n---\n\n### 6) Optimizer, gradient clipping, accumulation\n- Create the optimizer **after** sharding so it holds DTensor params.\n- If you need gradient accumulation / no_sync:\n  - use the FSDP2 mechanism (`set_requires_gradient_sync`) instead of FSDP1’s `no_sync()`.\n\nGradient clipping:\n- Use the approach shown in the FSDP2 tutorial (“Gradient Clipping and Optimizer with DTensor”), because parameters/gradients are DTensors.\n\nReference: `references/pytorch_fsdp2_tutorial.md`.\n\n---\n\n### 7) Checkpointing: prefer DCP or distributed state dict helpers\nTwo recommended approaches:\n\n**A) Distributed Checkpoint (DCP) — best default**\n- DCP saves/loads from multiple ranks in parallel and supports load-time resharding.\n- DCP produces **multiple files** (often at least one per rank) and operates “in place”.\n\n**B) Distributed state dict helpers**\n- `get_model_state_dict` / `set_model_state_dict` with `StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True, ...)`\n- For optimizer: `get_optimizer_state_dict` / `set_optimizer_state_dict`\n\nAvoid:\n- Saving DTensor state dicts with plain `torch.save` unless you intentionally convert with `DTensor.full_tensor()` and manage memory carefully.\n\nReferences:\n- `references/pytorch_dcp_overview.md` (DCP behavior and caveats)\n- `references/pytorch_dcp_recipe.md` and `references/pytorch_dcp_async_recipe.md` (end-to-end usage)\n- `references/pytorch_fsdp2_tutorial.md` (DTensor vs DCP state-dict flows)\n- `references/pytorch_examples_fsdp2.md` (working checkpoint scripts)\n\n---\n\n## Workflow checklists (copy-paste friendly)\n\n### Workflow A: Retrofit FSDP2 into an existing training script\n- [ ] Launch with `torchrun` and initialize the process group.\n- [ ] Set the CUDA device from `LOCAL_RANK`; create a `DeviceMesh` if you need multi-dim parallelism.\n- [ ] Build the model (use `meta` if needed), apply `fully_shard` bottom-up, then `fully_shard(model)`.\n- [ ] Create the optimizer after sharding so it captures DTensor parameters.\n- [ ] Use `model(inputs)` so hooks run; use `set_requires_gradient_sync` for accumulation.\n- [ ] Add DCP save/load via `torch.distributed.checkpoint` helpers.\n\nReference: `references/pytorch_fsdp2_tutorial.md`, `references/pytorch_fully_shard_api.md`, `references/pytorch_device_mesh_tutorial.md`, `references/pytorch_dcp_recipe.md`.\n\n### Workflow B: Add DCP save/load (minimal pattern)\n- [ ] Wrap state in `Stateful` or assemble state via `get_state_dict`.\n- [ ] Call `dcp.save(...)` from all ranks to a shared path.\n- [ ] Call `dcp.load(...)` and restore with `set_state_dict`.\n- [ ] Validate any resharding assumptions when loading into a different mesh.\n\nReference: `references/pytorch_dcp_recipe.md`.\n\n## Debug checklist (what the agent should check first)\n\n1. **All ranks on distinct GPUs?**  \n   If not, verify `torch.cuda.set_device(LOCAL_RANK)` and your `torchrun` flags.\n2. **Did you accidentally call `forward()` directly?**  \n   Use `model(input)` or explicitly `unshard()` / register forward.\n3. **Is `fully_shard()` applied bottom-up?**  \n   If only root is sharded, expect worse memory/perf and possible confusion.\n4. **Optimizer created at the right time?**  \n   Must be built on DTensor parameters *after* sharding.\n5. **Checkpointing path consistent?**  \n   - If using DCP, don’t mix with ad-hoc `torch.save` unless you understand conversions.\n   - Be mindful of PyTorch-version compatibility warnings for DCP.\n\n---\n\n## Common issues and fixes\n\n- **Forward hooks not running** → Call `model(inputs)` (or `unshard()` explicitly) instead of `model.forward(...)`.\n- **Optimizer sees non-DTensor params** → Create optimizer after all `fully_shard` calls.\n- **Only root module sharded** → Apply `fully_shard` bottom-up on submodules before the root.\n- **Memory spikes after forward** → Set `reshard_after_forward=True` for more modules.\n- **Gradient accumulation desync** → Use `set_requires_gradient_sync` instead of FSDP1’s `no_sync()`.\n\nReference: `references/pytorch_fully_shard_api.md`, `references/pytorch_fsdp2_tutorial.md`.\n\n---\n\n## Minimal reference implementation outline (agent-friendly)\n\nThe coding agent should implement a script with these labeled blocks:\n\n- `init_distributed()`: init process group, set device\n- `build_model_meta()`: model on meta, apply `fully_shard`, materialize weights\n- `build_optimizer()`: optimizer created after sharding\n- `train_step()`: forward/backward/step with `model(inputs)` and DTensor-aware patterns\n- `checkpoint_save/load()`: DCP or distributed state dict helpers\n\nConcrete examples live in `references/pytorch_examples_fsdp2.md` and the official tutorial reference.\n\n---\n\n## References\n- `references/pytorch_fsdp2_tutorial.md`\n- `references/pytorch_fully_shard_api.md`\n- `references/pytorch_ddp_notes.md`\n- `references/pytorch_fsdp1_api.md`\n- `references/pytorch_device_mesh_tutorial.md`\n- `references/pytorch_tp_tutorial.md`\n- `references/pytorch_dcp_overview.md`\n- `references/pytorch_dcp_recipe.md`\n- `references/pytorch_dcp_async_recipe.md`\n- `references/pytorch_examples_fsdp2.md`\n- `references/torchtitan_fsdp_notes.md` (optional, production notes)\n- `references/ray_train_fsdp2_example.md` (optional, integration example)","author":"@Orchestra-Research","ownerProfile":null,"authorContacts":null,"sourceUrl":"https://github.com/Orchestra-Research/AI-Research-SKILLs/tree/main/08-distributed-training/pytorch-fsdp2","license":"MIT","category":null,"lang":"en","tokens":2544,"stars":0,"calls30d":1,"claimed":false,"visibility":"public","origin":"crawler","version":"0.1.0","createdAt":"2026-08-22","updatedAt":"2026-08-22","files":[{"path":"references/pytorch_dcp_async_recipe.md","size":808,"sha256":"1d59c50d255dca746e0b6a04286eba1ce1de55627c73230862b7038f1214fbfd"},{"path":"references/pytorch_dcp_overview.md","size":1022,"sha256":"81ae7c08701ad657c1fa2775a035631072775266b5387a8c52c74ca535be2567"},{"path":"references/pytorch_dcp_recipe.md","size":1043,"sha256":"94f7f6d50817386395a717ff14eca95e762831daeac1e5a7aa169d6f64698d7b"},{"path":"references/pytorch_ddp_notes.md","size":461,"sha256":"cbf7ca5f02440e864b59fb5095979b9e1d3404d7e5af94f9a89da364c31b6975"},{"path":"references/pytorch_device_mesh_tutorial.md","size":1222,"sha256":"af3941d5347faa986fd8aa1ef7a5b4ca4fb8f8a89f2380fe1edea0ecc4323863"},{"path":"references/pytorch_examples_fsdp2.md","size":742,"sha256":"1b966ff1655ed1280ec22b648b6d285f3794e4425b5301ecda941e4078c8766d"},{"path":"references/pytorch_fsdp1_api.md","size":396,"sha256":"cf4c0228216ae4cbcff2b57b35f4c677af8d2a803ce8bc38cb8df5602e11b716"},{"path":"references/pytorch_fsdp2_tutorial.md","size":2495,"sha256":"fe062caeacddbfee7eaafd7cad83c22300ee5450f2b04db02db5696408712e2a"},{"path":"references/pytorch_fully_shard_api.md","size":2884,"sha256":"1481a3f56062c55c54f1db2b3a1828dca14224a871b67be93b9b515f49cb12f8"},{"path":"references/pytorch_tp_tutorial.md","size":994,"sha256":"2c5857146d422427ac1ca3924aa87abd89a55d94ffb6f66a28a7d9a45c975cbb"},{"path":"references/ray_train_fsdp2_example.md","size":592,"sha256":"12f3f3f618b9ae59361bbbb0006439f9c5132d802f5faeb900c3f9cd97f0c3a9"},{"path":"references/torchtitan_fsdp_notes.md","size":666,"sha256":"989d3b724b551f150726e38ed17a297a1ca99b00ba9c77acc4aded3ea741e89c"}],"requires":{"mcp":[],"tools":[]},"safety":{"flags":[],"scannedAt":"2026-08-22","hasScripts":false,"networkEndpoints":["docs.pytorch.org","docs.ray.io"]}}