NVIDIA introduces Reshard, a collective communication primitive that formulates reinforcement‑learning (RL) weight refit as a DTensor‑to‑DTensor resharding problem. An NCCL‑based reference implementation, NCCL M2N Reshard, is available in the NVIDIA/nccl-extensions GitHub repository and is demonstrated inside the NVIDIA NeMo RL framework. Reshard reduces redundant transfers by sending only overlapping shard regions and exploiting NVLink domain‑aware replica exchange, yielding 2.7×–5.7× speedups versus gather‑and‑broadcast in NeMo RL benchmarks and enabling more efficient refit for trillion‑parameter models.
Why refit is challenging in RL post‑training
During RL post‑training, inference trajectories are generated on the fly with the current policy model while training runs in parallel on separate GPUs. Although only the training GPUs update model parameters, the inference GPUs must periodically receive the latest weights — a process called refit. The common implementation uses a gather‑and‑broadcast pattern: training GPUs all‑gather their local shards to reconstruct each full parameter tensor, then broadcast that full tensor to inference GPUs, which keep the shard required by their local layout.
This approach is straightforward but inefficient: inference GPUs often receive the entire parameter even when they need only a shard, producing redundant traffic in both the gather and the broadcast stages. Communication payload scales with model size: for instance, Kimi K3 (≈2.8 trillion parameters) corresponds to roughly 5.6 TB of BF16 weights, and Qwen 3.8 2.4T A95B (≈2.4 trillion parameters) to roughly 4.8 TB. The two‑stage all‑gather plus broadcast logical payload is therefore on the order of twice the model weight size (≈11.2 TB and ≈9.6 TB respectively), which must be transferred at every refit.
Recasting refit as a DTensor‑to‑DTensor transfer
To minimize transferred data, NVIDIA frames refit as resharding a DTensor from the training layout to the inference layout. DTensor (PyTorch distributed tensor) encodes a tensor’s global shape, a device mesh, and placements (how the tensor is sharded or replicated). Given source and destination DTensor descriptions, the exact overlapping regions between source shards and destination shards are computable: a destination shard may be assembled from parts of multiple source shards, and replication means the same data may be available from several source GPUs.
Example: an 8×8 weight matrix sharded by rows across two training GPUs (each holds a 4×8 shard) and by columns across four inference GPUs (each needs an 8×2 shard). Each destination shard can be assembled by receiving 4×2 regions from each source. Representing both layouts as DTensors makes these overlaps explicit and enables direct shard‑to‑shard transfers.
The Reshard collective
Reshard is a new collective operation that transfers the logical tensor values from a source DTensor distribution to a destination DTensor distribution. A single Reshard call is specified by the source and destination DTensor descriptions, encapsulating the communication logic in a library so that algorithms can improve independently of applications.
NCCL M2N Reshard is the first open‑source reference implementation. Built on NCCL, it maps a tensor distributed across M source GPUs into the layout required by N destination GPUs. In RL refit, training GPUs typically form the source group and inference GPUs the destination group. The implementation computes exact source/destination intersections, maps them onto the physical network topology, and plans transfers accordingly.
How NCCL M2N Reshard maximizes network bandwidth
Rather than sending every source contribution separately to every destination replica (which duplicates network traffic), NCCL M2N Reshard groups destination replicas by NVIDIA NVLink domain. For each domain, it evenly distributes incoming contributions across the replica ranks so that each receiving rank gets a distinct subset of the total data over the network. The ranks then exchange their subsets over NVLink, allowing each replica to assemble the complete shard without duplicating network transfers. If replicas span multiple NVLink domains, one copy per contribution is forwarded across domains and the local balanced exchange repeats in each domain. Network transfers and NVLink exchanges are pipelined to hide latency.
This design ensures each source shard is transferred across the network only once and that NVLink is used efficiently to complete shards locally, instead of redundantly saturating network links.
NeMo RL integration and a hybrid practical design
NVIDIA provided a NeMo RL reference implementation to demonstrate Reshard in RL refit. Supplying DTensor descriptions for each parameter on training and inference instances is straightforward in principle, but the diversity of layer types and sharding strategies makes full coverage challenging.
To be both efficient and broadly usable, NeMo RL uses a hybrid approach: the feed‑forward projection weights (which account for about 95% of parameters in recent LLMs) follow the Reshard path, while smaller or more specialized weights (for example QKV projection weights) use the conventional gather‑and‑broadcast path, preserving coverage for edge cases while making the bulk of the payload much more efficient.
For runtime changes in GPU membership (nodes added/removed during failure recovery), frameworks can either recreate the communication group or use NCCL shrink/grow APIs. NeMo RL currently recreates groups for resiliency; shrink/grow‑based support is planned.
Benchmark results
NeMo RL benchmarks on NVIDIA DGX H100 (NVL8) and DGX GB200 (NVL72) systems demonstrate significant latency reductions for refit when using Reshard versus gather‑and‑broadcast. Selected results:
- NVIDIA DGX H100, Qwen 3 235B, TP4 x PP8 x EP16 → TP8 x DP16: gather‑and‑broadcast 6.4 s → Reshard 1.9 s
- NVIDIA DGX H100, DeepSeek v3 671B, PP16 x EP16 → TP32 x DP8: 14.9 s → 2.6 s
- NVIDIA DGX GB200, DeepSeek v3 671B, PP16 x EP16 → TP32 x DP8: 6.1 s → 2.3 s
- NVIDIA DGX GB200, Nemotron Ultra v3, TP8 x EP32 x PP2 → TP8 x DP8: 8.2 s → 2.3 s
These measurements correspond to a 2.7×–5.7× speedup range. In a refit‑only microbenchmark, Reshard completed refit for the Qwen3.8 2.4T BF16 model in 4.3 seconds, showing the approach scales to trillion‑parameter models.
How to get started
The NCCL M2N Reshard reference code is available in the NVIDIA/nccl-extensions GitHub repository. In NeMo RL you can enable the Reshard refit path by configuring:
policy:
generation:
colocated:
enabled: false
refit_transport: nccl_reshard
The NeMo RL refit guide and NCCL Reshard configuration requirements document which backends, layouts, and precision modes are supported.
Acknowledgements
NVIDIA acknowledges contributions from Ke Wen, Ching‑Hsiang Chu, Guyue Huang, Seonjin Na, and Nitin to NCCL M2N Reshard and the nccl‑reshard weight refit implementation in NeMo RL.
Significance
By reframing refit as a DTensor reshard problem and implementing a network‑aware NCCL collective that minimizes duplicated transfers, Reshard addresses a growing bottleneck in RL post‑training as models reach multi‑trillion parameter scales. The approach reduces per‑step refit traffic and latency, which has direct implications for accelerating RL training workflows that separate training and inference resources.
