Faster RL Weight Transfers

21 Aug, 2026

Intro

In RL, every policy update requires moving newly updated weights from trainer to inference ranks. As models grow and training runs occur over larger numbers of nodes, this weight transfer is becoming a larger bottleneck in the training loop. This blog looks into how prime-rl deals with this issue.

NIXL

NIXL is fundamentally a library that sits on top of different network or storage backends and offers applications a standard interface with which to move data between memory locations. It handles a bunch of transfer-related bookkeeping, (non-exhaustively) including memory registration, transfer preparation, completion polling, notifications, etc. Each NIXL "agent" (i.e., our inference and trainer ranks) registers the memory that it owns, while an application-level orchestrator (in prime-rl, they call theirs ModelExpress) exchanges the resulting metadata between agents and tells them which transfers to perform and whether or not they can perform said transfer.

The ModelExpress "control plane" code is located here. NIXL then provides the data plane used to actually move the weights. This image does a good job of visualizing the communication relationship during runtime.

UCX is the default NIXL backend used in prime-rl and is itself another abstraction layer. It sits above transports such as InfiniBand, RoCE, and TCP and selects among the hardware and protocols available on a particular system. On an InfiniBand system with GPUDirect RDMA support, for instance, UCX can move data directly between GPU memory on different nodes.

Version 1

The code for this version can be found on this branch.

My original notion was that bytes over the wire were the main bottleneck. For some back-of-the-napkin math, consider moving the 755.6 GB GLM-5.2-FP8 checkpoint from one eight-GPU trainer pod to one eight-GPU inference pod. If you assume some optimistic, aggregate InfiniBand bandwidth of 400 GB/s, transferring the weights alone has a floor of roughly 755.6 / 400 = 1.9 seconds before accounting for any auxiliary computation.

So, my first implementation attempted to compress the weight updates and send the compressed update to the inference server. The compression scheme was quite simple. You take the previous weights and XOR them with the current weights. Assuming that most bits (certainly at least the exponent and most mantissa bits) would not change due to the small learning rate, the XOR buffer should contain mostly zeroes and thus be highly compressible. For dense models, this proved to be true. Compression rates were near 90x for BF16 and 115x for FP8 Qwen models. For MoE models, however, compression fell to around 40x and 60x for BF16 and FP8, respectively. It's not clear to me why an MoE might compress less well or if it's just the Qwen model itself. My (unsubstantiated) belief is that since each expert receives gradients from far fewer tokens than a dense model, this might produce higher-variance updates across expert tensors and therefore change more bits per update.

Regardless, the obvious problem was that compression was not free. We had to preserve the old weights to compute the XOR, compress that resulting buffer, transfer it, decompress it on the inference worker, apply the layout expected by vLLM, etc. Furthermore, for FP8, you have the additional issue of how and where to quantize. Reducing the number of bytes sent over the network only helped if all of this additional work cost less than the transfer time we saved.

For the compression itself, we used LZ4 through NVIDIA's nvCOMP library. We also wrapped AdamW so that it generated the XOR while updating the weights. During an update step, we would copy the old parameters into a separate buffer, update the parameters, then apply the XOR and compression. We performed this operation over a circular buffer to be able to better overlap all this overhead.

On the inference side, each rank would pull only the compressed frames it actually needed under its parallelism layout, and we had the workers double-buffer their receive and decode arenas so that one group of tensors could be transferred/decompressed while another group could be applied by vLLM. We then also used CUDA graphs to reduce the overhead in applying the updates on the inference GPUs. Because we could decompress the weights into the same address for vLLM, we could capture the XOR operations with CUDA once and replay them for each subsequent update.

Even with all of this, the dense BF16 end-to-end speedup was only about 1.2x and the dense FP8 speedup was about 1.6x. For the MoE, the Qwen3-30B-A3B speedups were only about 1x for BF16 and 1.25x for FP8. Lower compression ratios meant that we saved less transfer time while paying roughly the same fixed costs for compression, decompression, routing, and XOR application.

Version 2

The code for this version can be found on this branch with these changes.

Version 1 made it fairly clear that sending fewer bytes does not imply transferring weights faster. Also, I didn't feel the added complexity was worth it. The compressed payloads were a lot smaller relative to the uncompressed ones, but the compression pipeline also put a lot of work on the critical path.

The main change for version 2 is really the direction of the NIXL operation. The current prime-rl implementation is receiver (pull) driven in that the trainer marks some group of weights ready then inference worker notice readiness and issues NIXL reads to pull the weights it needs. Our changes allow for a trainer-driven pipeline in which each trainer rank issues NIXL writes directly into inference-owned buffers.

At initialization, for both the push and pull architectures, the trainer first publishes its agent metadata through ModelExpress, then each inference worker determines which trainer ranks' source tensors it needs and how vLLM applies changes to those tensors. The push path goes further and has each inference worker register a receiving circular buffer of some size and depth then publish those mappings from trainer shards to inference buffers back to the trainer. This allows the trainer at runtime to NIXL write directly from its local staging buffer into some corresponding inference buffer. One optimization that we didn't opt for was having the trainer write directly into the vLLM weight addresses themselves. This wouldn't be fully model-agnostic (think of vLLM needing to quantize, for example), but it would definitely be faster where the trainer and inference layouts match.

Regardless, the weights are transferred layer-by-layer (really layer group by layer group), and we create a trainer-side circular buffer as well so that the trainer can "write ahead" while the receiver applies some group of weights for vLLM. Diminishing returns (in terms of performance) occur around 4 or 8 slots, and in our setup, this circular buffer adds only up to 1 or 2 GB of memory usage. Then, unlike the current prime-rl architecture, which communicates trainer/inference buffer readiness over ModelExpress, we opt to communicate directly over NIXL.

In the pull path, every group requires the trainer to publish readiness, the inference workers to observe it, the inference workers to initiate their reads, and the trainer to observe that the buffers can be reused. Repeating this control-plane (ModelExpress) dialogue for each layer group seemed wasteful when NIXL could just transfer both the weights and a completion notification over the same path.

So, we replaced the repeated ModelExpress dialogue with something mbarrier-inspired. Each trainer rank will attach a NIXL notification to its write, which gets delivered after that rank’s bytes have arrived. Because each inference worker knows how many trainer ranks contribute to a group, it can begin applying the weights once all of those notifications have arrived (in our mbarrier analogy, the phase completing once its expected transaction bytes have arrived). After applying the group, it acknowledges consumption so the trainer can reuse the corresponding ring-buffer slot. The generation counter serves the role of an mbarrier phase, keeping both sides aligned as the same synchronization points are reused across policy updates.

This benchmark used two eight-GPU H100 SXM nodes connected over InfiniBand. One node ran eight trainer ranks, and the other ran four TP=2 inference replicas, with expert parallelism enabled for Qwen3-30B-A3B. The models were in BF16.

The number I cared about most here was the wall-clock time around broadcast_weights(), which is the function that encompasses the full trainer-to-inference weight update.

Version 2 is not only simpler than the compression approach but faster for weight transfer. There is no retained old model, custom AdamW, XOR representation, nvCOMP dependency, or complicated decoding or delta-application logic on the inference workers.

Conclusion

Bytes over the wire matter but hiding communication latency behind better scheduling mattered more. By changing who initiates the transfer and removing the repeated ModelExpress communications, our push architecture let us more aggressively overlap weight transfer with applying the weights on the inference workers.