Signed-off-by: Connor Carpenter <[email protected]> Co-authored-by: Nick Hill <[email protected]>
7.9 KiB
Weight Transfer
vLLM provides a pluggable weight transfer system for synchronizing model weights from a training process to the inference engine during reinforcement learning (RL) workflows. This is essential for RLHF, GRPO, and other online RL methods where the policy model is iteratively updated during training and the updated weights must be reflected in the inference engine for rollout generation.
Architecture
Weight transfer has two engines, one per process, and they are symmetric:
| Trainer process | Inference workers | |
|---|---|---|
| Class | TrainerWeightTransferEngine |
WeightTransferEngine |
| Built by | WeightTransferTrainerFactory.trainer_init(...) |
vLLM, from WeightTransferConfig |
| Driven by | send_weights() |
the four-phase protocol below |
| Holds | communicators, transfer plans, wire params | communicators, the target model |
The trainer engine is stateful: it owns its communicators and wire params, pulls
weights from a WeightSource, and drives the inference
side through a VLLMWeightSyncClient. Trainer
code never has to know the transport, and never has to thread transfer state
back in on every round — one send_weights() call per sync.
Under the hood every round is the same four-phase protocol, which the trainer engine drives on your behalf:
- Initialization (
init_weight_transfer_engine): establishes the communication channel between the trainer and inference workers. Called once, fromtrainer_init, before the training loop begins. - Start (
start_weight_update): prepares the inference engine for a weight update. - Weight Update (
update_weights): transfers updated weights. May be called one or more times (e.g. for chunked transfers). - Finish (
finish_weight_update): finalizes the update (e.g. runs post-processing for checkpoint-format weights). Called once after all weights have been transferred.
Available Backends
| Backend | Transport | Use Case |
|---|---|---|
| NCCL | NCCL broadcast | Separate GPUs for training and inference |
| IPC | CUDA IPC handles | Colocated training and inference on same GPU |
| sparse_nccl | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
Quickstart
Inference Side
The inference side takes only a backend name. Everything else about the transfer is decided by the trainer and shipped over at the init handshake.
from vllm import LLM
from vllm.config import WeightTransferConfig
llm = LLM(
model="my-model",
weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc", "sparse_nccl"
)
Or, for online serving:
vllm serve my-model \
--weight-transfer-config '{"backend": "nccl"}'
Trainer Side
Build the engine once, then call send_weights() once per sync:
from vllm.distributed.weight_transfer import (
ModuleSource,
HTTPVLLMWeightSyncClient,
WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
# Once, before the training loop.
engine = WeightTransferTrainerFactory.trainer_init(
init_info=NCCLTrainerInitInfo(
master_address=master_address,
master_port=master_port,
world_size=world_size, # trainer + all inference workers
rank=0, # this trainer rank; rank 0 is the sender
packed=True,
),
client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
source=ModuleSource(model),
)
# Once per weight sync.
for step in range(num_steps):
train_one_step(model)
engine.send_weights()
send_weights() drives start → update → finish on the inference side and the
data-plane transfer, including any concurrency the backend needs (NCCL, for
example, must run the workers' update_weights at the same time as the
trainer's broadcast, since both rendezvous inside the same NCCL calls).
There is no backend= argument on the trainer side: each TrainerInitInfo
subclass declares its own backend, and the factory dispatches on it.
Where Each Setting Lives
The trainer is the single source of truth for anything the two sides must agree on. This is the main rule to internalize:
| Setting | Lives on | Notes |
|---|---|---|
| Backend, inference side | WeightTransferConfig(backend=...) |
A plain string selector — that is all this config holds |
| Backend, trainer side | init_info.backend |
A ClassVar on each init info; never passed by hand |
Wire params (packed, buffer sizes) |
the backend's TrainerInitInfo |
The trainer propagates them to the worker during trainer_init, so the two sides cannot disagree |
| Weights to send | a WeightSource |
ModuleSource(model) covers plain and FSDP-sharded modules |
| Per-round payload | the backend's update info | Built by the engine (from source.metadata(), or from a delta backend's per-round payload); you never construct it |
!!! note
Wire params are deliberately not on WeightTransferConfig and not on the
per-round update info. Because the trainer ships them at init and the worker
reads what it was told, a mismatched packed flag between the two sides is
unrepresentable rather than merely discouraged.
Multi-Rank Trainers
For sharded trainers (FSDP, TP/PP/EP), every trainer rank builds an engine and
calls send_weights(). Rank 0 is the sender: only it holds a communicator,
talks to the client, and puts bytes on the wire. Non-sender ranks still iterate
the WeightSource, because materializing a parameter is usually itself a
collective (an FSDP full_tensor() all-gather, a Megatron export) that would
deadlock if some ranks skipped it.
Pass each process's own rank on the init info. It is explicit rather than read
from a global process group, which is ambiguous once several groups (FSDP / TP /
PP / EP) exist.
engine = WeightTransferTrainerFactory.trainer_init(
init_info=NCCLTrainerInitInfo(..., rank=torch.distributed.get_rank()),
client=client,
source=ModuleSource(model),
)
engine.send_weights() # called on every rank
API Endpoints
When running vLLM as an HTTP server, the following endpoints are available for weight transfer. HTTPVLLMWeightSyncClient speaks the first four for you.
| Endpoint | Method | Description |
|---|---|---|
/init_weight_transfer_engine |
POST | Initialize the weight transfer engine with backend-specific info |
/start_weight_update |
POST | Start a weight update |
/update_weights |
POST | Transfer a batch of weights with backend-specific metadata |
/finish_weight_update |
POST | Finish the update and optionally commit its weight_version |
/update_weight_version |
POST | Update weight_version without changing model weights |
/weight_info |
GET | Get the latest committed weight version |
/pause |
POST | Pause generation before weight sync to handle inflight requests |
/resume |
POST | Resume generation after weight sync |
/get_world_size |
GET | Get the number of inference workers (useful for NCCL world size calculation) |
!!! note
The HTTP weight transfer endpoints require VLLM_SERVER_DEV_MODE=1 to be set.
The Rust frontend's optional gRPC Control service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The ServerInfo.rl_capabilities response reports whether weight transfer and sleep mode were configured. Backend-specific init_info and update_info remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, or sparse-NCCL transport.
Extending the System
Every piece of the system is replaceable: the weights you send (WeightSource),
the control-plane transport (VLLMWeightSyncClient), and the transports
themselves (both engine ABCs, each with its own factory registry). See
Base Classes.