pytorch-fsdp

Call fully_shard on each submodule bottom-up before applying it to the root module, so each layer forms its own all-gather group. The optimizer must be initialized with the resulting DTensor parameters, and forward must be invoked via model(input) to trigger the all-gather hooks.

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, as it provides the best performance and supports InfiniBand and GPUDirect. Use Gloo for CPU training or as a fallback if NCCL encounters problems.

Use the generic Join context manager from torch.distributed.algorithms, wrapping your Joinable objects such as DDP models and ZeroRedundancyOptimizer. It shadows collective communications of joined processes so training on uneven inputs does not hang or error.

Hangs occur when destroy_process_group is not called by all ranks, since ProcessGroupNCCL's destructor calls ncclCommAbort which must run collectively in a consistent order. Call destroy_process_group() once per trainer process near the end of main().