What problem does it solve?
This Skill enables the correct and efficient use of PyTorch's Fully Sharded Data Parallelism (FSDP2) for training large models that exceed single-GPU memory, by guiding the integration of its advanced features like DTensor-based sharding and distributed checkpointing.
Core Features & Use Cases
- FSDP2 Integration: Correctly initializes and applies
torch.distributed.fsdp.fully_shard bottom-up.
- Memory Optimization: Configures
reshard_after_forward, mixed precision, and CPU offload for memory-bound scenarios.
- Distributed Checkpointing: Implements robust saving and loading using Distributed Checkpoint (DCP) or state dict helpers.
- Use Case: Integrate FSDP2 into a large language model training script to shard parameters, gradients, and optimizer states across multiple GPUs, allowing models that would otherwise OOM to be trained effectively.
Quick Start
Integrate PyTorch FSDP2 into your existing training script by following the step-by-step procedure outlined in the skill's documentation.