pytorch-fsdp2_skill

Overview

This skill adds PyTorch FSDP2 (fully_shard) to training scripts so very large models can train across multiple GPUs with correct init, sharding, mixed precision/offload options, and distributed checkpointing. It codifies the required launch pattern, bottom-up sharding, DTensor-aware optimizer creation, and recommended Distributed Checkpoint (DCP) flows. The goal is a reliable retrofit or new training loop that avoids common FSDP2 pitfalls and memory spikes.

How this skill works

The skill initializes torch.distributed and sets the CUDA device per process, optionally creates a DeviceMesh, and then applies fully_shard bottom-up (shard submodules first, then root). It supports meta-device model construction, materialization with to_empty/reset_parameters, DTensor-friendly optimizer creation after sharding, mixed-precision/offload policies, and DCP or distributed state-dict checkpointing. The runtime enforces using model(input) so FSDP2 hooks run and offers patterns for gradient accumulation and clipping with DTensors.

When to use it

• Model does not fit on a single GPU (params + grads + optimizer state).
• You prefer DTensor-based per-parameter sharding and inspectability over older FSDP1 semantics.
• You plan to combine data-parallel with tensor-parallel via DeviceMesh later.
• You need efficient multi-rank checkpointing and load-time resharding (DCP).
• You have a recent PyTorch release that supports FSDP2 and DCP.

Best practices

• Launch with torchrun and set torch.cuda.set_device(int(os.environ['LOCAL_RANK'])).
• Apply fully_shard bottom-up: shard Transformer blocks/submodules before the root module.
• Create the optimizer after sharding so it holds DTensor parameters.
• Use model(inputs) (not model.forward) so FSDP2 hooks and unshard/reshard behavior run.
• Prefer DCP or distributed state-dict helpers for checkpointing; avoid naive torch.save on DTensors.

Example use cases

• Retrofit an existing training loop to enable multi-GPU training for models that exceed single-GPU memory.
• Train a transformer initialized on meta, then materialize weights to GPUs with to_empty and reset_parameters.
• Use BF16 mixed precision on H100/A100 with an explicit MixedPrecisionPolicy and optional CPU offload for parameter-heavy models.
• Save and restore checkpoints across ranks with DCP, allowing load-time resharding onto a different DeviceMesh.
• Implement gradient accumulation using FSDP2’s set_requires_gradient_sync instead of FSDP1 no_sync().