Accelerating Reinforcement Learning Weight Refit with the Reshard Collective

During reinforcement learning (RL) post-training, training trajectories are generated on the fly by running inference with the current policy model. For training efficiency, it is common to run inference and training in parallel on separate GPU resources in a disaggregated setup. Although the policy model is updated only on the training GPUs, the system must periodically refresh the weights on the inference GPUs with the latest parameters generated on the training GPUs. This transfer of model weights between training and inference GPUs is called refit.

Refit involves two defining characteristics that pose challenges. First, inference and training have different computational requirements, so the weights are distributed differently across inference and training instances by having different parallel mappings (data, context, tensor, expert, and pipeline parallelism).

Second, because refit must update all model weights on the inference GPUs, the communication payload is on the order of the total size of the model’s weights. This payload is already very large for the latest frontier models and is expected to continue growing. Together, these characteristics mean that refit requires high-performance communication across a wide variety of distributed source and destination layouts.

This post explains how AI training frameworks can address these refit challenges using the Reshard collective, a new type of collective communication primarily designed to accelerate RL weight refit. NCCL M2N Reshard, an open source reference implementation of the Reshard collective, is available through the NVIDIA/nccl-extensions GitHub repo.

Existing refit implementations commonly use a gather-and-broadcast approach to handle differences between training and inference layouts.

Training GPUs all-gather their local shards to reconstruct the complete parameter.

The complete parameter tensor is broadcast to the inference GPUs.

Each inference GPU retains the shard required by its local layout.

Using the complete tensor as an intermediate representation makes it straightforward to transfer weights between different layouts. However, an inference GPU receives the full parameter even when it needs only a shard. This process creates redundant communication in both the gathering and broadcasting phases, which is suboptimal.

The size of this communication grows with the model size, which continues to scale. For example, Kimi K3 contains 2.8 trillion parameters, while Qwen 3.8 2.4T A95B contains 2.4 trillion parameters. Assuming BF16 weights at two bytes per parameter, their weight sizes are approximately 5.6 TB and 4.8 TB, respectively.

Accounting for both the all-gather and broadcast stages gives an estimated logical payload of approximately twice the size of the model’s weights (Table 1).

Qwen 3.8 2.4T A95B

Table 1. Two-stage logical payload estimate assuming BF16 weights

For today’s largest open weight models, the communication payload already reaches around 10 TB per refit, which must be performed at every training step. As LLMs continue to scale, the overhead associated with this large payload will become even more significant. As a result, a more sophisticated approach is needed to perform weight refit that reduces the actual payload size while properly accounting for diverse weight layouts.

To overcome these limitations, start from the principle that refit should transfer only the theoretical minimum amount of weight data. However, achieving this presents several challenges. Training and inference instances can use different parallelism configurations and GPU counts, producing different shard boundaries. A destination shard may therefore need to be assembled from portions of several source shards, and vice versa.

Replication further complicates the mapping since the same weights may be available from multiple source GPUs or required by multiple destination GPUs. It is nontrivial to determine which source GPUs should send which tensor regions to each destination while achieving high performance across all possible configurations.

The key to addressing this challenge is to identify the right abstraction for this unique communication pattern. The team turned to the existing PyTorch distributed tensor (DTensor) abstraction, which is designed to represent tensors that are sharded or replicated across GPUs under different forms of data and model parallelism. The key insight is that weight refit can be reformulated as resharding a DTensor from the training layout to the inference layout, regardless of the parallelism configurations used by the two instances.

A DTensor layout is described by the tensor’s global shape, a device mesh specifying the logical arrangement of GPUs, and placements specifying how the tensor is sharded or replicated across that mesh. Refit can use the source description to specify where the updated weights reside and the destination description to specify which weights each inference GPU needs.

For example, consider an 8 x 8 weight matrix split by rows across two training GPUs and by columns across four inference GPUs. Assuming the tensor is sharded in the row dimension on the training GPUs and sharded in the column dimension on the inference GPUs, each training GPU holds a 4 x 8 shard, while each inference GPU needs an 8 x 2 shard. Each destination can obtain its required shard by receiving a 4 x 2 region from each source. These regions follow directly from the overlap between the source and destination shards, as illustrated in Figure 1.

Figure 1. Direct shard-to-shard transfer of an 8 x 8 tensor from two training GPUs to four inference GPUs. The pattern can be cleanly abstracted as a DTensor-to-DTensor reshard

As the figure shows, any kind of sharded tensor can be easily represented as a DTensor format, and the refit can be redefined as a DTensor-to-DTensor reshard problem. This finding provides the foundation for the Reshard collective.

Reshard is a new type of collective communication call for arbitrary DTensor-to-DTensor transfers. Like all-reduce, reduce-scatter, all-gather, broadcast, and so on, reshard defines a communication pattern that applications can invoke as one operation. It transfers the values of a logical tensor from its source distribution into its destination distribution.

A Reshard call can be specified simply by the source and destination DTensor descriptions. This abstraction not only simplifies the refit problem, but also encapsulates the communication logic in a library, allowing the underlying algorithms to evolve and improve independently of the application.

NCCL M2N Reshard is open source, provided to the community as the first reference implementation of the Reshard collective. Built on NCCL, this implementation transfers a tensor distributed across M source GPUs into the layout required by N destination GPUs. In RL refit, the training GPUs form the source group and the inference GPUs form the destination group.

Figure 2 shows how this implementation accepts the layout descriptions and local buffers to carry out the Reshard operation.

Figure 2. The nccl.m2n.reshard example that transfers a tensor sharded across four training GPUs (TP4) into replicas on four inference GPUs (DP4)

NCCL M2N Reshard first derives the exact intersections between source and destination shards, then maps those transfers onto the physical GPU network topology. Instead of sending every source contribution independently to every destination replica, NCCL M2N groups replicas by NVIDIA NVLink domain. Within each domain, it distributes the incoming contributions as evenly as possible across the replica ranks, so each receiving rank receives a distinct subset of the total data over the network.

The receiving ranks then exchange those subsets over NVLink, enabling every replica to assemble the complete shard without duplicating network traffic. If replicas span multiple NVLink domains, one copy of each contribution is forwarded through the domains, and the same balanced local exchange occurs within each domain. Network transfer and NVLink replication are pipelined to hide each other’s latency.

Figure 3 illustrates this data movement scheme for a tensor moving from TP4 on the source to TP2 x DP2 on the destination, with all four destination GPUs in one NVLink domain. With direct point-to-point transfers (Figure 3a), each source shard is sent to both data parallel (DP) replicas. As a result, each shard is transferred twice over the network redundantly. This means that each shard gets roughly half the network bandwidth (assuming each shard saturates the network bandwidth).

NCCL M2N Reshard transfers each source shard across the network just once, supplying each destination GPU with half of its required data (Figure 3b). The destination GPUs exchange the missing half over NVLink. For example, destination GPU 0 receives S0 from source GPU 0 over the network and S1 from destination GPU 2 over NVLink. Network transfers and NVLink exchanges proceed concurrently. This way, each shard gets the full network bandwidth.

Figure 3. Naive Direct P2P sends each source shard across the two DP replicas (left). With NCCL M2N Reshard, each destination GPU receives part of its data directly from the source over the network and the rest from another destination GPU over NVLink (right)

To demonstrate the effectiveness of the Reshard collective for weight refit, a reference implementation is provided in the NVIDIA NeMo RL framework. NeMo RL is an open source framework designed to streamline and scale RL post-training. The clear abstraction of the Reshard collective makes the required implementation within the framework relatively simple. For a given parallelism configuration, the framework only needs to provide a DTensor description for each parameter for training and inference instances, which should be straightforward because DTensor was designed for this purpose.

In practice, however, the variety of layer types and the diversity in how training and inference engines shard parameters under each parallelism configuration make it challenging to handle every edge case through the Reshard collective.

To develop an efficient and sustainable design, we built on the key observation that feed-forward projection weights in the transformer block dominate the weight refit payload. In the latest LLMs, these weights account for around 95% of the model’s parameters.

Based on this observation, our design uses a hybrid approach in which dominant feed-forward projection weights take the reshard path, while all other small weights (for example, QKV projection or mamba weights) use the conventional gather-and-broadcast refit path, which is optimized for broad functional coverage. This approach combines efficient reshard communication for most of the payload with the same functional coverage as conventional gather-and-broadcast refit.

The sharding layout can change at runtime as nodes are added or removed during failure recovery. To accommodate changes in GPU membership, frameworks can either destroy and recreate the communication group or use NCCL shrink and grow APIs to adjust the existing communicator group. NeMo RL currently uses the first approach for resiliency support. Resiliency support using NCCL shrink and grow APIs is planned for future updates.

Table 2 summarizes reshard-based refit latencies measured using the NeMo RL implementation introduced above. The benchmarks cover a range of models on NVIDIA DGX H100 (NVL8) and NVIDIA DGX GB200 (NVL72) systems. Across various benchmarks, the reshard refit achieves a speedup of 2.7x to 5.7x over conventional gather-and-broadcast refit.

Table 2. Refit latencies measured in NeMo RL

This new refit is designed for trillion-parameter models. In our refit-only microbenchmark, Reshard refit took only 4.3 seconds for the Qwen3.8 2.4T BF16 model, demonstrating its effectiveness at the trillion-parameter scale for which it was designed.

The reference implementation in the NVIDIA NeMo RL framework enables you to adopt and extend this approach for your own RL workloads.

The NeMo RL refit guide and NCCL Reshard configuration requirements describe the supported backends, layouts, and precision modes for this feature.

We sincerely appreciate the contributions of Ke Wen, Ching-Hsiang Chu, Guyue Huang, Seonjin Na, and Nitin to NCCL M2N Reshard and nccl-reshard weight refit in NeMo RL.