vLLM Large-Scale Sharded Weight Transfer with Ray Direct Transport

Overview

vLLM has implemented a native sharded weight-transfer engine leveraging Ray Direct Transport (RDT) to optimize the synchronization of model weights between trainers and inference workers in online Reinforcement Learning (RL) setups. This system replaces traditional NCCL broadcast methods, which often suffer from memory bottlenecks and synchronization stalls at trillion-parameter scales, with a pull-based, sharded approach that reduces peak memory usage and improves transfer speeds.

The Limitations of Broadcast-Based Sync

Standard weight synchronization typically relies on NCCL broadcasts, where a trainer all-gathers parameters into HuggingFace format and broadcasts them to every inference worker. This approach presents two primary challenges for large-scale models:

Memory Inefficiency: In configurations like Tensor Parallelism 8 (TP8), workers receive the full model but only retain 1/8th of the weights, discarding the rest. For large Mixture-of-Experts (MoE) models, this creates massive peak memory overhead.
Collective Synchronization: NCCL requires all ranks to participate synchronously. Straggler ranks or replica failures can stall the entire collective, making it problematic for dynamic, large-scale environments.

Technical Implementation: Sharded Weight Transfer

The Recording Tensor Dry Run

To ensure compatibility across diverse model architectures (such as Llama-4's fused experts or GQA in various models), vLLM uses a "recording tensor" dry run during initialization. vLLM's loaders are provided with a tensor subclass that reports shape and dtype but contains no data. Every transformation—including views, narrows, transposes, and reshapes—is recorded as a chain of operations (a "sharding plan").

This plan allows the trainer to perform the initial layout operations (fusion, relayout, splitting, and sharding) and transfer only the specific sharded weights in BF16 format that each vLLM rank requires, ensuring the process is correct by construction.

Ray Direct Transport (RDT) and NIXL

The engine utilizes Ray Direct Transport (RDT) with a NIXL backend to enable direct GPU-to-GPU communication between Ray actors. This architecture enables a pull-based system where inference ranks pull only the required sharded tensors from mapped trainer ranks.

The initialization flow consists of five steps:

Trainer ranks all-gather ownership metadata (parameter names, dtypes, and shapes).
Rank 0 transmits this metadata and trainer Ray actor names to inference workers.
Each vLLM worker creates its sharding plan via the recording-tensor dry run.
Workers map themselves to source trainer ranks in a load-balanced manner.
Both producers and consumers allocate and register RDT buffers upfront.

Performance Optimizations

vLLM iterated through three versions of the engine to optimize end-to-end latency for a Qwen3-235B-A22B model (TP4/PP2/EP8 trainer to DP16/EP16 vLLM server):

V1 (Simple Iterator): Gathered parameters across all dimensions (TP, PP, EP) one by one. This resulted in thousands of tiny collectives and redundant memory usage, with a sync time of 25.02s.
V2 (PP/EP-Local): Implemented PP-local gathers (only within the same pipeline stage) and EP-local transfers (experts are not gathered at all; inference ranks pull directly from the rank holding the expert). This reduced sync time to 5.61s.
V3 (Pipelined Execution): Introduced overlapping of all-gather, replay operations, and RDMA transfers. By gathering weights in decoder block groups and processing them in the background, sync latency dropped to 3.49s.

Large-Scale Validation: Kimi K2

Validation on the Kimi K2 model across 48 nodes of 8xH100 (32 trainer nodes, 16 inference nodes) demonstrated the following results:

Metric	Value
Bytes moved per sync	7.9 TB
Weight sync time	7.53s
Aggregate bandwidth	1,049 GB/s

This performance is approximately 1.5x the expected "speed of light" (SoL) transfer time for this specific setup, given the constraints of vLLM's layerwise reloading logic.

Fault Tolerance and Integration

By utilizing NIXL instead of broadcast collectives, the system is inherently more resilient to failures. If an inference engine fails, the router continues directing traffic to remaining engines, and the trainer only communicates with live engines during the next sync. Once a failed replica is restored, it rejoins at the next sync boundary and receives updated weights without affecting overall convergence.

Framework Integration

The engine is integrated into SkyRL. Other RL frameworks can adopt it by implementing a WeightSource iterator that provides parameter metadata and yields materialized tensors, with an optional held_names method to enable EP/PP-local optimizations.

Current Limitations

Loader Constraints: Loaders must use recordable operations; those that inspect real values during loading will fail.
Memory Budget: RDT destination buffers exist outside vLLM's gpu_memory_utilization budget.
Compatibility: The current implementation is incompatible with EPLB in vLLM.
Serialization: Transfers are currently serial across trainer PP groups to prevent OOMs during layerwise reloading.