skills.rest
qcmuu
qcmuuCOMMUNITY
Explore all skills by this author →
GitHub

pytorch-fsdp2

Add PyTorch FSDP2 fully_shard to distributed training scripts with DTensor sharding and DCP checkpointing.

Updated May 2, 2026
Search Tags:#pytorch#distributed training#mixed precision#fsdp2#dtensor#dcp checkpointing#device mesh
One-click install
npx skills add https://github.com/qcmuu/AI-Research-Skills --skill pytorch-fsdp2-qcmuu
Or copy as Structured Prompt for Agent▼
Please help me install this Agent Skill.
Skill: pytorch-fsdp2
Source: https://github.com/qcmuu/AI-Research-Skills/tree/main/08-distributed-training/pytorch-fsdp2
Command: npx skills add https://github.com/qcmuu/AI-Research-Skills --skill pytorch-fsdp2-qcmuu

SYSTEM DOCUMENTATION & REQUIREMENTS

💡 This Skill requires torch, and includes references (resource) components.

What problem does it solve?

This Skill prevents out-of-memory failures and subtle FSDP2 misconfigurations by showing how to add PyTorch FSDP2 (fully_shard) to a distributed training script with the correct initialization, sharding order, precision/offload settings, and checkpoint workflow.

Core Features & Use Cases

  • Correct FSDP2 initialization and contract adherence: guides required launch, device setup, hook-safe forward invocation, and optimizer construction on DTensor parameters.
  • Bottom-up parameter sharding with controlled resharding: explains why fully_shard must be applied to submodules first and how to tune reshard_after_forward for memory/performance trade-offs.
  • Robust distributed checkpointing: recommends Distributed Checkpoint (DCP) or distributed state-dict helpers instead of naive torch.save for DTensor/FSDP2 models.
  • Use Case: you are training a transformer-sized model that exceeds single-GPU memory and need sharded training that can later compose with other parallelism (e.g., DeviceMesh for hybrid strategies).

Quick Start

Use pytorch-fsdp2 to retrofit a training script so it launches with torchrun, applies fully_shard bottom-up with the right DeviceMesh options, builds the optimizer after sharding, and checkpoints with DCP safely.

Frequently Asked Questions about pytorch-fsdp2

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I add PyTorch FSDP2 fully_shard to a distributed training script?▼

To add PyTorch FSDP2 fully_shard, apply it bottom-up to submodules first, build the optimizer on DTensor parameters, and use model(input) for hook execution. This ensures correct sharding and prevents out-of-memory failures.

Why does FSDP2 require applying fully_shard bottom-up to submodules?▼

Applying fully_shard bottom-up is required because it controls parameter sharding and resharding after forward passes. This order optimizes memory and performance trade-offs for transformer-like workloads exceeding single-GPU memory.

Does FSDP2 distributed checkpointing work with torch.save?▼

FSDP2 distributed checkpointing does not work with naive torch.save. You must use Distributed Checkpoint (DCP) or distributed state-dict helpers to correctly save and restore DTensor parameters across multi-rank execution.

Can I create an optimizer before applying fully_shard in PyTorch?▼

You cannot create an optimizer before applying fully_shard. The optimizer must be constructed after sharding on the resulting DTensor parameters to ensure gradients and states are correctly distributed across the device mesh.

What is the best way to configure mixed precision and offload for FSDP2?▼

The best way to configure mixed precision and CPU offload for FSDP2 is by passing the appropriate settings during fully_shard application. This manages memory usage while maintaining performance for large transformer models.

Can I use FSDP2 with DeviceMesh for hybrid parallelism strategies?▼

FSDP2 composes with DeviceMesh for hybrid parallelism strategies. You can configure DeviceMesh options during fully_shard application to integrate sharded training with other distributed techniques across multi-rank execution.

Related Skills

Explore More Software Engineering Skills→
ovachieverovachiever
52

pytorch-fsdp

Apply PyTorch FSDP for sharded, mixed-precision distributed training of large models.

Community
Advanced
ihatesea69ihatesea69
3

pytorch-fsdp2

Integrate PyTorch FSDP2 fully_shard into training scripts for large models.

Community
Advanced
handsomelong922handsomelong922

pytorch-fsdp2

Integrate PyTorch FSDP2 into training scripts for per-parameter sharding.

Community
Advanced
Orchestra-ResearchOrchestra-Research
11.5k

pytorch-fsdp2

Integrate PyTorch FSDP2 with DTensor sharding and DCP checkpointing into training scripts.

Official
Advanced