NVIDIA says Reshard can cut reinforcement learning weight refit time

Why reinforcement learning refit is becoming a bottleneck

In reinforcement learning post-training, systems generate training trajectories by running inference with the current policy model. Training and inference can run in parallel on separate GPU resources, but the inference GPUs must periodically receive the latest weights produced by the training GPUs. NVIDIA refers to this process as refit.

The two sides often distribute model weights differently because they have different computational requirements. Their configurations can use different combinations of data, context, tensor, expert and pipeline parallelism. Refit must also update every model weight on the inference GPUs, so the communication payload grows with the total size of the model.

Existing implementations commonly use gather and broadcast. Training GPUs first all-gather their local shards to reconstruct a complete parameter, then broadcast that full tensor to inference GPUs. Each inference GPU keeps only the portion it needs, meaning the process moves data that is not required by the destination.

NVIDIA gives Kimi K3 and Qwen 3.8 2.4T A95B as examples. Assuming BF16 weights at two bytes per parameter, the company estimates weight sizes of approximately 5.6 TB and 4.8 TB, respectively. Including both all-gather and broadcast produces estimated logical payloads of approximately 11.2 TB and 9.6 TB.

Reshard sends overlapping regions directly

NVIDIA's approach reformulates refit as a DTensor-to-DTensor transfer. PyTorch's distributed tensor, or DTensor, abstraction describes a tensor's global shape, the logical arrangement of GPUs, and how the tensor is sharded or replicated across that arrangement.

With this description, the training layout becomes the source and the inference layout becomes the destination. The Reshard collective calculates which parts of source shards overlap with destination shards, allowing the required regions to move directly instead of reconstructing a complete intermediate tensor.

NVIDIA illustrates the idea with an 8 x 8 weight matrix. Two training GPUs hold row shards measuring 4 x 8, while four inference GPUs need column shards measuring 8 x 2. Each destination can receive a 4 x 2 region from each source, assembling the exact shard it requires.

The company describes Reshard as a collective communication call for arbitrary DTensor-to-DTensor transfers. Like other collective operations, it lets an application invoke the communication pattern as one operation while the underlying library manages the transfer logic.

Open-source implementation targets varied GPU layouts

NCCL M2N Reshard is NVIDIA's open-source reference implementation of the collective, available through the NVIDIA/nccl-extensions GitHub repository. Built on NCCL, it transfers a tensor distributed across M source GPUs into the layout required by N destination GPUs. In reinforcement learning refit, the training GPUs form the source group and the inference GPUs form the destination group.

According to NVIDIA's AI-generated summary of the post, benchmarks in NeMo RL on NVIDIA DGX H100 and DGX GB200 systems showed Reshard refit achieving a 2.7x to 5.7x speedup over conventional gather-and-broadcast refit across models including Qwen 3 235B and DeepSeek v3 671B.

The same summary says a hybrid approach in NeMo RL applies Reshard to feed-forward projection weights, which dominate the payload, while retaining gather-and-broadcast for smaller weights. NVIDIA directs developers to the NeMo RL refit guide for supported backends, layouts and precision modes, and to the NCCL Reshard configuration requirements for implementation details.

For users building distributed reinforcement learning systems, the announcement offers a way to reduce unnecessary weight-transfer traffic as models and GPU deployments grow. Its practical scope depends on the layouts and implementation options supported by the relevant NeMo RL and NCCL configurations.