Stop Waiting for Weights: Efficient RL Updates with WeightBridge

A library for efficient weight transfer across diverse LLM reinforcement learning settings

Xuanlin Jiang¹*, Samuel Hsia², Michael Kuchnik², Zachary DeVito², Minlan Yu¹, and Carole-Jean Wu²

¹ Harvard University · ² FAIR at Meta

Work carried out during an internship at Meta.

Open-source code · Paper · PyTorch Conference talk

Diverse RL Configurations Challenge Weight Transfer

RL alternates between rollout workers, which generate trajectories, and trainer workers, which use them to update the policy. Weight transfer sends the updated parameters back to rollouts. A typical system has one training engine and multiple independently running inference engines.

Sparse MoE models increase the parameters transferred relative to computation per token, while GPU compute improves faster than inter-node bandwidth. These two trends substantially increase weight transfer overhead in the RL pipeline and have drawn attention from Perplexity, ByteDance, Anyscale, Miles / LMSYS, Prime Intellect, and Ant Group.

Weight transfer involves three key questions:

1. What weights go where? — Layout discovery: identify the elements each trainer holds and each rollout needs.
2. How do they move? — Transfer planning: choose routes, RDMA operations, and layout-conversion kernels.
3. When are they available? — Synchronization and scheduling: coordinate updates with computation while preserving RL semantics.

Why diversity makes these questions challenging

Supporting these three steps becomes harder as RL configurations diversify. Existing mechanisms are efficient only in subsets of the layout and synchronization space: other configurations expose redundant transfers, sender bottlenecks, or layout-conversion overhead. Details of existing solutions

This diversity creates three challenges.

1. Heterogeneous layouts complicate correspondence. Different models and hardware drive diverse parallelism strategies and model layouts for both trainers and rollouts. Trainers and rollouts usually use mismatched layouts because of their inherently different performance goals and algorithmic flows. For example, in the figure below, the trainer uses larger tensor-parallel (TP) and pipeline-parallel (PP) degrees because gradients and optimizer states increase its memory requirements.

Specialized kernels in training and inference frameworks further increase differences in how tensors are stored on the GPU: they may split, transpose, or pad the same parameter tensor in entirely different ways. Together, these mismatches substantially complicate weight correspondence between trainers and rollouts.

2. Replication and topology complicate routing. Multiple trainers may hold the same weights, and multiple rollout workers may need them. Full-model broadcasts can deliver unneeded elements, causing redundancy, while independent fetches can overload one sender, causing imbalance. Efficient planning must eliminate redundant traffic and balance aggregate load across replica groups, using fast intra-node NVLink and slower inter-node RDMA appropriately.

3. Different staleness preferences of RL applications complicate worker coordination. Alignment tasks such as chat quality, instruction following, safety, and refusal behavior often have short rollouts that favor synchronous execution. Mathematics and competitive programming produce longer, more variable rollouts, making blocking updates more expensive and motivating asynchronous execution. Algorithms also differ in stale-trajectory tolerance. The weight transfer mechanism must coordinate workers carefully to support these different synchronization modes while keeping the APIs simple enough to integrate into different backend engines.

WeightBridge Unifies Layout Discovery, Transfer Planning, and Worker Coordination

WeightBridge answers the three questions through automatic correspondence extraction, global transfer planning and execution, and backend-independent worker coordination. Extraction and planning happen once, while the resulting routing plan is reused across policy updates.

1. Extracting correspondence through backend-native weight loaders

The model’s checkpoint can act as a common coordinate system between trainer and rollouts, and each backend’s existing load_weights function already maps checkpoint tensors into its runtime parameters. WeightBridge probes that function with identifiable index values, recovers each worker’s weight source map, and intersects trainer and rollout maps to find the correspondence between weight elements. Layout discovery details

2. Planning non-redundant, load-balanced transfers

Two principles guide the routing plan:

1. Trainers collectively transmit only one copy of the model, avoiding redundant sends and a single-source bottleneck.
2. Each rollout node receives exactly one copy of each weight element it needs over RDMA, with additional local replicas distributed over NVLink.

The planner assigns distinct slices to replicated workers, then executes three stages: receive slices from trainers, exchange across rollout nodes over RDMA, and distribute within each node over NVLink. Transfer planning details

3. Coordinating workers across synchronization modes under minimal APIs

Sender and receiver adapters expose three key APIs:

• connect() establishes the trainer–rollout connection and prepares the reusable transfer plan.
• send_weights() initiates transfer of a new model version from trainer workers.
• poll_requests() services connection and update requests at rollout inference boundaries, coordinating when workers apply new weights.

Independent sender- and receiver-staging switches provide four synchronization modes: fully synchronous, fully asynchronous, or overlap on only one side. Host staging enables overlap but adds copies and memory costs; the application chooses its throughput-staleness tradeoff. Worker coordination details

import wbridge.frontend.adapters as wbridge

def trainer_worker_process():
    adapter = wbridge.SenderAdapter(...)      # WeightBridge
    adapter.connect()                         # WeightBridge

    for _ in range(num_iters):
        train_one_step()
        adapter.send_weights()                # WeightBridge

def rollout_worker_process():
    adapter = wbridge.ReceiverAdapter(...)    # WeightBridge

    while True:
        generate_one_step()
        updated = adapter.poll_requests()     # WeightBridge

Minimal integration. Trainers connect once and send updated weights; rollouts poll between inference steps.

See the integration guide for API details and current implementation constraints.

Performance

Substantial Reduction in GPU Stall Time Across Diverse RL Settings

We evaluated WeightBridge with models from 16B to 1.03T parameters on 2–64 H100 nodes. Each node has eight GPUs. The primary metric is average GPU stall time (AGST) per weight transfer. Across the four deployments, WeightBridge lowers AGST by 16.6–29.4× over Miles Broadcast and 5.4–42.0× over Miles P2P. Layout and replica-count studies show that global planning avoids the single-source bottleneck while keeping transfer time close to the bandwidth-based lower bound.

Fast Integration into Different Backends

A coding agent integrated WeightBridge into Miles on its first attempt in 24.5 minutes of active work, excluding experiment waits. The integration reduced end-to-end weight-transfer time by 19.7× for Qwen3-30B-A3B and 8.84× for Qwen3-32B, with correctness checks passing.

WeightBridge also integrates with an internal framework using different training, rollout, and RDMA backends from Miles, demonstrating API generality across different backend stacks. With minimal manual integration effort, it reduced end-to-end weight-transfer time by 6.29–11.00× across one to four rollout replicas.

Explore the open-source code · Read the paper · Integration guide · PyTorch Conference talk