domains / families / data / pytorch-fsdp2
pytorch-fsdp2
separate
Adds PyTorch FSDP2 (fully_shard) to training scripts with correct init, sharding, mixed precision/offload config, and distributed checkpointing. Use when models exceed single-GPU memory or when you need DTensor-based sharding with DeviceMesh.
When to use
- Your model **doesn’t fit** on one GPU (parameters + gradients + optimizer state).
- You want an eager-mode sharding approach that is **DTensor-based per-parameter sharding** (more inspectable, simpler sharded state dicts) than FSDP1.
- You may later compose DP with **Tensor Parallel** using **DeviceMesh**.
- You need strict backwards-compatible checkpoints across PyTorch versions (DCP warns against this).
- You’re forced onto older PyTorch versions without the FSDP2 stack.
When not to use
- Strict backwards-compatible checkpoints across PyTorch versions
- Older PyTorch versions without the FSDP2 stack
- Classic distributed data parallel training (DDP)
- Original FSDP wrapper (FSDP1)
Skill page
Source