pytorch-fsdp

Apply fully_shard bottom-up: call it on each sublayer first, then on the root model. Each call groups parameters for one all-gather collective, and the root call groups remaining parameters like embeddings. The optimizer must be initialized with the resulting DTensor parameters.

FSDP2 uses DTensor-based per-parameter sharding on dim-0 instead of FSDP1's flat-parameter sharding, giving a simpler sharding representation 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, 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 DDP model and optimizer in Join([model, optimizer]). It shadows collective communications for joined processes, preventing hangs when ranks have different numbers of inputs.

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

Yes, passing a 2D mesh to fully_shard produces HSDP: parameters are sharded across the first mesh dimension and replicated across the second, with (Replicate(), Shard(0)) placement. A 1D mesh gives standard fully sharded FSDP.