vLLM shards weight transfer for large-scale online reinforcement learning
A Ray Direct Transport engine cuts repeated model-weight synchronization by sending each inference rank only the shard it needs.
Online reinforcement learning repeatedly moves updated model weights from training workers to inference workers. At trillion-parameter scale, that synchronization can become a memory and latency bottleneck in its own right. vLLM’s new sharded weight-transfer engine changes the transport pattern: inference ranks pull only the weight slices they need from trainer ranks over Ray Direct Transport, rather than receiving a full-model broadcast and discarding most of it.
What changed
The implementation uses Ray Direct Transport with its NIXL backend for direct GPU-to-GPU transfers. It supports dense and mixture-of-experts models, fused or per-expert checkpoints, and quantized inference. A “recording tensor” dry run observes vLLM’s own model loaders and builds a sharding plan for each worker, avoiding a separate set of model-specific tensor-layout rules in the RL framework.
During synchronization, trainer ranks gather one transformer block at a time. They avoid gathering across pipeline stages and leave distributed experts local where possible. Inference workers pull their assigned slices, while gathering, transfer and post-processing are pipelined. vLLM says that this bounded-memory design also preserves CUDA graphs through its layerwise-reloading path.
The reported result
In a 48-node test using eight H100 GPUs per node, the project transferred Kimi K2 BF16 weights from 32 trainer nodes to 16 inference nodes in 7.53 seconds. The synchronization moved 7.9 TB and reached 1,049 GB/s aggregate bandwidth, according to the project’s benchmark. The authors estimate 5 seconds as a more practical lower bound for that topology because layerwise reloading serializes transfers across trainer pipeline-parallel groups.
A smaller Qwen3-235B-A22B test shows why the layout work matters. vLLM reports that the sharded path reduced end-to-end synchronization from 64.72 seconds for an NCCL broadcast to 25.02 seconds in its first implementation, then to 5.61 seconds after pipeline- and expert-local gathering, and finally to 3.49 seconds with pipelining. The accompanying SkyRL examples reproduce the roughly 3.5-second path and document important constraints, including separate trainer and inference placement, Ray 2.56 or newer, and specific cross-node NIXL requirements on AWS.
Why it matters
This is infrastructure for training systems rather than a serving feature for ordinary inference deployments. Its practical value is that model size no longer forces every inference worker to absorb a full broadcast during each online-RL update. The pull-based design also changes failure behavior: the vLLM team demonstrates surviving an inference-engine failure in degraded mode, then allowing the recovered replica to rejoin at the next weight-sync boundary.
The feature is available in vLLM with developer-preview documentation. Teams evaluating it should treat the published numbers as topology-specific engineering results, not a universal expectation, and validate network backend selection, model layout and memory headroom on their own clusters.
sources
comments · 0