pytorch-fsdp

Apply fully_shard to each submodule bottom-up before applying it to the root model, passing a device mesh that defines the data-parallel workers. FSDP shards parameters, gradients, and optimizer states across workers and all-gathers parameters before each forward pass.

FSDP2 uses DTensor-based dim-0 per-parameter sharding instead of FSDP1's flat-parameter sharding, giving a simpler sharding representation and communication-free sharded state dicts. FSDP2 also avoids 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, and the Gloo backend for CPU training. Gloo can serve as a fallback for GPUs if NCCL encounters problems, though it runs slower.

FSDP2 produces sharded state dicts containing DTensors without requiring all-gathers. Use DTensor APIs like DTensor.full_tensor() or PyTorch Distributed Checkpoint's distributed state dict APIs to reshard them into full state dicts.

Hangs occur when ranks run different numbers of iterations and collective communications mismatch. Wrap the training loop in the Join context manager from torch.distributed.algorithms, which shadows collectives for joined processes until all ranks finish.