NCCL M2N speeds up tensor reshaping for large GPU clusters
NCCL M2N: A Layout- and Topology-Aware Collective for Distributed Tensor Resharding
Distributed, Parallel, and Cluster Computing
Summary
When computers train large AI models, they often need to rearrange pieces of data across many GPUs, which can be slow and inefficient. The authors designed NCCL M2N, a smarter way to shuffle data by understanding the layout of data and the network connections between GPUs. This new approach shares data more evenly and overlaps communication steps to save time. As a result, it can transfer data much faster and reduces the time taken in a real AI training task by over 12%.
What this means in practice
- •For distributed ai engineers: Accelerate weight synchronization in large-scale distributed AI training by reducing communication overhead during tensor reshaping.
- •For high-performance computing teams: Improve data redistribution speeds in applications requiring complex multi-GPU data layouts using hierarchical and topology-aware communication schedules.
Authors
Kaushik Kandadi, Youngeun Kwon, Sreeram Potluri, Ching-Hsiang Chu, Ke Wen, Pouya Kousha, Sangkug Lym, Nitin Nitin, Manjunath Gorentla Venkata
Abstract
Distributed training and rollout generation often use different tensor layouts, requiring model weights to be resharded across distinct process groups. This M-to-N redistribution is not directly expressed by standard collectives. Flat direct sends duplicate traffic across destination replicas, while gather-then-broadcast concentrates network injection at one root and transfers data that destinations do not need. For DeepSeek-V3 on 256 GPUs, legacy all-gather plus broadcast accounts for 29.4% of the reported reinforcement-learning step time. We present NCCL M2N, a layout- and topology-aware collective primitive for distributed tensor resharding. Given source and destination meshes and placements, it derives the required transfer regions and a global communication schedule. Its hierarchical route balances source contributions across eligible destination ranks, forwards one copy between destination NVLink domains, and completes local replication over NVLink. Network transfer and local replication overlap, avoiding replica-multiplied source egress. An aggregate data-movement model captures the limits of both stages. We evaluate NCCL M2N on up to 256 GB200 GPUs in an NVL72 cluster with NDR InfiniBand. A single FFN-MoE layer transfer achieves up to 7.9x speedup over flat direct sends (9.8 ms versus 77.3 ms). In a separate 256-GPU DeepSeek-V3 NeMo-RL experiment, NCCL M2N reduces reported weight-sync time from 5.78 s to 2.77 s, a 2.09x speedup over legacy all-gather plus broadcast, and reduces step time by 12.7%.