Source-linked AI summary
NCCL M2N: A Layout- and Topology-Aware Collective for Distributed Tensor Resharding
Kaushik Kandadi, Youngeun Kwon, Sreeram Potluri, Ching-Hsiang Chu, Ke Wen, Pouya Kousha, Sangkug Lym, Nitin Nitin, Manjunath Gorentla Venkata
TL;DR
Distributed training and rollout generation use different tensor layouts, but standard collectives do not express the resulting cross-group resharding directly. NCCL M2N derives the needed transfer regions and topology-aware schedule; in the reported experiments, it improved layer-transfer and DeepSeek-V3 weight-sync performance, while reducing step time.
Problem
M-to-N resharding across distinct process groups is not directly expressed by standard collectives, while common workarounds duplicate traffic or concentrate injection and over-deliver data.
Source: §I · p. 1Method
NCCL M2N derives transfer regions and a global schedule from both layouts, then routes one copy per destination NVLink domain and completes replication locally.
Source: §I · p. 2, B. Layout-Derived Schedule and… · p. 3Results
Up to 7.9× speedup over flat direct sends for a single FFN-MoE layer transfer; in a separate 256-GPU DeepSeek-V3 experiment, weight-sync time fell from 5.78 s to 2.77 s, a 2.09× speedup, and step time fell by 12.7%.
Source: §IX · p. 9Takeaways & Limitations
The reported results demonstrate the benefit of planning tensor resharding jointly across layouts and the interconnect hierarchy.
Source: §IX · p. 9Existing workarounds remain constrained: direct sends duplicate contributions across replicas, while gather-then-broadcast concentrates traffic at one root and transfers the full tensor.
Source: A. Limitations of Existing Approaches · p. 3
Abstract
from arXiv · showhide
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%.
I. INTRODUCTION
Different training and rollout layouts require weights to move across process groups during refit, but standard collectives do not express this M-to-N redistribution. NCCL M2N derives the transfer schedule from both layouts and routes data across the interconnect hierarchy.
I. INTRODUCTION
- 29.4% of the reported DeepSeek-V3 reinforcement-learning step time was spent on legacy all-gather-plus-broadcast weight sync.The refit transfer lies on the critical path before the next rollout can begin.Source: §I · p. 1, Table I · p. 1
- NCCL M2N derives exact transfer regions and a global schedule from source and destination layouts, including replica ownership and disjoint process groups.The collective is intended for M-to-N resharding across differing tensor layouts.Source: §I · p. 2
- Its hierarchical route sends one copy per destination NVLink domain, balances injection across eligible NICs, and completes replica fan-out locally over NVLink.This avoids replica-multiplied network traffic while preserving supported M-to-N layout semantics.Source: §I · p. 2
- On a GB200 NVL72 cluster, M2N achieved up to 7.9× speedup for a single FFN-MoE layer transfer and cut DeepSeek-V3 weight-sync time by 2.09× and step time by 12.7%.The DeepSeek-V3 result was measured in a separate 256-GPU NeMo-RL run against legacy all-gather plus broadcast.Source: §I · p. 2
A. RL Training Loop and Refit
- RL refit transfers updated policy weights from trainer to generator after each training step, before the next rollout begins.The transfer schedule affects the iteration critical path.Source: A. RL Training Loop and Refit · p. 2
B. Distributed Tensor Layouts
Distributed tensor layouts describe how a global tensor is sharded or replicated over a process mesh, and resharding must map overlapping regions between source and destination layouts. Existing collectives and lower-level communication operations do not themselves derive that cross-group schedule.
B. Distributed Tensor Layouts
- A layout is the tuple of global shape, process mesh, and placement list, whose entries specify sharding or replication and determine each rank’s local shape and offset.The described materialized parameter placements are SHARD and REPLICATE.Source: B. Distributed Tensor Layouts · p. 2
- An M-to-N reshard can change shard dimensions, shard counts, and replication factors across groups that may be disjoint, so the runtime must derive overlapping source and destination regions.TP and EP can shard different dimensions, requiring intersections of multidimensional regions.Source: B. Distributed Tensor Layouts · p. 2
C. Collective Communication and Interconnect Hierarchy
- Standard NCCL collectives do not accept arbitrary endpoint layouts or derive byte-level intersections for redistribution between different process groups.Lower-level mechanisms move data and signals but leave region selection and replica placement to the reshard schedule.Source: C. Collective Communication and… · p. 2
A. Limitations of Existing Approaches
- Flat direct sends duplicate source contributions across replicas, while gather-then-broadcast concentrates traffic at one root and over-delivers the full tensor.The missing capability is a collective that jointly plans intersections, byte offsets, and topology from both endpoint layouts.Source: A. Limitations of Existing Approaches · p. 3
B. Layout-Derived Schedule and Hierarchical Route
NCCL M2N derives a transfer plan from endpoint layouts and topology, then routes contributions across destination domains and replicates them locally. Its layout semantics cover intersecting shards, destination replicas, and per-stage resharding for pipeline parallelism.
B. Layout-Derived Schedule and Hierarchical Route
- A layout-derived schedule combines source and destination descriptors with runtime topology to determine overlapping regions and the leaders and links that carry them.The layout class formalized here supports the route’s stated traffic and load-balancing properties.Source: B. Layout-Derived Schedule and… · p. 3
- The hierarchical route injects each required contribution once, balances assignments across eligible leaders, and forwards copies across destination domains before local NVLink replication.This removes replica-multiplied source egress and avoids concentrating traffic at one receiver or over-delivering the full tensor.Source: B. Layout-Derived Schedule and… · p. 3
IV. RESHARD SEMANTICS AND DESCRIPTORS
- A reshard leaves each destination rank with its prescribed local shard, and the current implementation requires source and destination meshes to be disjoint communicator intervals.Every communicator rank supplies both endpoint layouts so all ranks derive the same schedule; participating ranks supply local buffers for their endpoint meshes.Source: §IV · p. 3
- In Example A, each destination TP slot consolidates two adjacent source shards, then DP=2 replication duplicates the resulting ownership across destination ranks.The flat direct-send realization sends each contribution to both replicas, transferring two copies of the unique payload over the network.Source: §IV · p. 4, Figure 1 · p. 4
- With PP>1, the caller issues a separate reshard per stage using that stage’s ranks plus the generator, rather than spanning the full trainer mesh.In Example B, generator rank 4 participates in both PP communicators and accumulates one full tensor per stage.Source: §IV · p. 4, Table IV · p. 4
V. LAYOUT-DERIVED DESIGN AND IMPLEMENTATION
NCCL M2N separates the derivation of bytes to move from the choice of route, with implementation support for per-stage PP communicators. The packaged API uses communicator-local meshes for each reshard.
V. LAYOUT-DERIVED DESIGN AND IMPLEMENTATION
- The planner derives required byte ranges from the endpoint layouts, while the topology-aware route determines how those bytes move.The implementation realizes that schedule using the mechanisms described in the section.Source: §V · pp. 4–5
- The packaged Python API invokes reshard with stage-local source and destination tensors, using meshes whose ranks are local to each PP communicator.The listed membership shows trainer ranks joining their stage communicator and rank 4 joining both.Source: §V · p. 5
- For pipeline parallelism, each stage uses its own communicator: ranks 0–1 form the source mesh and local rank 2 the destination mesh.The example uses PP comm 0 on global ranks {0, 1, 4} and PP comm 1 on {2, 3, 4}; rank 4 joins both.Source: §V · p. 5, Figure 2 · p. 5
A. Deterministic Schedule Derivation
The deterministic schedule derivation tests source–destination interval overlaps to identify transfer regions. A shared enumeration gives ranks the same schedule without coordination.
A. Deterministic Schedule Derivation
- For each source–destination pair, the planner intersects the ranks’ intervals along each tensor dimension and skips the pair if any overlap is empty.When all dimension overlaps are nonempty, their combined intervals define the exact contribution; replicated source copies collapse to one deterministic owner.Source: A. Deterministic Schedule Derivation · p. 5
- Every rank uses the same enumeration, so the schedule is derived without coordination.The planner applies the same source–destination tests across ranks.Source: A. Deterministic Schedule Derivation · p. 5
- In the illustrated DP=4 route, each network stage carries B bytes, while local NVLink exchange reconstructs each TP slice on both replicas.The figure distinguishes source injection and ring forwarding from local replica reconstruction.Source: Figure 3 · p. 5
B. Hierarchical Topology Route
The hierarchical route assigns contributions to destination leaders, forwards copies across domains, and replicates locally. Its data-movement model accounts for both network transfer and NVLink replication.
B. Hierarchical Topology Route
- Each contribution enters through one selected leader, is forwarded across destination domains, and is replicated locally over NVLink.Unlike flat direct sends, which transmit each contribution separately to every replica, the hierarchical route forwards copies across domains and then fans them out locally.Source: B. Hierarchical Topology Route · p. 5, C. Aggregate Data-Movement Model · p. 6
- When contributions are uneven, leaders receive loads differing by at most one, balancing eligible destination ranks.For each destination shard, replicas receive the base load and the first remainder-count replicas receive one extra contribution.Source: B. Hierarchical Topology Route · p. 5
C. Aggregate Data-Movement Model
- For the data-movement bound, the source injects one copy of B, while the destination domain has Kρ active network receivers.Across multiple domains, ring forwarding operates concurrently for large payloads, so the domain count affects startup and tail latency but does not multiply the steady-state B/BN term.Source: C. Aggregate Data-Movement Model · p. 6
- Within each destination domain, NVLink creates the remaining ρ −1 copies, and network movement overlaps local replication in the SOL bound.Local copy volume is (ρ −1)B; the model omits local preparation, chunking, protocol, and startup costs, so it is not a prediction of every end-to-end latency component.Source: C. Aggregate Data-Movement Model · pp. 5–6
T SOL
The SOL analysis contrasts root-injection-limited broadcast with M2N’s concurrent source injection. Its source-path count must reflect shared NIC bandwidth.
T SOL
- When broadcast crosses the inter-mesh network and root injection is the bottleneck, the model bounds broadcast time by root injection.The supplied passage introduces this root-limited case but does not give the complete expression.Source: T SOL · p. 6
- In the source-limited regime, M2N injects the Ns source shards concurrently, unlike a single root-injection path.The passage frames this as the M2N source-injection bound relative to broadcast.Source: T SOL · p. 6
- Ns counts independent source-injection paths, but should be replaced by the effective path count when shards share NIC bandwidth.Source: T SOL · p. 6
- Section VI-B instantiates the aggregate expressions and checks that NVLink is not on the critical path.Source: T SOL · p. 6
D. Implementation and Availability
NCCL M2N’s initial release implements layout-derived resharding through hierarchical transports. The evaluation covers a single-layer transfer and an end-to-end DeepSeek-V3 experiment.
D. Implementation and Availability
- The initial release supports one- and two-dimensional meshes, one- to three-dimensional tensors, and at most one sharded mesh axis per layout.PACK stages strided contributions contiguously before transfer and scatters them into destination buffers; beta PIPE pipelines bounded chunks through the same route.Source: D. Implementation and Availability · p. 6
A. Experimental Setup
- The single-layer experiment transfers a 7168 MB BF16 expert gate-projection tensor; the end-to-end experiment uses DeepSeek-V3 in BF16 on 256 GPUs.The single-layer setup uses a 16-GPU EP trainer group and a TP8 generator group with DP∈{2, 4, 8}; end-to-end trainer and generator groups each span two NVL72 domains.Source: A. Experimental Setup · p. 6
- The single-layer benchmark compares flat direct sends with the hierarchical route, while end-to-end measurements compare NCCL M2N against NeMo-RL’s legacy all-gather-plus-broadcast path.End-to-end metrics use one-step off-policy execution, sequence length 1536, and non-warmup step-time and weight-sync-time measurements.Source: A. Experimental Setup · p. 6
B. Single-Layer Microbenchmark
- 8.96 ms is the SOL bound for all three single-layer replication points, with NVLink below the network term even at eight-way local replication.This is a model bound, not a measured transport latency.Source: B. Single-Layer Microbenchmark · p. 7
T SOL
In the single-layer microbenchmark, hierarchical speedup over flat direct sends increases with replication. Flat direct-send latency rises with replication, while rank-0 all-gather plus broadcast is root-injection limited.
T SOL
- 10.251, 13.834, and 9.775 ms are the hierarchical transport latencies across the three points, despite a 4× increase in generator replication.They are 14%, 54%, and 9% above the 8.96 ms SOL, respectively; residual latency includes staging, pack/unpack, signalling, and startup costs.Source: T SOL · p. 7
- 77.258 ms at rd = 8 versus 20.859 ms at rd = 2: flat direct-send latency rises 3.7× as replication grows 4×.The values are 16%, 10%, and 8% above the corresponding {17.92, 35.84, 71.68} ms SOL values.Source: T SOL · p. 7
- 222.065 and 239.649 ms are the validated rank-0 packed AG+bcast latencies at rd = 2 and 4, or 1.55× and 1.67× the 143.36 ms SOL.The rd = 8 measurement is pending; Figure 4 shows only its analytical AG+bcast SOL value.Source: T SOL · p. 7
- 7.9× speedup over flat direct sends at rd = 8, up from 2.0× at rd = 2, with 2.8× at rd = 4.The endpoint speedups closely approach the analytical {2, 8}× prediction; the middle point retains more transport overhead.Source: T SOL · p. 7
C. End-to-End DeepSeek-V3 RL Weight Sync
In the 256-GPU DeepSeek-V3 NeMo-RL comparison, NCCL M2N lowers weight-sync time against legacy all-gather plus broadcast. Related work distinguishes M2N’s joint, layout-derived schedule from transfer interfaces, lower-level communication abstractions, and offline checkpoint resharding.
C. End-to-End DeepSeek-V3 RL Weight Sync
- Weight-sync time falls from 5.78 s to 2.77 s, a 2.09× speedup over legacy all-gather plus broadcast.The comparison uses non-colocated trainer and generator groups in asynchronous NeMo-RL.Source: C. End-to-End DeepSeek-V3 RL Weight Sync · pp. 7–8
- Reported step time decreases from 19.68 s to 17.18 s (12.7%), while similar KL errors provide only a limited stability check.The KL values are 0.001869 with M2N and 0.001912 with the legacy path, not evidence of convergence or training quality.Source: C. End-to-End DeepSeek-V3 RL Weight Sync · p. 8
VII. RELATED WORK
- Unlike DTensor-aware transfer interfaces and storage-backed reads, M2N derives one global schedule from both endpoint layouts and coordinates shard groups jointly.The related work describes torchstore and Monarch as bridging meshes, and TensorHub as tracking weight locations for topology-optimized reads.Source: §VII · p. 8
- M2N adds a layout-derived schedule above programmable communication abstractions, directing existing paths on what to move and which path to use.The passage names GIN, LSA, and multimem as examples of those paths.Source: §VII · p. 8
- Offline checkpoint resharding reconstructs target layouts across runs, whereas M2N transfers tensors between live process groups within a run; the approaches are complementary.Source: §VII · p. 8
VIII. DISCUSSION AND FUTURE WORK
The discussion identifies limits and possible extensions to M2N, including multi-axis layouts, elastic communicators, adaptive transport, overlapping meshes, recovery, and in-flight quantization. The conclusion reports speedups in a layer transfer and a separate DeepSeek-V3 NeMo-RL experiment.
VIII. DISCUSSION AND FUTURE WORK
- Current layouts permit at most one sharded mesh axis; simultaneous multi-axis sharding and meshes with more than two axes remain future work.Those extensions would cover placements such as TP×EP×DP natively, without composing two-dimensional meshes.Source: §VIII · p. 8
- Integrating shrink-and-grow communicator APIs could let M2N update its reshard schedule in-place as source or destination mesh sizes change.Source: §VIII · p. 8
- An adaptive policy could select staging or remote-memory access per transfer; reproducible evaluation should retain the explicit environment override.Source: §VIII · pp. 8–9
- Multi-rank source and destination meshes are currently disjoint, and protocol errors after host-RMA entry require communicator-level shutdown.Overlapping membership and distributed failure agreement would require protocol changes beyond descriptor validation.Source: §VIII · p. 9
- The staging kernel could apply quantization in-flight as chunks move through the staging ring, avoiding a separate pass and intermediate quantized buffer.Source: §VIII · p. 9
IX. CONCLUSION
- Its layout-derived route avoids repeated flat-send transfers and full-tensor gathering, using NVLink for replication within each destination domain.Source: §IX · p. 9
- M2N achieves up to 7.9× speedup for a single FFN-MoE layer transfer and reduces DeepSeek-V3 weight-sync time from 5.78 s to 2.77 s and step time by 12.7%.Source: §IX · p. 9