@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 172 lines 8.35 KB
Rendered Raw
1 # Weight Transfer
2
3 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.
4
5 ## Architecture
6
7 Weight transfer has **two engines, one per process**, and they are symmetric:
8
9 | | Trainer process | Inference workers |
10 | --- | --- | --- |
11 | Class | `TrainerWeightTransferEngine` | `WeightTransferEngine` |
12 | Built by | `WeightTransferTrainerFactory.trainer_init(...)` | vLLM, from `WeightTransferConfig` |
13 | Driven by | `send_weights()` | the four-phase protocol below |
14 | Holds | communicators, transfer plans, wire params | communicators, the target model |
15
16 The trainer engine is stateful: it owns its communicators and wire params, pulls
17 weights from a [`WeightSource`](base.md#weightsource), and drives the inference
18 side through a [`VLLMWeightSyncClient`](base.md#vllmweightsyncclient). Trainer
19 code never has to know the transport, and never has to thread transfer state
20 back in on every round — one `send_weights()` call per sync.
21
22 Under the hood every round is the same **four-phase protocol**, which the trainer
23 engine drives on your behalf:
24
25 1. **Initialization** (`init_weight_transfer_engine`): establishes the communication channel between the trainer and inference workers. Called once, from `trainer_init`, before the training loop begins.
26 2. **Start** (`start_weight_update`): prepares the inference engine for a weight update.
27 3. **Weight Update** (`update_weights`): transfers updated weights. May be called one or more times (e.g. for chunked transfers).
28 4. **Finish** (`finish_weight_update`): finalizes the update (e.g. runs post-processing for checkpoint-format weights). Called once after all weights have been transferred.
29
30 ## Available Backends
31
32 | Backend | Transport | Use Case |
33 | ------- | --------- | -------- |
34 | [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
35 | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
36 | [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches |
37 | [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) |
38 | [nccl_m2n](m2n.md) | NCCL M2N reshard | Trainer and inference use different sharding layouts |
39 | [ModelExpress](modelexpress.md) | ModelExpress version-based transfer | Install immutable published weight versions, including object-storage checkpoints and deltas |
40
41 ## Quickstart
42
43 ### Inference Side
44
45 The inference side takes only a backend name. Everything else about the transfer
46 is decided by the trainer and shipped over at the init handshake.
47
48 ```python
49 from vllm import LLM
50 from vllm.config import WeightTransferConfig
51
52 llm = LLM(
53 model="my-model",
54 # Other backends: "ipc", "sparse_nccl", "sharded_rdt", "nccl_m2n".
55 weight_transfer_config=WeightTransferConfig(backend="nccl"),
56 )
57 ```
58
59 Or, for online serving:
60
61 ```bash
62 vllm serve my-model \
63 --weight-transfer-config '{"backend": "nccl"}'
64 ```
65
66 ### Trainer Side
67
68 Build the engine once, then call `send_weights()` once per sync:
69
70 ```python
71 from vllm.distributed.weight_transfer import (
72 ModuleSource,
73 HTTPVLLMWeightSyncClient,
74 WeightTransferTrainerFactory,
75 )
76 from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
77
78 # Once, before the training loop.
79 engine = WeightTransferTrainerFactory.trainer_init(
80 init_info=NCCLTrainerInitInfo(
81 master_address=master_address,
82 master_port=master_port,
83 world_size=world_size, # trainer + all inference workers
84 rank=0, # this trainer rank; rank 0 is the sender
85 packed=True,
86 ),
87 client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
88 source=ModuleSource(model),
89 )
90
91 # Once per weight sync.
92 for step in range(num_steps):
93 train_one_step(model)
94 engine.send_weights()
95 ```
96
97 `send_weights()` drives start → update → finish on the inference side *and* the
98 data-plane transfer, including any concurrency the backend needs (NCCL, for
99 example, must run the workers' `update_weights` at the same time as the
100 trainer's broadcast, since both rendezvous inside the same NCCL calls).
101
102 There is no `backend=` argument on the trainer side: each `TrainerInitInfo`
103 subclass declares its own `backend`, and the factory dispatches on it.
104
105 ## Where Each Setting Lives
106
107 The trainer is the single source of truth for anything the two sides must agree
108 on. This is the main rule to internalize:
109
110 | Setting | Lives on | Notes |
111 | ------- | -------- | ----- |
112 | Backend, inference side | `WeightTransferConfig(backend=...)` | A plain string selector — that is all this config holds |
113 | Backend, trainer side | `init_info.backend` | A `ClassVar` on each init info; never passed by hand |
114 | 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 |
115 | Weights to send | a [`WeightSource`](base.md#weightsource) | `ModuleSource(model)` covers plain and FSDP-sharded modules |
116 | 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 |
117
118 !!! note
119 Wire params are deliberately *not* on `WeightTransferConfig` and *not* on the
120 per-round update info. Because the trainer ships them at init and the worker
121 reads what it was told, a mismatched `packed` flag between the two sides is
122 unrepresentable rather than merely discouraged.
123
124 ## Multi-Rank Trainers
125
126 For sharded trainers (FSDP, TP/PP/EP), **every trainer rank builds an engine and
127 calls `send_weights()`**. Rank 0 is the sender: only it holds a communicator,
128 talks to the client, and puts bytes on the wire. Non-sender ranks still iterate
129 the `WeightSource`, because materializing a parameter is usually itself a
130 collective (an FSDP `full_tensor()` all-gather, a Megatron export) that would
131 deadlock if some ranks skipped it.
132
133 Pass each process's own `rank` on the init info. It is explicit rather than read
134 from a global process group, which is ambiguous once several groups (FSDP / TP /
135 PP / EP) exist.
136
137 ```python
138 engine = WeightTransferTrainerFactory.trainer_init(
139 init_info=NCCLTrainerInitInfo(..., rank=torch.distributed.get_rank()),
140 client=client,
141 source=ModuleSource(model),
142 )
143 engine.send_weights() # called on every rank
144 ```
145
146 ## API Endpoints
147
148 When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first four for you.
149
150 | Endpoint | Method | Description |
151 | -------- | ------ | ----------- |
152 | `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info |
153 | `/start_weight_update` | POST | Start a weight update |
154 | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata |
155 | `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` |
156 | `/update_weight_version` | POST | Update `weight_version` without changing model weights |
157 | `/weight_info` | GET | Get the latest committed weight version |
158 | `/pause` | POST | Pause generation before weight sync to handle inflight requests |
159 | `/resume` | POST | Resume generation after weight sync |
160 | `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) |
161
162 !!! note
163 The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set.
164
165 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, sparse-NCCL, or sharded-RDT transport.
166
167 ## Extending the System
168
169 Every piece of the system is replaceable: the weights you send (`WeightSource`),
170 the control-plane transport (`VLLMWeightSyncClient`), and the transports
171 themselves (both engine ABCs, each with its own factory registry). See
172 [Base Classes](base.md).