pytorch-fsdp

Apply fully_shard bottom-up: call it on each sublayer (e.g., each Transformer layer) before calling it on the root model. Each call groups parameters for a single all-gather, and the root call groups remaining parameters like embeddings and output projections.

FSDP2 uses DTensor-based per-parameter sharding on dim-0 instead of FSDP1's flat-parameter sharding, giving simpler reasoning about sharded state and communication-free sharded state dicts. FSDP2 also removes record_stream memory management and exposes manual prefetching and collective scheduling APIs.

Use the NCCL backend for distributed training with CUDA GPUs, including InfiniBand hosts since NCCL supports GPUDirect. Use Gloo for CPU training or as a fallback if NCCL encounters problems, though Gloo runs slower on GPUs.

FSDP2 does not directly support full state dicts; it produces sharded state dicts containing DTensors. Convert them using DTensor.full_tensor() or PyTorch Distributed Checkpoint's distributed state dict APIs.

Hangs occur when ranks run different numbers of iterations, leaving collective calls unmatched. Use the generic Join context manager (torch.distributed.algorithms.Join) to shadow collectives of joined processes until all ranks finish.