[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"skill:01M74ZBT2WTWD9JMH3ZKP5969B":3,"skill-same:01M74ZBT2WTWD9JMH3ZKP5969B":104},{"skillNo":4,"name":5,"nameZh":6,"description":7,"descriptionZh":6,"repositoryFullName":8,"ownerLogin":9,"ownerAvatarUrl":10,"ownerAccountNo":11,"ownerDisplayName":12,"repositoryNo":13,"starsCount":14,"licenseSpdxId":15,"licenseFileUrl":6,"sourceHtmlUrl":16,"categories":17,"keywords":18,"publishedAt":19,"linkOnly":20,"markdownContent":21,"markdownBody":22,"frontMatter":23,"refName":30,"commitSha":31,"skillDirectory":32,"forksCount":33,"repositoryHtmlUrl":34,"repositoryHomepageUrl":35,"ownerHtmlUrl":36,"ownerBlogUrl":37,"ownerDescription":38,"updatedAt":39,"install":40,"resources":53},"01M74ZBT2WTWD9JMH3ZKP5969B","pytorch-fsdp2",null,"Adds PyTorch FSDP2 (fully_shard) to training scripts with correct init, sharding, mixed precision\u002Foffload config, and distributed checkpointing. Use when models exceed single-GPU memory or when you need DTensor-based sharding with DeviceMesh.","KalarisLabs\u002Fresearch-agent-skills","KalarisLabs","https:\u002F\u002Favatars.githubusercontent.com\u002Fu\u002F300878568?v=4","01M74Y7MT23GXXC56VM6TVSXZY","Kalarislabs","01M74Y7MTQM9REXZX4SNXWX95C",132,"MIT","https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002FSKILL.md",[],[],"2026-11-11 17:47:46",false,"---\nname: pytorch-fsdp2\ndescription: Adds PyTorch FSDP2 (fully_shard) to training scripts with correct init, sharding, mixed precision\u002Foffload config, and distributed checkpointing. Use when models exceed single-GPU memory or when you need DTensor-based sharding with DeviceMesh.\nlicense: MIT\nmetadata:\n  version: 1.0.0\n  category: ml-training\n  maintainer: Kalaris Labs\n  tags: PyTorch, FSDP2, Fully Sharded Data Parallel, Distributed Training, DTensor, Device Mesh, Sharded Checkpointing, Mixed Precision, Offload, Torch Distributed\n  dependencies: torch\n---\n\n# 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\u002Foffload 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\u002Fpytorch_fully_shard_api.md`, `references\u002Fpytorch_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\u002Fpytorch_ddp_notes.md`, `references\u002Fpytorch_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\u002Ftutorial; 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 \u003Cgpus_per_node> ...` and ensure `RANK`, `WORLD_SIZE`, `LOCAL_RANK` are visible.\n\nReference: `references\u002Fpytorch_fsdp2_tutorial.md` (launch commands and setup), `references\u002Fpytorch_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\u002Fpytorch_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\u002Fpytorch_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\u002Fpytorch_fully_shard_api.md` (bottom-up requirement and why).\n\n---\n\n### 4) Configure `reshard_after_forward` for memory\u002Fperf 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\u002Fpytorch_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\u002Freductions on H100\u002FA100-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\u002FNVLink traffic and runtime overhead.\n\nReference: `references\u002Fpytorch_fully_shard_api.md` (MixedPrecisionPolicy \u002F 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 \u002F 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\u002Fgradients are DTensors.\n\nReference: `references\u002Fpytorch_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\u002Floads 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` \u002F `set_model_state_dict` with `StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True, ...)`\n- For optimizer: `get_optimizer_state_dict` \u002F `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\u002Fpytorch_dcp_overview.md` (DCP behavior and caveats)\n- `references\u002Fpytorch_dcp_recipe.md` and `references\u002Fpytorch_dcp_async_recipe.md` (end-to-end usage)\n- `references\u002Fpytorch_fsdp2_tutorial.md` (DTensor vs DCP state-dict flows)\n- `references\u002Fpytorch_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\u002Fload via `torch.distributed.checkpoint` helpers.\n\nReference: `references\u002Fpytorch_fsdp2_tutorial.md`, `references\u002Fpytorch_fully_shard_api.md`, `references\u002Fpytorch_device_mesh_tutorial.md`, `references\u002Fpytorch_dcp_recipe.md`.\n\n### Workflow B: Add DCP save\u002Fload (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\u002Fpytorch_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()` \u002F register forward.\n3. **Is `fully_shard()` applied bottom-up?**  \n   If only root is sharded, expect worse memory\u002Fperf 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\u002Fpytorch_fully_shard_api.md`, `references\u002Fpytorch_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\u002Fbackward\u002Fstep with `model(inputs)` and DTensor-aware patterns\n- `checkpoint_save\u002Fload()`: DCP or distributed state dict helpers\n\nConcrete examples live in `references\u002Fpytorch_examples_fsdp2.md` and the official tutorial reference.\n\n---\n\n## References\n- `references\u002Fpytorch_fsdp2_tutorial.md`\n- `references\u002Fpytorch_fully_shard_api.md`\n- `references\u002Fpytorch_ddp_notes.md`\n- `references\u002Fpytorch_fsdp1_api.md`\n- `references\u002Fpytorch_device_mesh_tutorial.md`\n- `references\u002Fpytorch_tp_tutorial.md`\n- `references\u002Fpytorch_dcp_overview.md`\n- `references\u002Fpytorch_dcp_recipe.md`\n- `references\u002Fpytorch_dcp_async_recipe.md`\n- `references\u002Fpytorch_examples_fsdp2.md`\n- `references\u002Ftorchtitan_fsdp_notes.md` (optional, production notes)\n- `references\u002Fray_train_fsdp2_example.md` (optional, integration example)\n\n## Related skills\n\n- `huggingface-accelerate`: Wraps existing PyTorch training scripts with HuggingFace Accelerate (Accelerator class, accelerate config, accelerate launch) so the same c…\n- `distributed-llm-pretraining-torchtitan`: Provides PyTorch-native distributed LLM pretraining using torchtitan with 4D parallelism (FSDP2, TP, PP, CP).\n","# 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\u002Foffload 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\u002Fpytorch_fully_shard_api.md`, `references\u002Fpytorch_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\u002Fpytorch_ddp_notes.md`, `references\u002Fpytorch_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\u002Ftutorial; 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 \u003Cgpus_per_node> ...` and ensure `RANK`, `WORLD_SIZE`, `LOCAL_RANK` are visible.\n\nReference: `references\u002Fpytorch_fsdp2_tutorial.md` (launch commands and setup), `references\u002Fpytorch_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\u002Fpytorch_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\u002Fpytorch_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\u002Fpytorch_fully_shard_api.md` (bottom-up requirement and why).\n\n---\n\n### 4) Configure `reshard_after_forward` for memory\u002Fperf 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\u002Fpytorch_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\u002Freductions on H100\u002FA100-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\u002FNVLink traffic and runtime overhead.\n\nReference: `references\u002Fpytorch_fully_shard_api.md` (MixedPrecisionPolicy \u002F 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 \u002F 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\u002Fgradients are DTensors.\n\nReference: `references\u002Fpytorch_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\u002Floads 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` \u002F `set_model_state_dict` with `StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True, ...)`\n- For optimizer: `get_optimizer_state_dict` \u002F `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\u002Fpytorch_dcp_overview.md` (DCP behavior and caveats)\n- `references\u002Fpytorch_dcp_recipe.md` and `references\u002Fpytorch_dcp_async_recipe.md` (end-to-end usage)\n- `references\u002Fpytorch_fsdp2_tutorial.md` (DTensor vs DCP state-dict flows)\n- `references\u002Fpytorch_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\u002Fload via `torch.distributed.checkpoint` helpers.\n\nReference: `references\u002Fpytorch_fsdp2_tutorial.md`, `references\u002Fpytorch_fully_shard_api.md`, `references\u002Fpytorch_device_mesh_tutorial.md`, `references\u002Fpytorch_dcp_recipe.md`.\n\n### Workflow B: Add DCP save\u002Fload (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\u002Fpytorch_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()` \u002F register forward.\n3. **Is `fully_shard()` applied bottom-up?**  \n   If only root is sharded, expect worse memory\u002Fperf 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\u002Fpytorch_fully_shard_api.md`, `references\u002Fpytorch_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\u002Fbackward\u002Fstep with `model(inputs)` and DTensor-aware patterns\n- `checkpoint_save\u002Fload()`: DCP or distributed state dict helpers\n\nConcrete examples live in `references\u002Fpytorch_examples_fsdp2.md` and the official tutorial reference.\n\n---\n\n## References\n- `references\u002Fpytorch_fsdp2_tutorial.md`\n- `references\u002Fpytorch_fully_shard_api.md`\n- `references\u002Fpytorch_ddp_notes.md`\n- `references\u002Fpytorch_fsdp1_api.md`\n- `references\u002Fpytorch_device_mesh_tutorial.md`\n- `references\u002Fpytorch_tp_tutorial.md`\n- `references\u002Fpytorch_dcp_overview.md`\n- `references\u002Fpytorch_dcp_recipe.md`\n- `references\u002Fpytorch_dcp_async_recipe.md`\n- `references\u002Fpytorch_examples_fsdp2.md`\n- `references\u002Ftorchtitan_fsdp_notes.md` (optional, production notes)\n- `references\u002Fray_train_fsdp2_example.md` (optional, integration example)\n\n## Related skills\n\n- `huggingface-accelerate`: Wraps existing PyTorch training scripts with HuggingFace Accelerate (Accelerator class, accelerate config, accelerate launch) so the same c…\n- `distributed-llm-pretraining-torchtitan`: Provides PyTorch-native distributed LLM pretraining using torchtitan with 4D parallelism (FSDP2, TP, PP, CP).\n",{"name":5,"license":15,"metadata":24,"description":7},{"tags":25,"version":26,"category":27,"maintainer":28,"dependencies":29},"PyTorch, FSDP2, Fully Sharded Data Parallel, Distributed Training, DTensor, Device Mesh, Sharded Checkpointing, Mixed Precision, Offload, Torch Distributed","1.0.0","ml-training","Kalaris Labs","torch","main","29fe36d26127205057d5e0ed8425f6b0e9a8c13a","skills\u002Fpytorch-fsdp2",16,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills","https:\u002F\u002Fdocs.kalarislabs.com\u002Fresearch-agent-skills","https:\u002F\u002Fgithub.com\u002FKalarisLabs","kalarislabs.com","Applied AI lab building the infrastructure for the agentic era","2026-11-14 17:49:08",{"source":8,"skillName":5,"commands":41,"prompt":51,"zipAvailable":52},[42,45,48],{"tool":43,"command":44},"npx","npx skills add KalarisLabs\u002Fresearch-agent-skills --skill pytorch-fsdp2",{"tool":46,"command":47},"bunx","bunx skills add KalarisLabs\u002Fresearch-agent-skills --skill pytorch-fsdp2",{"tool":49,"command":50},"pnpm","pnpm dlx skills add KalarisLabs\u002Fresearch-agent-skills --skill pytorch-fsdp2","安装技能，技能地址：https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002FSKILL.md",true,[54,60,64,68,72,76,80,84,88,92,96,100],{"path":55,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":58,"sourceHtmlUrl":59},"references\u002Fpytorch_dcp_async_recipe.md","REFERENCE","参考文档",808,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_dcp_async_recipe.md",{"path":61,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":62,"sourceHtmlUrl":63},"references\u002Fpytorch_dcp_overview.md",1022,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_dcp_overview.md",{"path":65,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":66,"sourceHtmlUrl":67},"references\u002Fpytorch_dcp_recipe.md",1043,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_dcp_recipe.md",{"path":69,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":70,"sourceHtmlUrl":71},"references\u002Fpytorch_ddp_notes.md",461,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_ddp_notes.md",{"path":73,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":74,"sourceHtmlUrl":75},"references\u002Fpytorch_device_mesh_tutorial.md",1222,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_device_mesh_tutorial.md",{"path":77,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":78,"sourceHtmlUrl":79},"references\u002Fpytorch_examples_fsdp2.md",742,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_examples_fsdp2.md",{"path":81,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":82,"sourceHtmlUrl":83},"references\u002Fpytorch_fsdp1_api.md",396,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_fsdp1_api.md",{"path":85,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":86,"sourceHtmlUrl":87},"references\u002Fpytorch_fsdp2_tutorial.md",2495,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_fsdp2_tutorial.md",{"path":89,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":90,"sourceHtmlUrl":91},"references\u002Fpytorch_fully_shard_api.md",2884,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_fully_shard_api.md",{"path":93,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":94,"sourceHtmlUrl":95},"references\u002Fpytorch_tp_tutorial.md",994,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fpytorch_tp_tutorial.md",{"path":97,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":98,"sourceHtmlUrl":99},"references\u002Fray_train_fsdp2_example.md",592,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Fray_train_fsdp2_example.md",{"path":101,"resourceType":56,"resourceTypeDesc":57,"sizeBytes":102,"sourceHtmlUrl":103},"references\u002Ftorchtitan_fsdp_notes.md",666,"https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fpytorch-fsdp2\u002Freferences\u002Ftorchtitan_fsdp_notes.md",[105,112,119,126,133,140],{"skillNo":106,"name":107,"nameZh":6,"description":108,"descriptionZh":6,"repositoryFullName":8,"ownerLogin":9,"ownerAvatarUrl":10,"ownerAccountNo":11,"ownerDisplayName":12,"repositoryNo":13,"starsCount":14,"licenseSpdxId":15,"licenseFileUrl":6,"sourceHtmlUrl":109,"categories":110,"keywords":111,"publishedAt":19},"01M74ZSRSWD92NYM8TNJRH35DK","zarr-python","Guides use of Zarr-Python 3 for storing chunked, compressed N-dimensional arrays and groups, with local, in-memory, ZIP, and fsspec-backed S3\u002FGCS\u002FHTTP stores, plus NumPy, Dask, and Xarray integration. Covers array creation, resizing and appending, attributes, chunk and shard sizing, codecs, consolidated metadata, and v2-to-v3 migration. Use when creating or opening Zarr arrays or groups, choosing chunk sizes or compression for large datasets, reading or writing arrays on S3 or GCS, appending to time-series arrays, or migrating code from Zarr-Python 2 to 3. Use when setting up parallel I\u002FO with Dask or Xarray. For labeled multi-dimensional datasets, use Xarray instead.","https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fzarr-python\u002FSKILL.md",[],[],{"skillNo":113,"name":114,"nameZh":6,"description":115,"descriptionZh":6,"repositoryFullName":8,"ownerLogin":9,"ownerAvatarUrl":10,"ownerAccountNo":11,"ownerDisplayName":12,"repositoryNo":13,"starsCount":14,"licenseSpdxId":15,"licenseFileUrl":6,"sourceHtmlUrl":116,"categories":117,"keywords":118,"publishedAt":19},"01M74ZSK6JG5HKNC1AHF4D97G9","whisper","Transcribes and translates audio with OpenAI's Whisper (openai-whisper Python package and whisper CLI), covering model sizes from tiny to large plus turbo, language specification, initial prompts, word and segment timestamps, temperature fallback, batch processing, and subtitle generation. Use when transcribing speech, podcasts, meetings, or video audio to text. Use when translating non-English speech to English. Use when working with noisy or multilingual audio in any of 99 languages. Use when choosing a Whisper model size for available VRAM. Use when generating subtitles from audio. Not for speaker diarization or live captioning; use AssemblyAI or Deepgram instead.","https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fwhisper\u002FSKILL.md",[],[],{"skillNo":120,"name":121,"nameZh":6,"description":122,"descriptionZh":6,"repositoryFullName":8,"ownerLogin":9,"ownerAvatarUrl":10,"ownerAccountNo":11,"ownerDisplayName":12,"repositoryNo":13,"starsCount":14,"licenseSpdxId":15,"licenseFileUrl":6,"sourceHtmlUrl":123,"categories":124,"keywords":125,"publishedAt":19},"01M74ZSBYQ2EGDFTTA4ZK6K5TR","weights-and-biases","Logs and tracks machine learning experiments with Weights & Biases (W&B, wandb): metrics, hyperparameters, checkpoints, sweeps, artifacts with lineage, model registry, custom charts and shareable reports. Includes integrations for PyTorch, HuggingFace Transformers, PyTorch Lightning and Keras\u002FTensorFlow. Use when logging training runs and comparing them across configurations, running automated hyperparameter sweeps, versioning datasets and models as artifacts, registering model versions, or sharing results with a team workspace. Use when setting up wandb.init for offline or unstable-connection training. Not for general data versioning outside ML runs.","https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fweights-and-biases\u002FSKILL.md",[],[],{"skillNo":127,"name":128,"nameZh":6,"description":129,"descriptionZh":6,"repositoryFullName":8,"ownerLogin":9,"ownerAvatarUrl":10,"ownerAccountNo":11,"ownerDisplayName":12,"repositoryNo":13,"starsCount":14,"licenseSpdxId":15,"licenseFileUrl":6,"sourceHtmlUrl":130,"categories":131,"keywords":132,"publishedAt":19},"01M74ZS5MZTK84WHEQXCVDZ6GB","waypoint-bio","Use when working with Outpost Bio's open microbiome foundation models - the Waypoint checkpoints (Waypoint-6m, Waypoint-45m, Waypoint-170m), the Atlas pretraining corpus, the Compass eight-task benchmark, or the `waypoint` CLI from the `waypoint-bio` package. Covers embedding microbiome samples, fine-tuning on taxonomic abundance data, benchmarking a checkpoint on Compass, pretraining a GPT-2 model on taxonomic abundance profiles, and converting MetaPhlAn, Kraken2, QIIME 2, or MGnify abundance tables into waypoint format.","https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fwaypoint-bio\u002FSKILL.md",[],[],{"skillNo":134,"name":135,"nameZh":6,"description":136,"descriptionZh":6,"repositoryFullName":8,"ownerLogin":9,"ownerAvatarUrl":10,"ownerAccountNo":11,"ownerDisplayName":12,"repositoryNo":13,"starsCount":14,"licenseSpdxId":15,"licenseFileUrl":6,"sourceHtmlUrl":137,"categories":138,"keywords":139,"publishedAt":19},"01M74ZS0TCW96WJAKSG07C5Z7A","verl-rl-training","Provides guidance for training LLMs with reinforcement learning using verl (Volcano Engine RL). Use when implementing RLHF, GRPO, PPO, or other RL algorithms for LLM post-training at scale with flexible infrastructure backends.","https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fverl-rl-training\u002FSKILL.md",[],[],{"skillNo":141,"name":142,"nameZh":6,"description":143,"descriptionZh":6,"repositoryFullName":8,"ownerLogin":9,"ownerAvatarUrl":10,"ownerAccountNo":11,"ownerDisplayName":12,"repositoryNo":13,"starsCount":14,"licenseSpdxId":15,"licenseFileUrl":6,"sourceHtmlUrl":144,"categories":145,"keywords":146,"publishedAt":19},"01M74ZRVG6SSWGT1N3Q55134KA","venue-templates","Prepare journal manuscripts, conference papers, research posters, and grant documents using venue-specific formatting guidance and bundled LaTeX scaffolds. Use when selecting an official template, checking current page or anonymity rules, adapting academic writing to a venue, or inspecting a submission PDF.","https:\u002F\u002Fgithub.com\u002FKalarisLabs\u002Fresearch-agent-skills\u002Fblob\u002Fmain\u002Fskills\u002Fvenue-templates\u002FSKILL.md",[],[]]