Fast layout-aware communication reduces training delays on large gpu clusters
The Endgame of FSDP Communication: Layout-Aware Collectives on Copy Engines, RDMA, and NVSwitch
Distributed, Parallel, and Cluster ComputingMachine LearningPerformance
Summary
Training very large AI models on many GPUs takes time partly because data must be reshaped and communicated inefficiently between parts of the system. The authors show that these reshaping steps and communication slow down the training significantly. They introduce TIDE, a new way to communicate data without extra copying and with better use of system hardware like RDMA and NVSwitch. Their approach speeds up training by up to 27% on large models and reduces delays caused by communication interfering with computation.
What this means in practice
- •For deep learning engineers: Optimize distributed training of large language models by reducing communication overhead and memory usage across many GPUs.
- •For high performance computing teams: Improve efficiency and scalability of GPU cluster communications using layout-aware collectives for faster parallel computations.
Authors
Zejian Wang, Yifan Zhao, Kaihui Ren, Xiaoling Zhou, Yubin Wang, Chenhao Guan, Lin Zhang, Yufeng Li, Zhikang Mo, Zhuoru Ma, Yingtong Xiong, Jihu Guo, Yiqiao Liao, Chengru Song
Abstract
Fully Sharded Data Parallelism (FSDP2) maintains parameters and gradients as per-parameter tensors, whereas its collective path operates on flat, rank-major buffers. Converting between these layouts copies all-gather outputs into parameter storage and packs gradients into full-group fp32 buffers. Layout conversion alone accounts for 6.4% of step time in Llama-3.1 70B training at 8K tokens per GPU. Communication kernels impose a second cost: a default NCCL reduce-scatter can slow a concurrent GEMM that would otherwise complete in a single wave by 69%. We present TIDE, a set of zero-copy collectives built on public NCCL 2.31+ APIs. TIDE preserves FSDP2's per-parameter storage and compute kernels while folding layout placement into collective addressing. TIDE-AG delivers shards directly to parameter storage using NVSwitch multicast and paced recursive-doubling RDMA, with no SMs used for transport or placement. TIDE-RS consumes per-parameter bf16 gradients directly and reduces them into fp32 shards (a mixed-precision mode that NCCL does not provide) using copy engines and short-lived thread blocks (CTAs) for node-local accumulation and NVSwitch for cross-node reduction. A two-phase pipeline separates gradient reads from network progress; Slot-Reuse optionally stages gradients in vacated all-gather slots. On 64 GPUs, TIDE reduces step time by up to 27.1% relative to Stock FSDP2 and 6.0% relative to FSDP2+SymmMem across Llama-3.1 70B, Wan2.2, and Qwen3-30B-A3B. It reduces reduce-scatter SM time by factors of 12.5-14.5 and cuts interference with concurrent compute by up to an order of magnitude. Its reduce-scatter buffer remains nearly constant from 16 to 64 GPUs at a fixed per-GPU output size, and its peak memory is up to 4.7 GiB below FSDP2+SymmMem's and, with Slot-Reuse on Wan2.2, below Stock FSDP2's.