[Docs][RL] Rewrite weight-transfer docs; standardize examples (#51729)
Signed-off-by: haoaaron <ahao@anyscale.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Aaron Hao committed
Aug 11, 2026 at 18:53 UTC
793ca6998adfa5a3ab22d6012dee78044b8ba901
16 files changed
+1226
-1157
.buildkite/test-amd.yaml
+2
-2
@@ -847,8 +847,8 @@ steps:
847
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
848
- python3 ../examples/features/data_parallel/data_parallel_offline.py --enforce-eager
849
# rlhf examples
850
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_nccl.py
851
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_ipc.py
850
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_nccl.py
851
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_ipc.py
852
853
- label: Elastic EP Scaling Test # TBD
854
timeout_in_minutes: 180
.buildkite/test_areas/distributed.yaml
+3
-3
@@ -133,9 +133,9 @@ steps:
133
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
134
# test with internal dp
135
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
136
- # rlhf examples
137
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
138
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
136
+ # rlhf examples (each launches its own `vllm serve` and tears it down)
137
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
138
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
139
140
- label: Distributed DP Tests (4 GPUs)
141
key: distributed-dp-tests-4-gpus
docs/training/weight_transfer/README.md
+111
-36
@@ -4,12 +4,28 @@ vLLM provides a pluggable weight transfer system for synchronizing model weights
4
5
## Architecture
6
7
-The weight transfer system follows a **four-phase protocol** with a pluggable backend design:
7
+Weight transfer has **two engines, one per process**, and they are symmetric:
8
9
-1. **Initialization** (`init_weight_transfer_engine`): Establishes the communication channel between the trainer and inference workers. Called once before the training loop begins.
10
-2. **Start** (`start_weight_update`): Prepares the inference engine for a weight update.
11
-3. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. May be called one or more times (e.g., for chunked transfers).
12
-4. **Finish** (`finish_weight_update`): Finalizes the weight update (e.g., runs post-processing for checkpoint-format weights). Called once after all weights have been transferred.
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
@@ -19,11 +35,12 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba
35
| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
36
| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
37
22
-## Configuration
38
+## Quickstart
39
24
-Specify the weight transfer backend through `WeightTransferConfig`. The backend determines which engine handles the weight synchronization.
40
+### Inference Side
41
26
-### Programmatic (Offline Inference)
42
+The inference side takes only a backend name. Everything else about the transfer
43
+is decided by the trainer and shipped over at the init handshake.
44
45
```python
46
from vllm import LLM
@@ -31,22 +48,100 @@ from vllm.config import WeightTransferConfig
48
49
llm = LLM(
50
model="my-model",
34
- weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc"
51
+ weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc", "sparse_nccl"
52
)
53
```
54
38
-### CLI (Online Serving)
55
+Or, for online serving:
56
57
```bash
58
vllm serve my-model \
59
--weight-transfer-config '{"backend": "nccl"}'
60
```
61
45
-The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
62
+### Trainer Side
63
+
64
+Build the engine once, then call `send_weights()` once per sync:
65
+
66
+```python
67
+from vllm.distributed.weight_transfer import (
68
+ ModuleSource,
69
+ HTTPVLLMWeightSyncClient,
70
+ WeightTransferTrainerFactory,
71
+)
72
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
73
+
74
+# Once, before the training loop.
75
+engine = WeightTransferTrainerFactory.trainer_init(
76
+ init_info=NCCLTrainerInitInfo(
77
+ master_address=master_address,
78
+ master_port=master_port,
79
+ world_size=world_size, # trainer + all inference workers
80
+ rank=0, # this trainer rank; rank 0 is the sender
81
+ packed=True,
82
+ ),
83
+ client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
84
+ source=ModuleSource(model),
85
+)
86
+
87
+# Once per weight sync.
88
+for step in range(num_steps):
89
+ train_one_step(model)
90
+ engine.send_weights()
91
+```
92
+
93
+`send_weights()` drives start → update → finish on the inference side *and* the
94
+data-plane transfer, including any concurrency the backend needs (NCCL, for
95
+example, must run the workers' `update_weights` at the same time as the
96
+trainer's broadcast, since both rendezvous inside the same NCCL calls).
97
+
98
+There is no `backend=` argument on the trainer side: each `TrainerInitInfo`
99
+subclass declares its own `backend`, and the factory dispatches on it.
100
+
101
+## Where Each Setting Lives
102
+
103
+The trainer is the single source of truth for anything the two sides must agree
104
+on. This is the main rule to internalize:
105
+
106
+| Setting | Lives on | Notes |
107
+| ------- | -------- | ----- |
108
+| Backend, inference side | `WeightTransferConfig(backend=...)` | A plain string selector — that is all this config holds |
109
+| Backend, trainer side | `init_info.backend` | A `ClassVar` on each init info; never passed by hand |
110
+| 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 |
111
+| Weights to send | a [`WeightSource`](base.md#weightsource) | `ModuleSource(model)` covers plain and FSDP-sharded modules |
112
+| 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 |
113
+
114
+!!! note
115
+ Wire params are deliberately *not* on `WeightTransferConfig` and *not* on the
116
+ per-round update info. Because the trainer ships them at init and the worker
117
+ reads what it was told, a mismatched `packed` flag between the two sides is
118
+ unrepresentable rather than merely discouraged.
119
+
120
+## Multi-Rank Trainers
121
+
122
+For sharded trainers (FSDP, TP/PP/EP), **every trainer rank builds an engine and
123
+calls `send_weights()`**. Rank 0 is the sender: only it holds a communicator,
124
+talks to the client, and puts bytes on the wire. Non-sender ranks still iterate
125
+the `WeightSource`, because materializing a parameter is usually itself a
126
+collective (an FSDP `full_tensor()` all-gather, a Megatron export) that would
127
+deadlock if some ranks skipped it.
128
+
129
+Pass each process's own `rank` on the init info. It is explicit rather than read
130
+from a global process group, which is ambiguous once several groups (FSDP / TP /
131
+PP / EP) exist.
132
+
133
+```python
134
+engine = WeightTransferTrainerFactory.trainer_init(
135
+ init_info=NCCLTrainerInitInfo(..., rank=torch.distributed.get_rank()),
136
+ client=client,
137
+ source=ModuleSource(model),
138
+)
139
+engine.send_weights() # called on every rank
140
+```
141
142
## API Endpoints
143
49
-When running vLLM as an HTTP server, the following endpoints are available for weight transfer:
144
+When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first four for you.
145
146
| Endpoint | Method | Description |
147
| -------- | ------ | ----------- |
@@ -63,29 +158,9 @@ When running vLLM as an HTTP server, the following endpoints are available for w
158
!!! note
159
The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set.
160
66
-## Trainer-Side API
67
-
68
-Both backends provide static methods that the trainer calls to send weights. The general pattern is:
69
-
70
-```python
71
-# 1. Initialize the transfer engine (backend-specific)
72
-EngineClass.trainer_init(init_info)
73
-
74
-# 2. Start weight update on inference side
75
-llm.start_weight_update()
76
-
77
-# 3. Send weights to inference workers
78
-EngineClass.trainer_send_weights(
79
- iterator=model.named_parameters(),
80
- trainer_args=backend_specific_args,
81
-)
82
-
83
-# 4. Finish weight update on inference side
84
-llm.finish_weight_update(weight_version="step-42")
85
-```
86
-
87
-See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples.
88
-
161
## Extending the System
162
91
-The weight transfer system is designed to be extensible. You can implement custom backends by subclassing `WeightTransferEngine` and registering them with the factory. See the [Base Class](base.md) page for details.
163
+Every piece of the system is replaceable: the weights you send (`WeightSource`),
164
+the control-plane transport (`VLLMWeightSyncClient`), and the transports
165
+themselves (both engine ABCs, each with its own factory registry). See
166
+[Base Classes](base.md).
docs/training/weight_transfer/base.md
+396
-59
@@ -1,36 +1,391 @@
1
-# Base Class and Custom Engines
1
+# Base Classes and Custom Engines
2
3
-The weight transfer system is built on an abstract base class that defines the contract between vLLM's worker infrastructure and the transport backend. You can implement custom backends by subclassing `WeightTransferEngine` and registering them with the `WeightTransferEngineFactory`.
3
+The weight transfer system is built from four abstractions, each independently
4
+replaceable:
5
5
-## WeightTransferEngine
6
+| Abstraction | Side | Answers |
7
+| ----------- | ---- | ------- |
8
+| [`WeightSource`](#weightsource) | Trainer | *What* weights to send |
9
+| [`VLLMWeightSyncClient`](#vllmweightsyncclient) | Trainer | *How to reach* the inference engine — the adapter for your RL stack's own vLLM wrapper |
10
+| [`TrainerWeightTransferEngine`](#trainerweighttransferengine) | Trainer | *How to transmit* the bytes |
11
+| [`WeightTransferEngine`](#weighttransferengine) | Inference | *How to receive* them and load them |
12
7
-The `WeightTransferEngine` is a generic abstract class parameterized by two dataclass types:
13
+The two engines are registered in two separate factories,
14
+[`WeightTransferTrainerFactory`](#weighttransfertrainerfactory) and
15
+[`WeightTransferEngineFactory`](#weighttransferenginefactory). They share backend
16
+names by convention, but a trainer process never instantiates a worker engine or
17
+vice versa, so the registries stay independent.
18
9
-- **`TInitInfo`** (extends `WeightTransferInitInfo`): Backend-specific initialization parameters.
10
-- **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): Backend-specific weight update metadata.
19
+## Trainer Side
20
12
-### Abstract Methods
21
+### WeightSource
22
14
-Subclasses must implement these methods:
23
+A `WeightSource` is a **re-iterable** source of the trainer's weights with two
24
+channels:
25
16
-| Method | Side | Description |
17
-| ------ | ---- | ----------- |
18
-| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
19
-| `start_weight_update()` | Inference | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
20
-| `finish_weight_update()` | Inference | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
21
-| `receive_weights(update_info)` | Inference | Receive weights and load them into `self.model` |
22
-| `shutdown()` | Inference | Clean up resources |
23
-| `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
26
+- **`metadata() -> list[ParamMeta]`** — the name, wire dtype, and full shape of
27
+ every parameter, *without transferring anything*. Cheap when shapes are known
28
+ locally (an FSDP `DTensor` knows its global shape); may be expensive on the
29
+ first call for producers that must materialize to learn shapes (a
30
+ Megatron-Bridge export), in which case it should cache.
31
+- **iteration** — yields fully-materialized `(name, tensor)` pairs, one at a
32
+ time.
33
25
-The base class provides two methods:
34
+```python
35
+@dataclass(frozen=True)
36
+class ParamMeta:
37
+ name: str
38
+ dtype: torch.dtype
39
+ shape: tuple[int, ...]
40
+```
41
+
42
+!!! warning "The two channels must agree, element for element"
43
+ `metadata()` must declare exactly what iteration will yield: the same
44
+ parameters, in the same order, with the same dtypes and shapes. This is an
45
+ invariant of the ABC, not of any one backend — a source that reorders, omits,
46
+ or re-dtypes a parameter between the two channels is broken even if the
47
+ backend you happen to test against never notices. Backends are free to read
48
+ both channels and to trust that they match; dense NCCL does, and
49
+ [enforces it](nccl.md#the-two-weightsource-channels-must-agree).
50
+
51
+Materializing is typically a collective, so **every trainer rank must iterate the
52
+same source in the same order, in lockstep**, or ranks deadlock. `metadata()` can
53
+itself be a collective for custom producers, so it too runs on every rank — only
54
+the sender ships the result. Under pipeline parallelism a rank may not own a
55
+parameter at all; iterating still drives the collective, and the yielded tensor
56
+is only meaningful on the sender.
57
+
58
+`iter(source)` must yield a *fresh* pass each round.
59
+
60
+#### ModuleSource
61
+
62
+`ModuleSource(module)` is the common case, over `module.named_parameters()`. It
63
+handles plain and FSDP-sharded modules with no special casing: iteration
64
+all-gathers each `DTensor` via `full_tensor()`, while `metadata()` reads the
65
+*global* `.shape` / `.dtype` and so never triggers a gather.
66
+
67
+```python
68
+from vllm.distributed.weight_transfer import ModuleSource
69
+
70
+source = ModuleSource(model)
71
+```
72
+
73
+#### Custom Sources
74
+
75
+Subclass `WeightSource` when the weights you want to send require additional processing to convert to a HF compatible format.
76
+
77
+```python
78
+from vllm.distributed.weight_transfer import ParamMeta, WeightSource
79
+from vllm.distributed.weight_transfer.base import materialize_full_tensor
80
+
81
+
82
+class MyExportSource(WeightSource):
83
+ def __init__(self, model):
84
+ self._model = model
85
+ self._meta: list[ParamMeta] | None = None
86
+
87
+ def metadata(self) -> list[ParamMeta]:
88
+ # Cache: for producers that must materialize to learn shapes, this is
89
+ # the expensive channel. Runs on every rank (it may be a collective).
90
+ if self._meta is None:
91
+ self._meta = [
92
+ ParamMeta(name, t.dtype, tuple(t.shape)) for name, t in self._export()
93
+ ]
94
+ return self._meta
95
+
96
+ def __iter__(self):
97
+ # Must yield exactly what metadata() declared, in the same order.
98
+ for name, tensor in self._export():
99
+ yield name, materialize_full_tensor(tensor)
100
+```
101
+
102
+### VLLMWeightSyncClient
103
+
104
+**This is the adapter for however your RL stack reaches vLLM.** Many RL frameworks wrap
105
+inference engines in their own abstractions, and each reaches vLLM its own way.
106
+`VLLMWeightSyncClient` is the single seam where that bespoke shape is adapted, so
107
+weight sync engines remain control plane agnostic.
108
+
109
+The contract is only this: **however the wrapper is shaped, it must bottom out in
110
+the same four calls** — `init_weight_transfer_engine` once at setup, then
111
+`start_weight_update` → one or more `update_weights` → `finish_weight_update` per
112
+round. Everything a trainer engine needs from the inference side goes through them.
113
+
114
+```python
115
+class VLLMWeightSyncClient(Protocol):
116
+ def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ...
117
+ def start_weight_update(self) -> None: ...
118
+ def update_weights(self, update_info: dict[str, Any]) -> None: ...
119
+ def finish_weight_update(self, weight_version: str | None = None) -> None: ...
120
+```
121
+
122
+It is a `@runtime_checkable` structural `Protocol` (PEP 544), which is what makes
123
+adapting cheap: **any object with those four methods already satisfies it**. An existing
124
+wrapper in your framework can usually become a client by gaining four forwarding
125
+methods.
126
27
-1. `__init__` : Engines receive `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`) and `model` (`nn.Module`)
28
-2. `update_weights(update_info_dict)`: Thin wrapper for `receive_weights`: parses
29
-the dict into user-specified data type, calls `receive_weights`, and synchronizes the device. Subclasses implement `receive_weights`.
127
+Two implementations ship with vLLM:
128
+
129
+| Client | Talks to |
130
+| ------ | -------- |
131
+| `RayVLLMWeightSyncClient(handle)` | One or more `AsyncLLM`/`LLM` Ray actors. Accepts a list and fans each call out to every handle, blocking on all of them, so a multi-actor (e.g. multi-DP) deployment is driven as one unit |
132
+| `HTTPVLLMWeightSyncClient(base_url, timeout=300)` | A vLLM server over the RLHF HTTP routes |
133
+
134
+Custom weight sync clients can be implement like so:
135
+
136
+```python
137
+class MyFrameworkWeightSyncClient:
138
+ """Adapts one RL framework's rollout pool to the four weight-sync calls."""
139
+
140
+ def __init__(self, rollout_pool):
141
+ self.pool = rollout_pool # whatever your stack already has
142
+
143
+ def init_weight_transfer_engine(self, init_info):
144
+ # Fan out to every replica and block: all of them receive weights.
145
+ self.pool.broadcast_rpc("init_weight_transfer_engine", init_info=init_info)
146
+
147
+ def start_weight_update(self):
148
+ self.pool.broadcast_rpc("start_weight_update")
149
+
150
+ def update_weights(self, update_info):
151
+ self.pool.broadcast_rpc("update_weights", update_info=update_info)
152
+
153
+ def finish_weight_update(self, weight_version=None):
154
+ self.pool.broadcast_rpc("finish_weight_update")
155
+ if weight_version is not None:
156
+ self.pool.broadcast_rpc("update_weight_version", weight_version)
157
+```
158
+
159
+Two things to get right in any adapter:
160
+
161
+- **Reach every replica, and block until all of them are done.** A weight update
162
+ is not a load-balanced request: every worker holding a copy of the model must
163
+ receive it. Returning before they all finish lets the trainer race ahead of
164
+ workers still loading. (Both built-in clients do this — Ray by fanning out over
165
+ its handles, HTTP because the server's DP client broadcasts internally.)
166
+- **Raise on failure.** Trainer engines rely on exceptions to surface
167
+ inference-side errors; a client that swallows them turns a failed sync into
168
+ silently stale weights, or into a hang for backends whose transfer rendezvouses
169
+ with the worker.
170
+
171
+!!! note
172
+ HTTP cannot carry raw CUDA IPC handles, so `HTTPVLLMWeightSyncClient` pickles
173
+ and base64-encodes them into an `ipc_handles_pickled` field. The worker
174
+ deserializes it only when `VLLM_ALLOW_INSECURE_SERIALIZATION=1`. Backends
175
+ whose payloads are JSON-native (NCCL) pass through untouched.
176
+
177
+### TrainerWeightTransferEngine
178
+
179
+The trainer-side engine: it holds the transport state (NCCL communicators, IPC
180
+device info, transfer plans), pulls weights from a `WeightSource`, and drives the
181
+inference side through a `VLLMWeightSyncClient`. It is generic over its init info
182
+type, constructed by the `trainer_init` classmethod factory, and driven by `send_weights()`.
183
+
184
+| Method | Description |
185
+| ------ | ----------- |
186
+| `trainer_init(init_info, *, client, source=None)` | Classmethod. Rendezvous with the inference side and return a ready instance |
187
+| `send_weights()` | Push weights and drive the full update round trip |
188
+| `shutdown()` | Tear down communicators / process groups. Default no-op |
189
+
190
+Both `trainer_init` and `send_weights` are called on **every** trainer rank.
191
+`is_sender` is resolved once, at `trainer_init`, from `init_info.rank`. Each
192
+engine holds the real client on every rank but guards the control-plane RPCs and
193
+the transmit on `self.is_sender`, so only the sender touches the wire; non-sender
194
+ranks still run every collective so the group stays aligned.
195
+
196
+The trainer side takes **no `WeightTransferConfig`**. The backend comes from the
197
+init info's `backend` `ClassVar`, and the wire params ride the init info too.
198
+
199
+#### TrainerInitInfo
200
+
201
+The `init_info` passed to `trainer_init` above. It is how a caller configures a
202
+transfer: it selects the backend, says which rank this process is, and carries the
203
+wire params. Each backend subclasses it; the base class holds the one field every
204
+backend needs.
205
+
206
+```python
207
+@dataclass
208
+class TrainerInitInfo:
209
+ backend: ClassVar[str] # factory dispatch key
210
+ rank: int = field(kw_only=True)
211
+
212
+ @property
213
+ def is_sender(self) -> bool:
214
+ return self.rank == 0
215
+```
216
+
217
+- **`rank`** is this trainer process's rank, supplied **explicitly**. The engine
218
+ does not read it from a global process group, which is ambiguous once several
219
+ groups (FSDP / TP / PP / EP) exist. **Rank 0 is always the sender** — this is
220
+ what `trainer_init` resolves into `is_sender`. It is keyword-only, so backend
221
+ subclasses can add positional fields freely.
222
+- **`backend`** is a `ClassVar`, not an `__init__` field: it is a fixed
223
+ per-backend constant that the factory reads to dispatch, which is why callers
224
+ never pass a `backend=` argument. Every subclass must set it —
225
+ `__init_subclass__` raises otherwise.
226
+
227
+Subclasses also carry the transfer's **wire params** (`packed`, buffer sizes).
228
+The sender propagates them to the worker inside `trainer_init`, so the two sides
229
+cannot disagree. See [`NCCLTrainerInitInfo`](nccl.md#nccltrainerinitinfo) and
230
+[`IPCTrainerInitInfo`](ipc.md#ipctrainerinitinfo) for the concrete fields.
231
+
232
+#### Full-Resync vs. Delta Backends
233
+
234
+`source` is optional, which splits the backends into two shapes:
235
+
236
+- **Full resync** (NCCL, IPC) — a stable `WeightSource` is fixed at
237
+ `trainer_init` and re-iterated each round; `send_weights()` takes no
238
+ arguments. These backends validate that `source` is non-null themselves.
239
+- **Delta** (sparse NCCL) — the payload differs every round, so there is no
240
+ stable source. The engine takes no `source` and each round's payload is passed
241
+ straight to `send_weights(patches)`.
242
+
243
+#### Implementing a Custom Trainer Engine
244
+
245
+```python
246
+from dataclasses import dataclass
247
+from typing import ClassVar
248
+
249
+from typing_extensions import Self
250
+
251
+from vllm.distributed.weight_transfer.base import (
252
+ TrainerInitInfo,
253
+ TrainerWeightTransferEngine,
254
+ VLLMWeightSyncClient,
255
+ WeightSource,
256
+)
257
+
258
+
259
+@dataclass
260
+class MyTrainerInitInfo(TrainerInitInfo):
261
+ backend: ClassVar[str] = "my_backend"
262
+
263
+ endpoint: str
264
+ chunk_size_bytes: int = 256 * 1024 * 1024 # a wire param: shipped to the worker
265
+
266
+
267
+class MyTrainerWeightTransferEngine(TrainerWeightTransferEngine[MyTrainerInitInfo]):
268
+ init_info_cls = MyTrainerInitInfo
269
+
270
+ def __init__(self, *, client, source, is_sender=True, chunk_size_bytes=0):
271
+ super().__init__(client=client, source=source, is_sender=is_sender)
272
+ self.chunk_size_bytes = chunk_size_bytes
273
+
274
+ @classmethod
275
+ def trainer_init(
276
+ cls,
277
+ init_info: MyTrainerInitInfo,
278
+ *,
279
+ client: VLLMWeightSyncClient,
280
+ source: WeightSource | None = None,
281
+ ) -> Self:
282
+ if source is None:
283
+ raise ValueError("my_backend requires a WeightSource.")
284
+ engine = cls(
285
+ client=client,
286
+ source=source,
287
+ is_sender=init_info.is_sender,
288
+ chunk_size_bytes=init_info.chunk_size_bytes,
289
+ )
290
+ if engine.is_sender:
291
+ # Ship the must-agree wire params so the worker decodes exactly as
292
+ # this trainer encodes, then open the trainer-side endpoint.
293
+ engine.client.init_weight_transfer_engine(
294
+ {"chunk_size_bytes": init_info.chunk_size_bytes}
295
+ )
296
+ return engine
297
+
298
+ def send_weights(self) -> None:
299
+ assert self.source is not None
300
+ meta = self.source.metadata() # every rank: may be a collective
301
+ if not self.is_sender:
302
+ for _ in self.source: # stay in the trainer-side collective
303
+ pass
304
+ return
305
+
306
+ self.client.start_weight_update()
307
+ self.client.update_weights(
308
+ {
309
+ "names": [m.name for m in meta],
310
+ "dtype_names": [str(m.dtype).split(".")[-1] for m in meta],
311
+ "shapes": [list(m.shape) for m in meta],
312
+ }
313
+ )
314
+ for name, tensor in self.source:
315
+ ... # transmit
316
+ self.client.finish_weight_update()
317
+```
318
+
319
+Two things to get right, both of which have bitten the built-in backends:
320
+
321
+- **Drain before returning.** `send_weights` must not return with transfers
322
+ still in flight. Anything keeping a send buffer alive dies with the frame, and
323
+ the inference side's `finish_weight_update` post-processing can otherwise
324
+ finalize weights that have not landed.
325
+- **Never join a control-plane thread on the error path.** If you run
326
+ `update_weights` on a side thread concurrently with a transmit (as NCCL does)
327
+ and the transmit raises, the worker is still blocked in the matching
328
+ collective and will never return. Shut the executor down without waiting, so
329
+ the real exception surfaces instead of hanging.
330
+
331
+### WeightTransferTrainerFactory
332
+
333
+```python
334
+from vllm.distributed.weight_transfer import WeightTransferTrainerFactory
335
+
336
+# Lazy loading (recommended): the module is imported only when the backend is used
337
+WeightTransferTrainerFactory.register_engine(
338
+ "my_backend",
339
+ "my_package.my_module",
340
+ "MyTrainerWeightTransferEngine",
341
+)
342
+
343
+# Or register the class directly
344
+WeightTransferTrainerFactory.register_engine("my_backend", MyTrainerWeightTransferEngine)
345
+
346
+engine = WeightTransferTrainerFactory.trainer_init(
347
+ init_info=MyTrainerInitInfo(rank=0, endpoint="..."), # `backend` selects the engine
348
+ client=client,
349
+ source=source,
350
+)
351
+```
352
+
353
+## Inference Side
354
+
355
+### WeightTransferEngine
356
+
357
+A generic abstract class parameterized by two dataclass types:
358
+
359
+- **`TInitInfo`** (extends `WeightTransferInitInfo`): backend-specific initialization parameters.
360
+- **`TUpdateInfo`** (extends `WeightTransferUpdateInfo`): backend-specific weight update metadata.
361
+
362
+Subclasses must implement five methods:
363
+
364
+| Method | Description |
365
+| ------ | ----------- |
366
+| `init_transfer_engine(init_info)` | Initialize the communication channel on each inference worker, and record the trainer-supplied wire params |
367
+| `start_weight_update()` | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
368
+| `finish_weight_update()` | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
369
+| `receive_weights(update_info)` | Receive weights and load them into `self.model` |
370
+| `shutdown()` | Clean up resources |
371
+
372
+The base class provides:
373
+
374
+1. `__init__`, taking `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`), and `model` (`nn.Module`).
375
+2. `update_weights(update_info_dict)`, a thin wrapper for `receive_weights`: it parses the dict into the typed dataclass, calls `receive_weights`, and synchronizes the device.
376
+3. `parse_init_info` / `parse_update_info`, which convert API-level dicts into the typed dataclasses and raise `ValueError` on a bad payload.
377
+4. `set_weight_update_target` / `reset_weight_update_target`, used to retarget an update at the speculative draft model.
378
+
379
+!!! note "Read wire params from the handshake, not the payload"
380
+ Anything the two sides must agree on — `packed`, buffer geometry — arrives
381
+ on the **init info** and should be stored on `self` in
382
+ `init_transfer_engine`, then read from `self` in `receive_weights`. Per-round
383
+ update info carries only per-round metadata. This is what makes a
384
+ trainer/worker mismatch unrepresentable.
385
386
### Request Classes
387
33
-The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses.
388
+The API-level request classes provide backend-agnostic serialization using plain dictionaries.
389
390
```python
391
from vllm.distributed.weight_transfer.base import (
@@ -49,25 +404,17 @@ update_request = WeightTransferUpdateRequest(
404
)
405
```
406
407
+Using a built-in client, you never construct these by hand — `RayVLLMWeightSyncClient`
408
+wraps the dicts for you, and `HTTPVLLMWeightSyncClient` posts them as JSON.
409
+
410
At the LLM/API layer, call `start_draft_weight_update()` instead of
411
`start_weight_update()` to target the speculative draft model;
54
-`update_weights` / `finish_weight_update` are unchanged.
412
+`update_weights` / `finish_weight_update` are unchanged. Engines that cannot
413
+support this set `supports_draft_weight_update = False`.
414
56
-### WeightTransferUpdateInfo
415
+### Implementing a Custom Engine
416
58
-The base `WeightTransferUpdateInfo` is a marker class for backend-specific update info:
59
-
60
-```python
61
-@dataclass
62
-class WeightTransferUpdateInfo(ABC):
63
- pass
64
-```
65
-
66
-## Implementing a Custom Engine
67
-
68
-To create a custom weight transfer backend:
69
-
70
-### 1. Define Info Dataclasses
417
+#### 1. Define Info Dataclasses
418
419
```python
420
from dataclasses import dataclass
@@ -80,29 +427,26 @@ from vllm.distributed.weight_transfer.base import (
427
@dataclass
428
class MyInitInfo(WeightTransferInitInfo):
429
endpoint: str
83
- token: str
430
+ chunk_size_bytes: int = 256 * 1024 * 1024 # must-agree wire param
431
432
@dataclass
433
class MyUpdateInfo(WeightTransferUpdateInfo):
434
names: list[str]
435
dtype_names: list[str]
436
shapes: list[list[int]]
90
- # Add custom fields as needed
437
+ # Per-round metadata only.
438
```
439
93
-### 2. Implement the Engine
440
+#### 2. Implement the Engine
441
442
```python
96
-from collections.abc import Iterator
97
-from typing import Any
98
-import torch
99
-
443
class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
444
init_info_cls = MyInitInfo
445
update_info_cls = MyUpdateInfo
446
447
def init_transfer_engine(self, init_info: MyInitInfo) -> None:
105
- # Set up connection to trainer using init_info.endpoint, etc.
448
+ # Record the trainer's wire params, then set up the connection.
449
+ self.chunk_size_bytes = init_info.chunk_size_bytes
450
...
451
452
def start_weight_update(self) -> None:
@@ -128,22 +472,12 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
472
def shutdown(self) -> None:
473
# Clean up resources
474
...
131
-
132
- @staticmethod
133
- def trainer_send_weights(
134
- iterator: Iterator[tuple[str, torch.Tensor]],
135
- trainer_args: dict[str, Any],
136
- ) -> None:
137
- # Send weights from the trainer process
138
- for name, tensor in iterator:
139
- # Send tensor via custom transport
140
- ...
475
```
476
143
-### 3. Register with the Factory
477
+#### 3. Register with the Factory
478
479
```python
146
-from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
480
+from vllm.distributed.weight_transfer import WeightTransferEngineFactory
481
482
# Option 1: Lazy loading (recommended for built-in engines)
483
WeightTransferEngineFactory.register_engine(
@@ -159,14 +493,14 @@ WeightTransferEngineFactory.register_engine(
493
)
494
```
495
162
-Once registered, users can select your backend via `WeightTransferConfig(backend="my_backend")`.
496
+Once registered, users select your backend via `WeightTransferConfig(backend="my_backend")`.
497
164
-## WeightTransferEngineFactory
498
+### WeightTransferEngineFactory
499
500
The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, and `sparse_nccl`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
501
502
```python
169
-from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
503
+from vllm.distributed.weight_transfer import WeightTransferEngineFactory
504
505
# Create an engine from config
506
engine = WeightTransferEngineFactory.create_engine(
@@ -176,3 +510,6 @@ engine = WeightTransferEngineFactory.create_engine(
510
model=model,
511
)
512
```
513
+
514
+vLLM calls this for you during worker startup; you only need it directly when
515
+embedding the engine in your own worker.
docs/training/weight_transfer/ipc.md
+74
-76
@@ -8,111 +8,109 @@ The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) h
8
9
## How It Works
10
11
-1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`. In multi-GPU setups (e.g. FSDP), each trainer rank must all-gather the full tensor for each layer onto its own GPU before generating the IPC handle.
12
-2. IPC handles for each gpu are sent to the inference engine via **Ray**, **HTTP**, or a **custom callable**. Each rank only reads the handle corresponding to its own GPU.
11
+1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`. In multi-GPU setups (e.g. FSDP), each trainer rank materializes the full tensor for each parameter onto its own GPU before generating the handle — which `ModuleSource` does for you.
12
+2. Every trainer rank contributes its handles to an all-gather; the sender merges them so each payload maps every GPU UUID to its args, and ships the merged handles to the inference engine through the client. Each worker reads only the handle for its own GPU.
13
3. The inference worker reconstructs the tensors from the handles using `rebuild_cuda_tensor`, reading directly from the trainer's GPU memory.
14
15
-!!! warning
16
- IPC handles involve sending serialized Python objects. When using HTTP transport, you must set `VLLM_ALLOW_INSECURE_SERIALIZATION=1` on both the server and client. This is because IPC handles are pickled and base64-encoded for HTTP transmission.
15
+Unlike NCCL, IPC transfer is straight-line: `update_weights` *is* the transfer,
16
+and it rides the client, so there is no concurrent broadcast to overlap with.
17
18
-## Packed (Chunked) Transfer
18
+!!! note
19
+ The handle all-gather in step 2 runs over the **default** process group, so it
20
+ assumes that group is exactly the set of colocated trainer ranks and that the
21
+ sender is a member. It is a no-op when no distributed group exists.
22
20
-By default, all weights are sent in a single API call. For large models, this requires the full model to reside in GPU memory on both sides simultaneously. Setting `packed=True` enables **chunked transfer** with bounded GPU memory:
23
+!!! warning
24
+ IPC handles involve sending serialized Python objects. When using HTTP transport, you must set `VLLM_ALLOW_INSECURE_SERIALIZATION=1` on both the server and client. This is because IPC handles are pickled and base64-encoded for HTTP transmission.
25
22
-- Weights are concatenated into fixed-size packed buffers (controlled by `packed_buffer_size_bytes`).
23
-- Each chunk is sent as a separate `update_weights` call within a single `start_weight_update` / `finish_weight_update` bracket, so the layerwise reload pass is initialized once at the start and finalized once at the end regardless of chunk count.
24
-- After each chunk is consumed, the GPU memory for that chunk can be reclaimed.
26
+## Inference Side
27
28
```python
27
-trainer_args = IPCTrainerSendWeightsArgs(
28
- send_mode="ray",
29
- llm_handle=llm_actor_handle,
30
- packed=True,
31
- packed_buffer_size_bytes=256 * 1024 * 1024, # 256 MB chunks
32
-)
33
-```
34
-
35
-## Initialization
29
+from vllm import LLM
30
+from vllm.config import WeightTransferConfig
31
37
-The IPC backend requires no initialization on either side. The `init_transfer_engine` call is a no-op for IPC.
38
-
39
-## Sending Weights
32
+llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="ipc"))
33
+```
34
41
-IPC supports two transport modes for delivering the handles:
35
+```bash
36
+vllm serve my-model --weight-transfer-config '{"backend": "ipc"}'
37
+```
38
43
-### Ray Mode
39
+IPC needs no data-plane rendezvous, so `init_transfer_engine` opens no channel —
40
+it only records the `packed` flag the trainer ships at the handshake, which
41
+`receive_weights` then reads. Whether a transfer is packed is therefore never
42
+something you configure on the inference side.
43
45
-Used when vLLM is running as a Ray actor:
44
+## Trainer Side
45
46
```python
48
-from vllm.distributed.weight_transfer.ipc_engine import (
49
- IPCTrainerSendWeightsArgs,
50
- IPCWeightTransferEngine,
47
+from vllm.distributed.weight_transfer import (
48
+ ModuleSource,
49
+ HTTPVLLMWeightSyncClient,
50
+ WeightTransferTrainerFactory,
51
)
52
+from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
53
53
-trainer_args = IPCTrainerSendWeightsArgs(
54
- send_mode="ray",
55
- llm_handle=llm_actor_handle,
54
+engine = WeightTransferTrainerFactory.trainer_init(
55
+ init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 is the sender
56
+ client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
57
+ source=ModuleSource(model),
58
)
57
-# start
58
-ray.get(llm_actor_handle.start_weight_update.remote())
59
-# send weights
60
-IPCWeightTransferEngine.trainer_send_weights(
61
- iterator=model.named_parameters(),
62
- trainer_args=trainer_args,
63
-)
64
-# finish
65
-ray.get(llm_actor_handle.finish_weight_update.remote())
59
+
60
+engine.send_weights() # once per sync
61
```
62
68
-In Ray mode, the engine calls `llm_handle.update_weights.remote(...)` directly, passing the IPC handles via Ray's serialization.
63
+`send_weights()` drives `start_weight_update`, the transfer itself, and
64
+`finish_weight_update`, and holds strong references to the IPC-shared copies
65
+until after the post-send barrier — the consumer's views would otherwise dangle.
66
70
-### HTTP Mode
67
+Any [`VLLMWeightSyncClient`](base.md#vllmweightsyncclient) works here — the
68
+built-in HTTP and Ray clients, or an adapter for your own stack. The engine is
69
+identical either way.
70
72
-Used when vLLM is running as an HTTP server:
71
+### `IPCTrainerInitInfo`
72
74
-```python
75
-trainer_args = IPCTrainerSendWeightsArgs(
76
- send_mode="http",
77
- url="http://localhost:8000",
78
-)
73
+| Field | Default | Description |
74
+| ----- | ------- | ----------- |
75
+| `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** |
76
+| `packed` | `False` | Chunked, bounded-memory transfer (see below) |
77
+| `packed_buffer_size_bytes` | 1 GiB | Chunk size when `packed=True` |
78
80
-# start
81
-base_url = "http://localhost:8000"
82
-url = f"{base_url}/start_weight_update"
83
-response = requests.post(url, json={}, timeout=60)
84
-response.raise_for_status()
85
-# send weights
86
-IPCWeightTransferEngine.trainer_send_weights(
87
- iterator=model.named_parameters(),
88
- trainer_args=trainer_args,
89
-)
90
-# finish
91
-url = f"{base_url}/finish_weight_update"
92
-response = requests.post(url, json={}, timeout=60)
93
-response.raise_for_status()
94
-```
79
+`packed` is a must-agree wire param: `trainer_init` ships it to the worker, which
80
+records it and decodes accordingly. It is not a `WeightTransferConfig` field and
81
+not a per-round `update_weights` field. `packed_buffer_size_bytes` is
82
+producer-only — the consumer rebuilds from the IPC handle plus the per-chunk
83
+`tensor_sizes`, so it never needs the buffer size.
84
96
-In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. Because the worker deserializes the payload via `pickle.loads`, the vLLM server must be started with `VLLM_ALLOW_INSECURE_SERIALIZATION=1`.
85
+## Packed (Chunked) Transfer
86
98
-```python
99
-def my_custom_sender(update_info: IPCWeightTransferUpdateInfo):
100
- # Custom logic to deliver update_info to vLLM
101
- ...
87
+By default, all weights are sent in a single `update_weights` call. For large
88
+models, this requires the full model to reside in GPU memory on both sides
89
+simultaneously. Setting `packed=True` enables **chunked transfer** with bounded
90
+GPU memory:
91
103
-trainer_args = IPCTrainerSendWeightsArgs(
104
- send_mode=my_custom_sender,
105
-)
92
+- Weights are concatenated into fixed-size packed buffers (`packed_buffer_size_bytes`).
93
+- Each chunk is sent as a separate `update_weights` call within a single `start_weight_update` / `finish_weight_update` bracket, so the layerwise reload pass is initialized once at the start and finalized once at the end regardless of chunk count.
94
+- After each chunk is consumed, the GPU memory for that chunk can be reclaimed.
95
107
-IPCWeightTransferEngine.trainer_send_weights(
108
- iterator=model.named_parameters(),
109
- trainer_args=trainer_args,
96
+```python
97
+engine = WeightTransferTrainerFactory.trainer_init(
98
+ init_info=IPCTrainerInitInfo(
99
+ rank=0,
100
+ packed=True,
101
+ packed_buffer_size_bytes=256 * 1024 * 1024, # 256 MB chunks
102
+ ),
103
+ client=client,
104
+ source=ModuleSource(model),
105
)
106
```
107
113
-See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields.
108
+On a multi-rank trainer the producer reuses one buffer across chunks, so packed
109
+mode carries a per-chunk barrier across ranks: without it, a rank could overwrite
110
+its buffer while its colocated worker is still reading the current chunk. This is
111
+handled inside `send_weights()`.
112
113
## Examples
114
117
-- [RLHF with IPC weight syncing (offline, Ray)](../../../examples/rl/rlhf_ipc.py) - Colocated training and inference on a single GPU using Ray placement groups and CUDA IPC handles
118
-- [RLHF with IPC weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - Weight transfer with a vLLM HTTP server where both server and trainer share the same GPU
115
+- [RLHF with IPC weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - **Start here.** Server and training model share a single GPU; HTTP control plane, CUDA IPC data plane. Launches and tears down its own server
116
+- [RLHF with IPC + FSDP2 and expert parallelism](../../../examples/rl/rlhf_ipc_fsdp_ep.py) - Multi-rank trainer colocated with a `--data-parallel-size 4` server on the same 4 GPUs: every FSDP rank builds an engine and joins the handle all-gather, with packed chunking and sleep/wake around the transfer
docs/training/weight_transfer/nccl.md
+143
-85
@@ -10,128 +10,186 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b
10
11
## How It Works
12
13
-1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
13
+1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). The trainer is rank 0; workers start at `rank_offset` (1).
14
2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
15
3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
16
17
-## Initialization
17
+The workers' `update_weights` and the trainer's broadcast run at the same time —
18
+both sides rendezvous inside the same NCCL calls. The trainer engine
19
+owns that concurrency internally.
20
19
-NCCL requires explicit process group setup. The trainer and inference workers must agree on a master address, port, and world size.
21
+## Inference Side
22
21
-### Inference Side
23
+The inference side takes a plain backend selector. The rendezvous parameters and
24
+the packing wire params arrive from the trainer at the init handshake.
25
26
```python
24
-from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
25
-
26
-# rank_offset accounts for the trainer occupying rank 0
27
-llm.init_weight_transfer_engine(
28
- WeightTransferInitRequest(
29
- init_info=dict(
30
- master_address=master_address,
31
- master_port=master_port,
32
- rank_offset=1,
33
- world_size=world_size, # trainer + all inference workers
34
- )
35
- )
36
-)
27
+from vllm import LLM
28
+from vllm.config import WeightTransferConfig
29
+
30
+llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="nccl"))
31
+```
32
+
33
+```bash
34
+vllm serve my-model --weight-transfer-config '{"backend": "nccl"}'
35
```
36
39
-### Trainer Side
37
+Nothing else is required: `init_weight_transfer_engine`, `start_weight_update`,
38
+`update_weights`, and `finish_weight_update` are all driven remotely by the
39
+trainer engine.
40
+
41
+## Trainer Side
42
43
```python
42
-from vllm.distributed.weight_transfer.nccl_engine import (
43
- NCCLWeightTransferEngine,
44
+from vllm.distributed.weight_transfer import (
45
+ ModuleSource,
46
+ HTTPVLLMWeightSyncClient,
47
+ WeightTransferTrainerFactory,
48
)
49
+from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
50
46
-group = NCCLWeightTransferEngine.trainer_init(
47
- dict(
51
+engine = WeightTransferTrainerFactory.trainer_init(
52
+ init_info=NCCLTrainerInitInfo(
53
master_address=master_address,
54
master_port=master_port,
50
- world_size=world_size,
51
- )
55
+ world_size=world_size, # trainer + all inference workers
56
+ rank=0, # this trainer rank; rank 0 is the sender
57
+ packed=True,
58
+ ),
59
+ client=HTTPVLLMWeightSyncClient("http://localhost:8000"), # or RayVLLMWeightSyncClient(llm)
60
+ source=ModuleSource(model),
61
)
53
-```
62
55
-!!! note
56
- `trainer_init` always assigns the trainer to rank 0. Inference workers start at `rank_offset` (typically 1).
63
+engine.send_weights() # once per sync
64
+```
65
58
-## Sending Weights
66
+`trainer_init` drives the whole handshake: it kicks off the inference side's
67
+`init_weight_transfer_engine` (through the client) on a side thread while opening
68
+the trainer's own rank-0 endpoint, since both ends must rendezvous together. It
69
+builds the worker's init info itself, with `rank_offset=1` and the same packed
70
+params, so the two sides cannot disagree.
71
60
-```python
61
-from vllm.distributed.weight_transfer.nccl_engine import (
62
- NCCLTrainerSendWeightsArgs,
63
- NCCLWeightTransferEngine,
64
-)
72
+`send_weights()` then drives `start_weight_update`, `update_weights` — run
73
+concurrently with the broadcast — and `finish_weight_update`, and returns only
74
+once every transfer has drained.
75
66
-trainer_args = NCCLTrainerSendWeightsArgs(
67
- group=group,
68
- packed=True, # use packed broadcasting for efficiency
69
-)
76
+### `NCCLTrainerInitInfo`
77
71
-NCCLWeightTransferEngine.trainer_send_weights(
72
- iterator=model.named_parameters(),
73
- trainer_args=trainer_args,
74
-)
75
-```
78
+| Field | Default | Description |
79
+| ----- | ------- | ----------- |
80
+| `master_address` | — | Rendezvous host |
81
+| `master_port` | — | Rendezvous port |
82
+| `world_size` | — | Full trainer + worker NCCL group size |
83
+| `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** |
84
+| `packed` | `True` | Use packed broadcasting |
85
+| `packed_buffer_size_bytes` | 1 GiB | Packed buffer size |
86
+| `packed_num_buffers` | 2 | Number of rotating buffers (double/triple buffering) |
87
77
-See [`NCCLTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/nccl_engine.py) for the full list of configurable fields.
88
+`packed` defaults to `True` here. The worker-side default is `False`, but it only
89
+applies when no trainer ships a value — which, on this path, never happens.
90
91
### Packed Tensor Broadcasting
92
81
-When `packed=True`, multiple weight tensors are packed into large contiguous buffers before broadcasting. This reduces the number of NCCL operations and uses double/triple buffering with dedicated CUDA streams for overlap between packing, broadcasting, and unpacking.
93
+When `packed=True`, weight tensors are packed into large contiguous buffers
94
+before broadcasting. This cuts the number of NCCL operations and uses
95
+double/triple buffering with dedicated CUDA streams to overlap packing,
96
+broadcasting, and unpacking.
97
+
98
+You set `packed`, `packed_buffer_size_bytes`, and `packed_num_buffers` **only on
99
+`NCCLTrainerInitInfo`**. The trainer propagates them to the worker inside
100
+`trainer_init`, the worker records them at the handshake, and `receive_weights`
101
+decodes with exactly the values the trainer encoded with. They are not per-round
102
+`update_weights` fields.
103
+
104
+!!! note "Memory"
105
+ The rotating buffers are live for the whole transfer:
106
+ `packed_buffer_size_bytes * packed_num_buffers` on each side (2 GiB at the
107
+ defaults). Lower `packed_buffer_size_bytes` if that is too much headroom.
108
+
109
+### The two `WeightSource` channels must agree
110
+
111
+Dense NCCL is the backend that reads *both*
112
+[`WeightSource`](base.md#weightsource) channels, so it is the one where a
113
+disagreement between them is fatal. The engine builds the per-round update info
114
+from `metadata()` and ships it ahead of the bytes; the worker sizes its receive
115
+buffers from that info, and in packed mode cuts its chunk boundaries from it. The
116
+bytes themselves come from iterating the source.
117
+
118
+If iteration disagrees with what `metadata()` declared — a reordered, omitted, or
119
+re-dtyped parameter — the two sides split the same byte stream differently. The
120
+transfer then either hangs in NCCL waiting for a length that never arrives, or
121
+loads garbage into the model.
122
+
123
+The sender therefore checks each pair against the declared metadata as it goes,
124
+one comparison per parameter, and raises naming the first divergent parameter
125
+rather than letting it reach the wire. `ModuleSource` satisfies this by
126
+construction. If you write a custom source — a Megatron export, an MoE re-fusing
127
+pass — this is the invariant to test first.
128
83
-Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightTransferUpdateInfo`) must use matching `packed_buffer_size_bytes` and `packed_num_buffers` values.
129
+!!! note
130
+ IPC does not read `metadata()` at all: it derives the update info from
131
+ iteration as it goes, so it cannot observe a divergence.
132
85
-## Receiving Weights (Inference Side)
133
+## Sparse NCCL
134
+
135
+Sparse, flat-index weight patches use a separate backend,
136
+`WeightTransferConfig(backend="sparse_nccl")`. It shares only NCCL process-group
137
+initialization with the dense engine; patches are applied directly in place to
138
+existing parameters, with no layerwise reload. The current sparse MVP requires
139
+`TP=1` and `PP=1`, and assumes a single-rank trainer.
140
87
-The inference side triggers weight reception using the four-phase protocol:
88
-`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
89
-`finish_weight_update`. The init phase is shown [above](#initialization). The
90
-remaining three steps are:
141
+Sparse is a **delta** backend: each round's patches are a fresh set of deltas
142
+from the latest optimizer step, not a stable stream of the model's parameters. So
143
+the engine takes **no `WeightSource`** — passing one raises — and each round's
144
+patches go straight to `send_weights(patches)`. A round with no patches is a
145
+no-op.
146
147
```python
93
-from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
94
-
95
-# 1. Start the weight update
96
-llm.start_weight_update()
97
-
98
-# 2. Receive weights (can be called multiple times for chunked transfers)
99
-llm.update_weights(
100
- WeightTransferUpdateRequest(
101
- update_info=dict(
102
- names=names,
103
- dtype_names=dtype_names,
104
- shapes=shapes,
105
- packed=True,
106
- )
107
- )
148
+from vllm.distributed.weight_transfer import (
149
+ RayVLLMWeightSyncClient,
150
+ WeightTransferTrainerFactory,
151
+)
152
+from vllm.distributed.weight_transfer.sparse_nccl_engine import (
153
+ SparseNCCLTrainerInitInfo,
154
+ SparseWeightPatch,
155
)
156
110
-# 3. Finish the weight update
111
-llm.finish_weight_update()
112
-```
157
+engine = WeightTransferTrainerFactory.trainer_init(
158
+ init_info=SparseNCCLTrainerInitInfo(
159
+ master_address=master_address,
160
+ master_port=master_port,
161
+ world_size=world_size,
162
+ rank=0,
163
+ ),
164
+ client=RayVLLMWeightSyncClient(llm),
165
+)
166
114
-The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
115
-must match the order in which the trainer iterates over its parameters.
167
+patches = [
168
+ SparseWeightPatch(
169
+ name="model.layers.0.mlp.down_proj.weight",
170
+ indices=flat_indices, # int32, 1-D
171
+ values=new_values, # same length as indices
172
+ full_shape=tuple(param.shape), # required when sending via the engine
173
+ )
174
+]
175
+engine.send_weights(patches)
176
+```
177
117
-`start_weight_update` must be called before `update_weights`, and
118
-`finish_weight_update` must be called after all weight chunks have been
119
-transferred. The NCCL engine receives checkpoint-format weights and applies
120
-layerwise reload processing automatically inside `start_weight_update` /
121
-`finish_weight_update`.
178
+`SparseNCCLTrainerInitInfo` takes the same rendezvous fields as the dense
179
+backend and no packed params — sparse transfers are never packed.
180
123
-## Sparse NCCL
181
+Patches are validated on the trainer *before* `start_weight_update`: int32
182
+indices, 1-D, and matching `indices`/`values` lengths, with `full_shape` set. The
183
+worker checks the same invariants when applying, but by then the broadcasts are
184
+under way, where a size mismatch wedges both sides instead of raising.
185
125
-Sparse, flat-index weight patches use a separate backend,
126
-`WeightTransferConfig(backend="sparse_nccl")`, implemented by
127
-`SparseNCCLWeightTransferEngine`. It shares only NCCL process-group
128
-initialization with the dense engine; patches are applied directly in place to
129
-existing parameters (no layerwise reload). The current sparse MVP requires
130
-`TP=1` and `PP=1`. See the example below.
186
+A trainer can hold a dense engine and a sparse engine at once — a full resync
187
+through the dense engine, cheap sparse deltas in between — as
188
+[`rlhf_sparse_nccl.py`](../../../examples/rl/rlhf_sparse_nccl.py) demonstrates.
189
190
## Examples
191
134
-- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
192
+- [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server
193
+- [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire
194
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
136
-- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
137
-- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
195
+- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient`
examples/rl/rlhf_http_ipc.py
+142
-118
@@ -1,45 +1,35 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""
4
-Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
5
-via HTTP API, with IPC-based weight syncing APIs.
6
-
7
-Unlike rlhf_nccl.py which uses NCCL and can use separate GPUs, this script
8
-uses CUDA IPC which requires the training model and vLLM server to be on the
9
-same GPU. Memory must be carefully managed to fit both models.
10
-
11
-Unlike rlhf.py which creates a vLLM instance programmatically, this script
12
-assumes you have already started a vLLM server using `vllm serve`. It uses:
13
-- OpenAI-compatible API for inference requests
14
-- HTTP endpoints for weight transfer control plane
15
-- CUDA IPC for actual weight data transfer
16
-
17
-Prerequisites:
18
- Start a vLLM server with weight transfer enabled and reduced GPU memory
19
- utilization to leave room for the training model:
20
-
21
- $ VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
22
- vllm serve facebook/opt-125m --enforce-eager \
23
- --weight-transfer-config '{"backend": "ipc"}' \
24
- --load-format dummy \
25
- --gpu-memory-utilization 0.5
26
-
27
- Then run this script:
28
-
29
- $ python rlhf_http_ipc.py
30
-
31
-The example performs the following steps:
32
-
33
-* Load the training model on GPU 0 (same GPU as the vLLM server).
34
-* Generate text using the vLLM server via OpenAI-compatible API. The output
35
- is expected to be nonsense because the server is initialized with dummy weights.
36
-* Initialize weight transfer via HTTP endpoint (no-op for IPC).
37
-* Broadcast the real weights from the training model to the vLLM server
38
- using CUDA IPC handles.
39
-* Generate text again to show normal output after the weight update.
4
+RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the
5
+data plane.
6
+
7
+ * OpenAI-compatible API for inference requests
8
+ * HTTP endpoints for the weight-transfer control plane
9
+ * CUDA IPC handles for the weight data plane
10
+
11
+1-GPU layout (single node): IPC shares GPU memory directly, so the server (TP=1)
12
+and the training model both live on GPU 0. The server is started with
13
+`--gpu-memory-utilization 0.5` to leave room for the training model.
14
+
15
+The script starts the server itself, then:
16
+
17
+ 1. Generate over HTTP → gibberish (server started with dummy weights).
18
+ 2. Pause generation, sync real weights trainer → server over IPC, resume.
19
+ 3. Generate again → sensible output.
20
+
21
+IPC handles are pickled for HTTP transport, so both sides need
22
+`VLLM_ALLOW_INSECURE_SERIALIZATION=1`; this script sets it for itself and for
23
+the server it spawns.
24
+
25
+Run:
26
+ $ python examples/rl/rlhf_http_ipc.py
27
"""
28
29
import os
30
+import subprocess
31
+import sys
32
+import time
33
34
import requests
35
import torch
@@ -53,12 +43,76 @@ from vllm.distributed.weight_transfer import (
43
)
44
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
45
56
-BASE_URL = "http://localhost:8000"
46
MODEL_NAME = "facebook/opt-125m"
47
59
-# Enable insecure serialization for IPC handle serialization
48
+SERVER_PORT = 8000
49
+BASE_URL = f"http://localhost:{SERVER_PORT}"
50
+
51
+# IPC requires colocation: the server and the training model share this GPU.
52
+SERVER_DEVICE_IDS = "0"
53
+TRAINER_DEVICE = "cuda:0"
54
+# Leave room on the shared GPU for the training model.
55
+SERVER_GPU_MEMORY_UTILIZATION = 0.5
56
+
57
+# Needed to (de)serialize IPC handles across the HTTP boundary.
58
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
59
60
+PROMPTS = [
61
+ "Hello, my name is",
62
+ "The president of the United States is",
63
+ "The capital of France is",
64
+ "The future of AI is",
65
+]
66
+
67
+
68
+def start_vllm_server() -> subprocess.Popen:
69
+ """Spawn `vllm serve` and block until it is healthy."""
70
+ serve_args = [
71
+ "vllm",
72
+ "serve",
73
+ MODEL_NAME,
74
+ "--tensor-parallel-size",
75
+ "1",
76
+ "--device-ids",
77
+ SERVER_DEVICE_IDS,
78
+ "--enforce-eager",
79
+ "--load-format",
80
+ "dummy",
81
+ "--gpu-memory-utilization",
82
+ str(SERVER_GPU_MEMORY_UTILIZATION),
83
+ "--port",
84
+ str(SERVER_PORT),
85
+ "--weight-transfer-config",
86
+ '{"backend": "ipc"}',
87
+ ]
88
+ env = os.environ.copy()
89
+ # Exposes the weight-transfer and pause/resume endpoints.
90
+ env["VLLM_SERVER_DEV_MODE"] = "1"
91
+ env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
92
+ print(f"[server] Launching: {' '.join(serve_args)}")
93
+ proc = subprocess.Popen(
94
+ serve_args,
95
+ env=env,
96
+ stdout=sys.stdout,
97
+ stderr=sys.stderr,
98
+ start_new_session=True,
99
+ )
100
+
101
+ deadline = time.monotonic() + 900
102
+ while True:
103
+ if proc.poll() is not None:
104
+ raise RuntimeError("vLLM server exited before becoming ready.")
105
+ try:
106
+ if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200:
107
+ break
108
+ except requests.RequestException:
109
+ pass
110
+ if time.monotonic() > deadline:
111
+ raise RuntimeError("vLLM server failed to start in time.")
112
+ time.sleep(2)
113
+ print("[server] Ready.")
114
+ return proc
115
+
116
117
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
118
"""Generate completions using the OpenAI-compatible API."""
@@ -76,100 +130,70 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list
130
131
def pause_generation(base_url: str) -> None:
132
"""Pause generation via HTTP endpoint."""
79
- url = f"{base_url}/pause"
80
- response = requests.post(url, timeout=60)
81
- response.raise_for_status()
133
+ requests.post(f"{base_url}/pause", timeout=60).raise_for_status()
134
135
136
def resume_generation(base_url: str) -> None:
137
"""Resume generation via HTTP endpoint."""
86
- url = f"{base_url}/resume"
87
- response = requests.post(url, timeout=60)
88
- response.raise_for_status()
89
-
90
-
91
-def get_world_size(base_url: str) -> int:
92
- """Get world size from the vLLM server."""
93
- url = f"{base_url}/get_world_size"
94
- response = requests.get(url, timeout=10)
95
- response.raise_for_status()
96
- return response.json()["world_size"]
97
-
98
-
99
-def main():
100
- # IPC requires the training model to be on the same GPU as the vLLM server
101
- # The server should be started on GPU 0 with reduced memory utilization
102
- device = "cuda:0"
103
- torch.accelerator.set_device_index(device)
104
-
105
- # Load the training model on the same GPU as the server
106
- # Use bfloat16 to reduce memory footprint
107
- print(f"Loading training model: {MODEL_NAME} on {device}")
108
- print(
109
- "Note: Ensure the vLLM server was started with --gpu-memory-utilization 0.5 "
110
- "or lower to leave room for the training model."
111
- )
112
- train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
113
- train_model.to(device)
114
- train_model.eval() # Set to eval mode to save memory
115
-
116
- # Create OpenAI client pointing to the vLLM server
117
- client = OpenAI(
118
- base_url=f"{BASE_URL}/v1",
119
- api_key="EMPTY", # vLLM doesn't require an API key by default
120
- )
138
+ requests.post(f"{base_url}/resume", timeout=60).raise_for_status()
139
122
- # Test prompts
123
- prompts = [
124
- "Hello, my name is",
125
- "The president of the United States is",
126
- "The capital of France is",
127
- "The future of AI is",
128
- ]
140
130
- # Generate text before weight update. The output is expected to be nonsense
131
- # because the server is initialized with dummy weights.
141
+def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None:
142
print("-" * 50)
133
- print("Generating text BEFORE weight update (expect nonsense):")
143
+ print(label)
144
print("-" * 50)
135
- outputs = generate_completions(client, MODEL_NAME, prompts)
145
for prompt, generated_text in zip(prompts, outputs):
146
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
147
print("-" * 50)
148
140
- print("Initializing weight transfer (IPC backend)...")
141
-
142
- # The trainer engine drives the inference side over HTTP. init for IPC is a
143
- # no-op rendezvous; the same client carries start/update/finish.
144
- engine = WeightTransferTrainerFactory.trainer_init(
145
- init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
146
- client=HTTPVLLMWeightSyncClient(BASE_URL),
147
- source=ModuleSource(train_model),
148
- )
149
-
150
- # Pause generation before weight sync
151
- pause_generation(BASE_URL)
152
-
153
- print("Broadcasting weights via CUDA IPC (HTTP)...")
154
- # One call drives start_weight_update / update_weights / finish_weight_update.
155
- engine.send_weights()
149
157
- # Resume generation after weight sync
158
- resume_generation(BASE_URL)
159
-
160
- # Generate text after weight update. The output is expected to be normal
161
- # because the real weights are now loaded.
162
- print("-" * 50)
163
- print("Generating text AFTER weight update:")
164
- print("-" * 50)
165
- outputs_updated = generate_completions(client, MODEL_NAME, prompts)
166
- for prompt, generated_text in zip(prompts, outputs_updated):
167
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
168
- print("-" * 50)
150
+def main():
151
+ server_proc = start_vllm_server()
152
+ try:
153
+ # The training model must sit on the same physical GPU as the server.
154
+ torch.accelerator.set_device_index(TRAINER_DEVICE)
155
+
156
+ print(f"[trainer] Loading training model: {MODEL_NAME} on {TRAINER_DEVICE}")
157
+ train_model = AutoModelForCausalLM.from_pretrained(
158
+ MODEL_NAME, dtype=torch.bfloat16
159
+ )
160
+ train_model.to(TRAINER_DEVICE)
161
+ train_model.eval() # eval mode to save memory on the shared GPU
162
+
163
+ client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY")
164
+
165
+ # Generate with dummy weights — expect nonsense.
166
+ outputs = generate_completions(client, MODEL_NAME, PROMPTS)
167
+ print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs)
168
+
169
+ # IPC needs no data-plane rendezvous; `trainer_init` only ships the
170
+ # `packed` flag, which the server must decode with.
171
+ print("[transfer] Initializing IPC weight transfer...")
172
+ engine = WeightTransferTrainerFactory.trainer_init(
173
+ init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
174
+ client=HTTPVLLMWeightSyncClient(BASE_URL),
175
+ source=ModuleSource(train_model),
176
+ )
177
170
- # Note: The training model and IPC handles remain in memory.
171
- # In a real RLHF training loop, you would update the training model
172
- # and create new IPC handles for each weight update.
178
+ pause_generation(BASE_URL)
179
+
180
+ # Drives start_weight_update / update_weights / finish_weight_update.
181
+ print("[sync] Sharing weights via CUDA IPC...")
182
+ engine.send_weights()
183
+ print("[sync] Weight transfer complete.")
184
+
185
+ resume_generation(BASE_URL)
186
+
187
+ # Generate with the synced weights — expect sensible output.
188
+ outputs_updated = generate_completions(client, MODEL_NAME, PROMPTS)
189
+ print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated)
190
+ finally:
191
+ print("[server] Shutting down...")
192
+ server_proc.terminate()
193
+ try:
194
+ server_proc.wait(timeout=30)
195
+ except subprocess.TimeoutExpired:
196
+ server_proc.kill()
197
198
199
if __name__ == "__main__":
examples/rl/rlhf_http_nccl.py
+148
-104
@@ -1,38 +1,33 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""
4
-Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
5
-via HTTP API, with native weight syncing APIs.
4
+RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the
5
+data plane.
6
7
-Unlike rlhf.py which creates a vLLM instance programmatically, this script
8
-assumes you have already started a vLLM server using `vllm serve`. It uses:
9
-- OpenAI-compatible API for inference requests
10
-- HTTP endpoints for weight transfer control plane
11
-- NCCL for actual weight data transfer
7
+ * OpenAI-compatible API for inference requests
8
+ * HTTP endpoints for the weight-transfer control plane
9
+ * NCCL for the weight data plane
10
13
-Prerequisites:
14
- Start a vLLM server with weight transfer enabled:
11
+3-GPU layout (single node):
12
+ Inference — GPUs 0-1, `vllm serve` with TP=2 and fp8 quantization
13
+ Training — GPU 2, a bf16 Hugging Face model in this process
14
+ (the server quantizes to fp8 as it loads)
15
16
- $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
17
- --enforce-eager \
18
- --weight-transfer-config '{"backend": "nccl"}' \
19
- --load-format dummy
16
+The script starts the server itself, then:
17
21
- Then run this script:
18
+ 1. Generate over HTTP → gibberish (server started with dummy weights).
19
+ 2. Pause generation, sync real weights trainer → server over NCCL, resume.
20
+ 3. Generate again → sensible output.
21
23
- $ python rlhf_http.py
24
-
25
-The example performs the following steps:
26
-
27
-* Load the training model on GPU 0.
28
-* Generate text using the vLLM server via OpenAI-compatible API. The output
29
- is expected to be nonsense because the server is initialized with dummy weights.
30
-* Initialize weight transfer via HTTP endpoint.
31
-* Broadcast the real weights from the training model to the vLLM server
32
- using NCCL.
33
-* Generate text again to show normal output after the weight update.
22
+Run:
23
+ $ python examples/rl/rlhf_http_nccl.py
24
"""
25
26
+import os
27
+import subprocess
28
+import sys
29
+import time
30
+
31
import requests
32
import torch
33
from openai import OpenAI
@@ -46,9 +41,71 @@ from vllm.distributed.weight_transfer import (
41
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
42
from vllm.utils.network_utils import get_ip, get_open_port
43
49
-BASE_URL = "http://localhost:8000"
44
MODEL_NAME = "facebook/opt-125m"
45
46
+SERVER_PORT = 8000
47
+BASE_URL = f"http://localhost:{SERVER_PORT}"
48
+
49
+INFERENCE_TP_SIZE = 2
50
+# Physical GPUs for the server; the trainer takes the next one.
51
+SERVER_DEVICE_IDS = "0,1"
52
+TRAINER_DEVICE = "cuda:2"
53
+
54
+PROMPTS = [
55
+ "Hello, my name is",
56
+ "The president of the United States is",
57
+ "The capital of France is",
58
+ "The future of AI is",
59
+]
60
+
61
+
62
+def start_vllm_server() -> subprocess.Popen:
63
+ """Spawn `vllm serve` and block until it is healthy."""
64
+ serve_args = [
65
+ "vllm",
66
+ "serve",
67
+ MODEL_NAME,
68
+ "--tensor-parallel-size",
69
+ str(INFERENCE_TP_SIZE),
70
+ "--device-ids",
71
+ SERVER_DEVICE_IDS,
72
+ "--quantization",
73
+ "fp8",
74
+ "--enforce-eager",
75
+ "--load-format",
76
+ "dummy",
77
+ "--port",
78
+ str(SERVER_PORT),
79
+ "--weight-transfer-config",
80
+ '{"backend": "nccl"}',
81
+ ]
82
+ env = os.environ.copy()
83
+ # Exposes the weight-transfer and pause/resume endpoints.
84
+ env["VLLM_SERVER_DEV_MODE"] = "1"
85
+ print(f"[server] Launching: {' '.join(serve_args)}")
86
+ proc = subprocess.Popen(
87
+ serve_args,
88
+ env=env,
89
+ stdout=sys.stdout,
90
+ stderr=sys.stderr,
91
+ start_new_session=True,
92
+ )
93
+
94
+ deadline = time.monotonic() + 900
95
+ while True:
96
+ if proc.poll() is not None:
97
+ raise RuntimeError("vLLM server exited before becoming ready.")
98
+ try:
99
+ if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200:
100
+ break
101
+ except requests.RequestException:
102
+ pass
103
+ if time.monotonic() > deadline:
104
+ raise RuntimeError("vLLM server failed to start in time.")
105
+ time.sleep(2)
106
+ print("[server] Ready.")
107
+ return proc
108
+
109
110
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
111
"""Generate completions using the OpenAI-compatible API."""
@@ -66,105 +123,92 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list
123
124
def pause_generation(base_url: str) -> None:
125
"""Pause generation via HTTP endpoint."""
69
- url = f"{base_url}/pause"
70
- response = requests.post(url, timeout=60)
71
- response.raise_for_status()
126
+ requests.post(f"{base_url}/pause", timeout=60).raise_for_status()
127
128
129
def resume_generation(base_url: str) -> None:
130
"""Resume generation via HTTP endpoint."""
76
- url = f"{base_url}/resume"
77
- response = requests.post(url, timeout=60)
78
- response.raise_for_status()
131
+ requests.post(f"{base_url}/resume", timeout=60).raise_for_status()
132
133
134
def get_world_size(base_url: str) -> int:
82
- """Get world size from the vLLM server."""
83
- url = f"{base_url}/get_world_size"
84
- response = requests.get(url, timeout=10)
135
+ """Get the number of inference workers from the vLLM server."""
136
+ response = requests.get(f"{base_url}/get_world_size", timeout=10)
137
response.raise_for_status()
138
return response.json()["world_size"]
139
140
89
-def main():
90
- # Get the inference world size from the vLLM server
91
- inference_world_size = get_world_size(BASE_URL)
92
- world_size = inference_world_size + 1 # +1 for the trainer
93
- device = f"cuda:{inference_world_size}"
94
- torch.accelerator.set_device_index(device)
95
-
96
- # Load the training model
97
- print(f"Loading training model: {MODEL_NAME}")
98
- train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
99
- train_model.to(device)
100
-
101
- # Create OpenAI client pointing to the vLLM server
102
- client = OpenAI(
103
- base_url=f"{BASE_URL}/v1",
104
- api_key="EMPTY", # vLLM doesn't require an API key by default
105
- )
106
-
107
- # Test prompts
108
- prompts = [
109
- "Hello, my name is",
110
- "The president of the United States is",
111
- "The capital of France is",
112
- "The future of AI is",
113
- ]
114
-
115
- # Generate text before weight update. The output is expected to be nonsense
116
- # because the server is initialized with dummy weights.
141
+def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None:
142
print("-" * 50)
118
- print("Generating text BEFORE weight update (expect nonsense):")
143
+ print(label)
144
print("-" * 50)
120
- outputs = generate_completions(client, MODEL_NAME, prompts)
145
for prompt, generated_text in zip(prompts, outputs):
146
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
147
print("-" * 50)
148
125
- # Set up the communication channel between the training process and the
126
- # vLLM server. The trainer is rank 0, vLLM worker(s) start after it.
127
- master_address = get_ip()
128
- master_port = get_open_port()
129
-
130
- print(f"Initializing weight transfer: master={master_address}:{master_port}")
131
-
132
- # The trainer engine owns the full handshake. `trainer_init` kicks off the
133
- # server's init_weight_transfer_engine (via the HTTP client) on a worker
134
- # thread while opening the trainer's NCCL endpoint, so both ends rendezvous
135
- # together — no manual threading needed here.
136
- engine = WeightTransferTrainerFactory.trainer_init(
137
- init_info=NCCLTrainerInitInfo(
138
- master_address=master_address,
139
- master_port=master_port,
140
- world_size=world_size,
141
- rank=0, # single-GPU trainer is the sole (sender) rank
142
- packed=True,
143
- ),
144
- client=HTTPVLLMWeightSyncClient(BASE_URL),
145
- source=ModuleSource(train_model),
146
- )
149
148
- # Pause generation before weight sync
149
- pause_generation(BASE_URL)
150
+def main():
151
+ server_proc = start_vllm_server()
152
+ try:
153
+ # The trainer sits on the GPU after the server's, and is NCCL rank 0.
154
+ torch.accelerator.set_device_index(TRAINER_DEVICE)
155
+
156
+ print(f"[trainer] Loading training model: {MODEL_NAME} on {TRAINER_DEVICE}")
157
+ train_model = AutoModelForCausalLM.from_pretrained(
158
+ MODEL_NAME, dtype=torch.bfloat16
159
+ )
160
+ train_model.to(TRAINER_DEVICE)
161
151
- # One call drives start_weight_update / update_weights / finish_weight_update
152
- # over HTTP, concurrent with the NCCL broadcast.
153
- print("Broadcasting weights via NCCL...")
154
- engine.send_weights()
162
+ client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY")
163
156
- # Resume generation after weight sync
157
- resume_generation(BASE_URL)
164
+ # Generate with dummy weights — expect nonsense.
165
+ outputs = generate_completions(client, MODEL_NAME, PROMPTS)
166
+ print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs)
167
159
- # Generate text after weight update. The output is expected to be normal
160
- # because the real weights are now loaded.
161
- print("-" * 50)
162
- print("Generating text AFTER weight update:")
163
- print("-" * 50)
164
- outputs_updated = generate_completions(client, MODEL_NAME, prompts)
165
- for prompt, generated_text in zip(prompts, outputs_updated):
166
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
167
- print("-" * 50)
168
+ # The transfer NCCL group is the trainer plus every inference worker.
169
+ world_size = get_world_size(BASE_URL) + 1
170
+ master_address = get_ip()
171
+ master_port = get_open_port()
172
+ print(
173
+ f"[transfer] Rendezvous at {master_address}:{master_port}, "
174
+ f"world_size={world_size} (1 trainer + {world_size - 1} vLLM workers)"
175
+ )
176
+
177
+ # `trainer_init` drives the handshake: it initializes the server's
178
+ # transfer engine while opening the trainer's own NCCL endpoint, so both
179
+ # ends rendezvous together.
180
+ engine = WeightTransferTrainerFactory.trainer_init(
181
+ init_info=NCCLTrainerInitInfo(
182
+ master_address=master_address,
183
+ master_port=master_port,
184
+ world_size=world_size,
185
+ rank=0, # single-GPU trainer is the sole (sender) rank
186
+ packed=True,
187
+ ),
188
+ client=HTTPVLLMWeightSyncClient(BASE_URL),
189
+ source=ModuleSource(train_model),
190
+ )
191
+
192
+ pause_generation(BASE_URL)
193
+
194
+ # Drives start_weight_update / update_weights / finish_weight_update,
195
+ # concurrent with the NCCL broadcast.
196
+ print("[sync] Broadcasting weights via NCCL...")
197
+ engine.send_weights()
198
+ print("[sync] Weight broadcast complete.")
199
+
200
+ resume_generation(BASE_URL)
201
+
202
+ # Generate with the synced weights — expect sensible output.
203
+ outputs_updated = generate_completions(client, MODEL_NAME, PROMPTS)
204
+ print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated)
205
+ finally:
206
+ print("[server] Shutting down...")
207
+ server_proc.terminate()
208
+ try:
209
+ server_proc.wait(timeout=30)
210
+ except subprocess.TimeoutExpired:
211
+ server_proc.kill()
212
213
214
if __name__ == "__main__":
examples/rl/rlhf_ipc.py
deleted
-151
@@ -1,151 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray,
5
-with IPC-based weight syncing APIs
6
-
7
-The script colocates the training and inference workloads onto the same GPU using Ray.
8
-
9
-The example performs the following steps:
10
-
11
-* Request a placement group of 1 GPU.
12
-* Place the inference model on the above GPU using the placement group.
13
-* Place and load the training model on the same GPU using the placement group.
14
-* Generate text from a list of prompts using the inference engine.
15
-* Update the weights of the training model and broadcast the updated weights
16
- to the inference engine by using CUDA IPC handles. Note that
17
- for demonstration purposes we simply zero out the weights.
18
-
19
-This example assumes a single-node cluster with a single GPU,
20
-but can be extended to multiple GPUs.
21
-"""
22
-
23
-import os
24
-
25
-import ray
26
-from ray.util.placement_group import placement_group
27
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
28
-from transformers import AutoModelForCausalLM
29
-
30
-from vllm import LLM, SamplingParams
31
-from vllm.config import WeightTransferConfig
32
-from vllm.distributed.weight_transfer import (
33
- ModuleSource,
34
- RayVLLMWeightSyncClient,
35
- WeightTransferTrainerFactory,
36
-)
37
-from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
38
-
39
-
40
-class MyLLM(LLM):
41
- """Configure the vLLM worker for Ray placement group execution."""
42
-
43
- def __init__(self, *args, **kwargs):
44
- # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
45
- # so that vLLM can manage its own device placement within the worker.
46
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
47
- # Each worker uses 0.4 GPU so that two instances fit on the same GPU.
48
- os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
49
- os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
50
- # needed for ipc handle serialization
51
- os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
52
- super().__init__(*args, **kwargs)
53
-
54
-
55
-# Load the OPT-125M model onto GPU 0 for the training workload.
56
-
57
-MODEL_NAME = "facebook/opt-125m"
58
-
59
-
60
-@ray.remote
61
-class TrainModel:
62
- def __init__(self, llm_handle: ray.actor.ActorHandle):
63
- self.train_model = AutoModelForCausalLM.from_pretrained(
64
- MODEL_NAME,
65
- )
66
- self.train_model.to("cuda:0")
67
- self.llm_handle = llm_handle
68
-
69
- def init_weight_transfer(self):
70
- """Build the trainer-side IPC engine (no rendezvous needed for IPC)."""
71
- self.engine = WeightTransferTrainerFactory.trainer_init(
72
- init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender
73
- client=RayVLLMWeightSyncClient(self.llm_handle),
74
- source=ModuleSource(self.train_model),
75
- )
76
-
77
- def broadcast_weights(self):
78
- """Broadcast weights to the inference engine using IPC.
79
-
80
- Drives start/update/finish on the inference side internally.
81
- """
82
- self.engine.send_weights()
83
-
84
-
85
-ray.init()
86
-
87
-pg_colocate = placement_group([{"GPU": 1, "CPU": 0}])
88
-ray.get(pg_colocate.ready())
89
-
90
-
91
-llm = ray.remote(
92
- num_cpus=0,
93
- num_gpus=0,
94
- scheduling_strategy=PlacementGroupSchedulingStrategy(
95
- placement_group=pg_colocate,
96
- placement_group_capture_child_tasks=True,
97
- ),
98
-)(MyLLM).remote(
99
- model=MODEL_NAME,
100
- enforce_eager=True,
101
- tensor_parallel_size=1,
102
- distributed_executor_backend="ray",
103
- gpu_memory_utilization=0.7,
104
- weight_transfer_config=WeightTransferConfig(backend="ipc"),
105
- load_format="dummy",
106
-)
107
-
108
-train_model = TrainModel.options(
109
- num_gpus=0.1,
110
- num_cpus=0,
111
- scheduling_strategy=PlacementGroupSchedulingStrategy(
112
- placement_group=pg_colocate, placement_group_capture_child_tasks=True
113
- ),
114
-).remote(llm)
115
-
116
-
117
-# Generate text from the prompts.
118
-prompts = [
119
- "Hello, my name is",
120
- "The president of the United States is",
121
- "The capital of France is",
122
- "The future of AI is",
123
-]
124
-
125
-sampling_params = SamplingParams(temperature=0)
126
-
127
-outputs = ray.get(llm.generate.remote(prompts, sampling_params))
128
-
129
-print("-" * 50)
130
-for output in outputs:
131
- prompt = output.prompt
132
- generated_text = output.outputs[0].text
133
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
134
- print("-" * 50)
135
-
136
-ray.get(llm.sleep.remote(level=0))
137
-
138
-ray.get(train_model.init_weight_transfer.remote())
139
-# One call drives start_weight_update / update_weights / finish_weight_update.
140
-ray.get(train_model.broadcast_weights.remote())
141
-
142
-ray.get(llm.wake_up.remote(tags=["scheduling"]))
143
-
144
-outputs_packed = ray.get(llm.generate.remote(prompts, sampling_params))
145
-print("-" * 50)
146
-print("Results after packed/chunked IPC weight sync:")
147
-for output in outputs_packed:
148
- prompt = output.prompt
149
- generated_text = output.outputs[0].text
150
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
151
- print("-" * 50)
examples/rl/rlhf_ipc_fsdp_ep.py
+196
-247
@@ -4,96 +4,88 @@
4
RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC**
5
weight transfer and **packed** tensors.
6
7
-Layout (4 GPUs, TP=1, DP=4, EP):
8
- * One Ray placement group per GPU.
9
- * Each PG holds one FSDP training worker and one vLLM ``LLM`` instance
10
- (sync API) using fractional GPUs so both fit on the same device.
11
- * The 4 ``LLM`` instances form a DP group via env-var-based SPMD
12
- coordination (``VLLM_DP_RANK``, ``VLLM_DP_SIZE``, etc.), the same
13
- mechanism used by ``examples/offline_inference/data_parallel.py``.
14
- * A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
15
- waits for initialization, and orchestrates generation / weight-sync.
7
+Multi-rank version of `rlhf_http_ipc.py`: the trainer is 4 FSDP2 Ray actors
8
+colocated with a data-parallel `vllm serve` on the same 4 physical GPUs.
9
+
10
+4-GPU layout (single node), all colocated:
11
+ Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors
12
+ Inference — the same 4 GPUs, `vllm serve --data-parallel-size 4 -tp 1
13
+ --enable-expert-parallel` (EP_SIZE = TP x DP = 4)
14
+
15
+IPC requires the trainer and the server to sit on the same GPUs, so the script
16
+reserves the training GPUs through Ray first, asks Ray which ones it got, and
17
+pins the server to exactly those with `--device-ids`.
18
+
19
+Both sides share each GPU, so the server is capped with
20
+`--gpu-memory-utilization` and its weights are moved aside for the transfer:
21
+
22
+ 1. `/sleep?level=1` — offload server weights to CPU, drop KV cache
23
+ 2. `/wake_up?tags=weights` — weights back on GPU, KV cache still free
24
+ 3. packed IPC transfer — overwrite weights with room to spare
25
+ 4. `/wake_up?tags=kv_cache&tags=scheduling` — re-allocate KV cache, resume
26
27
Every FSDP rank builds an ``IPCTrainerWeightTransferEngine`` (via ``trainer_init``)
28
and calls ``send_weights()``; all ranks join the IPC handle all-gather, and only
19
-rank 0 (the sender) ships the merged handles and drives the DP LLM actors through
20
-its ``RayVLLMWeightSyncClient``.
29
+rank 0 (the sender) ships the merged handles and drives the server.
30
31
This example was run on 4xH100.
32
+
33
+Run:
34
+ $ python examples/rl/rlhf_ipc_fsdp_ep.py
35
"""
36
37
from __future__ import annotations
38
39
import os
40
+import subprocess
41
+import sys
42
+import time
43
44
import ray
45
+import requests
46
import torch
47
import torch.distributed as dist
48
from huggingface_hub import snapshot_download
33
-from ray.util.placement_group import placement_group
34
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
49
+from openai import OpenAI
50
from torch.distributed.fsdp import fully_shard
51
from transformers import AutoModelForCausalLM
52
38
-from vllm import LLM, SamplingParams
39
-from vllm.config import WeightTransferConfig
53
from vllm.distributed.weight_transfer import (
54
+ HTTPVLLMWeightSyncClient,
55
ModuleSource,
42
- RayVLLMWeightSyncClient,
56
WeightTransferTrainerFactory,
57
)
58
from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
59
from vllm.utils.network_utils import get_ip, get_open_port
60
48
-# The inference side only needs the backend; the packed wire params are
49
-# trainer-side and get propagated to the workers at the init handshake.
50
-WEIGHT_TRANSFER_CONFIG = WeightTransferConfig(backend="ipc")
51
-# Packed IPC transfer with a 1 GB buffer (matches the per-chunk buffer size).
52
-PACKED = True
53
-PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024
54
-
55
-TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
56
-VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))
57
-
61
MODEL_NAME = "Qwen/Qwen3-30B-A3B"
62
+SERVED_MODEL_NAME = "policy"
63
64
FSDP_WORLD_SIZE = 4
65
INFERENCE_TP_SIZE = 1
66
INFERENCE_DP_SIZE = 4
67
68
+# Packed IPC transfer with a 1 GB chunk buffer.
69
+PACKED = True
70
+PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024
71
65
-class MyLLM(LLM):
66
- """LLM subclass that configures DP env vars for SPMD coordination."""
67
-
68
- def __init__(
69
- self,
70
- *args,
71
- dp_rank: int = 0,
72
- dp_size: int = 1,
73
- dp_master_ip: str = "127.0.0.1",
74
- dp_master_port: int = 0,
75
- **kwargs,
76
- ):
77
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
78
- os.environ["VLLM_RAY_PER_WORKER_GPUS"] = str(VLLM_GPU_FRACTION)
79
- os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
80
- os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
81
-
82
- os.environ["VLLM_DP_RANK"] = str(dp_rank)
83
- os.environ["VLLM_DP_RANK_LOCAL"] = str(dp_rank)
84
- os.environ["VLLM_DP_SIZE"] = str(dp_size)
85
- os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
86
- os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
72
+# The server shares each GPU with a training rank, so cap what it reserves.
73
+SERVER_GPU_MEMORY_UTILIZATION = 0.35
74
88
- super().__init__(*args, **kwargs)
75
+SERVER_PORT = 8000
76
+BASE_URL = f"http://localhost:{SERVER_PORT}"
77
90
- def ready(self):
91
- return True
78
+PROMPTS = [
79
+ "Hello, my name is",
80
+ "The president of the United States is",
81
+ "The capital of France is",
82
+ "The future of AI is",
83
+]
84
85
94
-@ray.remote(num_cpus=0, num_gpus=TRAIN_GPU_FRACTION)
86
+@ray.remote(num_gpus=1)
87
class FSDPTrainWorker:
96
- """One FSDP2 worker per GPU; colocated with vLLM DP rank via placement group."""
88
+ """One FSDP2 worker per GPU; colocated with one vLLM DP rank."""
89
90
def __init__(
91
self,
@@ -115,12 +107,6 @@ class FSDPTrainWorker:
107
model_name, torch_dtype=torch.bfloat16
108
)
109
118
- self.weight_names = [n for n, _ in model.named_parameters()]
119
- self.weight_dtype_names = [
120
- str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
121
- ]
122
- self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
123
-
110
for layer in model.model.layers:
111
fully_shard(layer)
112
fully_shard(model)
@@ -130,237 +116,200 @@ class FSDPTrainWorker:
116
def get_rank(self):
117
return self.rank
118
133
- def get_weight_metadata(self):
134
- return self.weight_names, self.weight_dtype_names, self.weight_shapes
135
-
136
- def setup_engine(self, llm_handles):
137
- """Build the trainer IPC engine on every FSDP rank.
119
+ def get_gpu_ids(self):
120
+ """Physical GPU id(s) Ray assigned to this worker."""
121
+ return ray.get_gpu_ids()
122
139
- Called on all ranks: rank 0 becomes the sender (drives the inference
140
- side and performs the no-op IPC init handshake against all DP LLM
141
- actors); the other ranks hold a null-client engine and only join the
142
- IPC handle all-gather during send_weights.
143
- """
123
+ def setup_engine(self, base_url: str):
124
+ """Build the trainer IPC engine. Called on every FSDP rank."""
125
self.engine = WeightTransferTrainerFactory.trainer_init(
126
init_info=IPCTrainerInitInfo(
127
rank=self.rank, # FSDP rank; sender is 0
128
packed=PACKED,
129
packed_buffer_size_bytes=PACKED_BUFFER_SIZE_BYTES,
130
),
150
- client=RayVLLMWeightSyncClient(llm_handles),
131
+ client=HTTPVLLMWeightSyncClient(base_url),
132
source=ModuleSource(self.model),
133
)
134
135
def gather_and_broadcast_weights_ipc(self):
155
- """All-gather full params across FSDP ranks; rank 0 sends to vLLM.
156
-
157
- Called on all ranks concurrently. `send_weights` gathers each param and
158
- contributes to the IPC handle all-gather on every rank; only rank 0 (the
159
- sender) drives start/update/finish on the inference side (the other
160
- ranks' client RPCs no-op).
161
- """
136
+ """Send the current weights to vLLM. Called on every FSDP rank."""
137
self.engine.send_weights()
138
139
165
-@ray.remote(num_cpus=1)
166
-class DataParallelInferenceEngine:
167
- """Manages a pool of DP-sharded vLLM LLM actors.
140
+def start_vllm_server(model_path: str, device_ids: str) -> subprocess.Popen:
141
+ """Spawn a `vllm serve` HTTP server (DP+EP) pinned to `device_ids`."""
142
+ serve_args = [
143
+ "vllm",
144
+ "serve",
145
+ model_path,
146
+ "--served-model-name",
147
+ SERVED_MODEL_NAME,
148
+ "--tensor-parallel-size",
149
+ str(INFERENCE_TP_SIZE),
150
+ "--data-parallel-size",
151
+ str(INFERENCE_DP_SIZE),
152
+ "--enable-expert-parallel",
153
+ # Pins the server to the same physical GPUs as the training ranks.
154
+ "--device-ids",
155
+ device_ids,
156
+ "--enable-sleep-mode",
157
+ "--enforce-eager",
158
+ "--load-format",
159
+ "dummy",
160
+ "--gpu-memory-utilization",
161
+ str(SERVER_GPU_MEMORY_UTILIZATION),
162
+ "--port",
163
+ str(SERVER_PORT),
164
+ "--weight-transfer-config",
165
+ '{"backend": "ipc"}',
166
+ ]
167
+ env = os.environ.copy()
168
+ env["VLLM_SERVER_DEV_MODE"] = "1" # exposes weight-transfer + sleep endpoints
169
+ env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" # IPC handles over HTTP
170
+ env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
171
+ print(f"[server] Launching: {' '.join(serve_args)} (GPUs {device_ids})")
172
+ proc = subprocess.Popen(
173
+ serve_args,
174
+ env=env,
175
+ stdout=sys.stdout,
176
+ stderr=sys.stderr,
177
+ start_new_session=True,
178
+ )
179
169
- Spawns one MyLLM actor per placement group, waits for all engines to
170
- finish initializing, and exposes generation / weight-sync helpers.
171
- """
180
+ # Wait for the server to come up (model load can take a while).
181
+ deadline = time.monotonic() + 1800
182
+ while True:
183
+ if proc.poll() is not None:
184
+ raise RuntimeError("vLLM server exited before becoming ready.")
185
+ try:
186
+ if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200:
187
+ break
188
+ except requests.RequestException:
189
+ pass
190
+ if time.monotonic() > deadline:
191
+ raise RuntimeError("vLLM server failed to start in time.")
192
+ time.sleep(2)
193
+ print("[server] Ready.")
194
+ return proc
195
+
196
+
197
+def generate_completions(client: OpenAI, prompts: list[str]) -> list[str]:
198
+ """Generate completions via the OpenAI HTTP API."""
199
+ results = []
200
+ for prompt in prompts:
201
+ response = client.completions.create(
202
+ model=SERVED_MODEL_NAME,
203
+ prompt=prompt,
204
+ max_tokens=32,
205
+ temperature=0,
206
+ )
207
+ results.append(response.choices[0].text)
208
+ return results
209
173
- def __init__(
174
- self,
175
- model: str,
176
- pgs: list,
177
- dp_master_ip: str,
178
- dp_master_port: int,
179
- ):
180
- dp_size = len(pgs)
181
- self.llm_actors = []
182
- for r in range(dp_size):
183
- sched = PlacementGroupSchedulingStrategy(
184
- placement_group=pgs[r],
185
- placement_group_capture_child_tasks=True,
186
- )
187
- actor = (
188
- ray.remote(num_cpus=0, num_gpus=0)(MyLLM)
189
- .options(scheduling_strategy=sched)
190
- .remote(
191
- model=model,
192
- enforce_eager=True,
193
- tensor_parallel_size=INFERENCE_TP_SIZE,
194
- distributed_executor_backend="ray",
195
- enable_expert_parallel=True,
196
- gpu_memory_utilization=0.35,
197
- weight_transfer_config=WEIGHT_TRANSFER_CONFIG,
198
- enable_sleep_mode=True,
199
- load_format="dummy",
200
- dp_rank=r,
201
- dp_size=dp_size,
202
- dp_master_ip=dp_master_ip,
203
- dp_master_port=dp_master_port,
204
- )
205
- )
206
- self.llm_actors.append(actor)
207
-
208
- ray.get([actor.ready.remote() for actor in self.llm_actors])
209
-
210
- def get_llm_actors(self):
211
- return self.llm_actors
212
-
213
- def generate(self, prompts: list[str], sampling_params):
214
- """Distribute prompts round-robin across DP ranks and collect results."""
215
- dp_size = len(self.llm_actors)
216
- per_rank: list[list[str]] = [[] for _ in range(dp_size)]
217
- indices: list[list[int]] = [[] for _ in range(dp_size)]
218
-
219
- for i, prompt in enumerate(prompts):
220
- rank = i % dp_size
221
- per_rank[rank].append(prompt)
222
- indices[rank].append(i)
223
-
224
- refs = [
225
- actor.generate.remote(per_rank[r], sampling_params)
226
- for r, actor in enumerate(self.llm_actors)
227
- if per_rank[r]
228
- ]
229
- all_outputs = ray.get(refs)
230
-
231
- ordered = [None] * len(prompts)
232
- rank_idx = 0
233
- for r in range(dp_size):
234
- if per_rank[r]:
235
- for local_i, orig_i in enumerate(indices[r]):
236
- ordered[orig_i] = all_outputs[rank_idx][local_i]
237
- rank_idx += 1
238
- return ordered
239
-
240
- def sleep(self, level: int = 0):
241
- ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])
242
-
243
- def wake_up(self, tags: list[str] | None = None):
244
- ray.get([actor.wake_up.remote(tags=tags) for actor in self.llm_actors])
210
+
211
+def sleep_engine(base_url: str, level: int) -> None:
212
+ """Put the engine to sleep (level 1 offloads weights, drops KV cache)."""
213
+ response = requests.post(f"{base_url}/sleep", params={"level": level}, timeout=600)
214
+ response.raise_for_status()
215
+
216
+
217
+def wake_up_engine(base_url: str, tags: list[str] | None = None) -> None:
218
+ """Wake the engine, optionally only for specific memory tags."""
219
+ params = [("tags", tag) for tag in tags] if tags else None
220
+ response = requests.post(f"{base_url}/wake_up", params=params, timeout=600)
221
+ response.raise_for_status()
222
+
223
+
224
+def print_generations(label: str, prompts: list[str], outputs: list[str]) -> None:
225
+ print("-" * 60)
226
+ print(label)
227
+ print("-" * 60)
228
+ for prompt, text in zip(prompts, outputs):
229
+ print(f"Prompt: {prompt!r}")
230
+ print(f"Generated: {text!r}")
231
+ print("-" * 60)
232
233
234
def main():
235
ray.init(
236
runtime_env={
237
"env_vars": {
238
+ # The trainer pickles IPC handles for the HTTP client.
239
"VLLM_ALLOW_INSECURE_SERIALIZATION": "1",
240
}
241
}
242
)
243
256
- assert TRAIN_GPU_FRACTION + VLLM_GPU_FRACTION <= 1.0, (
257
- "Train + vLLM GPU fractions must sum to at most 1.0 per bundle."
258
- )
259
-
244
local_model_path = snapshot_download(MODEL_NAME)
245
print(f"[init] Model downloaded to {local_model_path}")
246
247
fsdp_master_addr = get_ip()
248
fsdp_master_port = get_open_port()
265
- dp_master_port = get_open_port()
266
- dp_master_ip = get_ip()
267
-
268
- # Create one placement group per DP rank (one GPU each).
269
- pgs = []
270
- for _ in range(INFERENCE_DP_SIZE):
271
- pg = placement_group([{"GPU": 1, "CPU": 1}])
272
- pgs.append(pg)
273
- ray.get([pg.ready() for pg in pgs])
274
- print(f"[init] {len(pgs)} placement groups ready.")
275
-
276
- # Launch FSDP training workers, one per PG.
277
- scheduling = [
278
- PlacementGroupSchedulingStrategy(
279
- placement_group=pgs[r],
280
- placement_group_capture_child_tasks=True,
281
- )
282
- for r in range(FSDP_WORLD_SIZE)
283
- ]
249
250
+ # Launch the training workers first so Ray reserves their GPUs; the server
251
+ # is then pinned to those same physical GPUs.
252
fsdp_workers = [
286
- FSDPTrainWorker.options(scheduling_strategy=scheduling[r]).remote(
253
+ FSDPTrainWorker.remote(
254
local_model_path,
288
- r,
255
+ rank,
256
FSDP_WORLD_SIZE,
257
fsdp_master_addr,
258
fsdp_master_port,
259
)
293
- for r in range(FSDP_WORLD_SIZE)
260
+ for rank in range(FSDP_WORLD_SIZE)
261
]
262
ray.get([w.get_rank.remote() for w in fsdp_workers])
296
- print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")
297
-
298
- # Launch DP inference engine (spawns and initializes all LLM actors).
299
- inference_engine = DataParallelInferenceEngine.remote(
300
- model=local_model_path,
301
- pgs=pgs,
302
- dp_master_ip=dp_master_ip,
303
- dp_master_port=dp_master_port,
304
- )
305
- llm_actors = ray.get(inference_engine.get_llm_actors.remote())
306
- print(f"[init] {INFERENCE_DP_SIZE} LLM actors ready.")
307
-
308
- prompts = [
309
- "Hello, my name is",
310
- "The president of the United States is",
311
- "The capital of France is",
312
- "The future of AI is",
313
- ]
314
- sampling_params = SamplingParams(temperature=0)
263
+ print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.")
264
316
- print("[generate] Generating with dummy weights...")
317
- outputs = ray.get(inference_engine.generate.remote(prompts, sampling_params))
318
- print("-" * 60)
319
- print("BEFORE weight sync (dummy weights):")
320
- print("-" * 60)
321
- for output in outputs:
322
- print(f"Prompt: {output.prompt!r}")
323
- print(f"Generated: {output.outputs[0].text!r}")
324
- print("-" * 60)
325
-
326
- # --- Weight transfer ---
327
- # The rank-0 FSDP worker owns the trainer engine and drives the inference
328
- # side (init/start/update/finish) for all DP actors via the Ray client.
329
- print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...")
330
- ray.get([w.setup_engine.remote(llm_actors) for w in fsdp_workers])
331
-
332
- # Two-phase sleep/wake pattern:
333
- # 1. sleep(level=1) — offload weights to CPU, discard KV cache
334
- # 2. wake_up(tags=["weights"]) — bring weights back to GPU (KV cache still free)
335
- # 3. IPC weight transfer — overwrite weights, plenty of room without KV cache
336
- # 4. wake_up(tags=["kv_cache"]) — re-allocate KV cache for inference
337
- print("[sync] Sleeping engines (offload weights + free KV cache)...")
338
- ray.get(inference_engine.sleep.remote(level=1))
339
-
340
- print("[sync] Waking weights (KV cache stays free)...")
341
- ray.get(inference_engine.wake_up.remote(tags=["weights"]))
342
-
343
- # All FSDP ranks participate in the IPC handle all-gather; rank 0's engine
344
- # additionally drives start_weight_update / update_weights /
345
- # finish_weight_update on the inference side.
346
- print("[sync] Packed IPC transfer FSDP → vLLM...")
347
- ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers])
348
- print("[sync] Weight transfer complete.")
349
-
350
- print("[sync] Waking KV cache + scheduling...")
351
- ray.get(inference_engine.wake_up.remote(tags=["kv_cache", "scheduling"]))
352
-
353
- print("[generate] Generating with synced weights...")
354
- outputs_updated = ray.get(
355
- inference_engine.generate.remote(prompts, sampling_params)
265
+ training_gpus = sorted(
266
+ int(g)
267
+ for ids in ray.get([w.get_gpu_ids.remote() for w in fsdp_workers])
268
+ for g in ids
269
)
357
- print("-" * 60)
358
- print("AFTER weight sync (real weights):")
359
- print("-" * 60)
360
- for output in outputs_updated:
361
- print(f"Prompt: {output.prompt!r}")
362
- print(f"Generated: {output.outputs[0].text!r}")
363
- print("-" * 60)
270
+ if len(training_gpus) != INFERENCE_TP_SIZE * INFERENCE_DP_SIZE:
271
+ raise RuntimeError(
272
+ f"Need {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} colocated GPUs but "
273
+ f"Ray assigned training to {training_gpus}."
274
+ )
275
+ device_ids = ",".join(str(g) for g in training_gpus)
276
+ print(f"[init] Colocating training and inference on GPUs [{device_ids}].")
277
+
278
+ server_proc = start_vllm_server(local_model_path, device_ids)
279
+ try:
280
+ client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY")
281
+
282
+ print("[generate] Generating with dummy weights...")
283
+ outputs = generate_completions(client, PROMPTS)
284
+ print_generations("BEFORE weight sync (dummy weights):", PROMPTS, outputs)
285
+
286
+ # --- Weight transfer ---
287
+ print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...")
288
+ ray.get([w.setup_engine.remote(BASE_URL) for w in fsdp_workers])
289
+
290
+ print("[sync] Sleeping engine (offload weights + free KV cache)...")
291
+ sleep_engine(BASE_URL, level=1)
292
+
293
+ print("[sync] Waking weights (KV cache stays free)...")
294
+ wake_up_engine(BASE_URL, tags=["weights"])
295
+
296
+ print("[sync] Packed IPC transfer FSDP -> vLLM...")
297
+ ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers])
298
+ print("[sync] Weight transfer complete.")
299
+
300
+ print("[sync] Waking KV cache + scheduling...")
301
+ wake_up_engine(BASE_URL, tags=["kv_cache", "scheduling"])
302
+
303
+ print("[generate] Generating with synced weights...")
304
+ outputs_updated = generate_completions(client, PROMPTS)
305
+ print_generations("AFTER weight sync (real weights):", PROMPTS, outputs_updated)
306
+ finally:
307
+ print("[server] Shutting down...")
308
+ server_proc.terminate()
309
+ try:
310
+ server_proc.wait(timeout=30)
311
+ except subprocess.TimeoutExpired:
312
+ server_proc.kill()
313
314
315
if __name__ == "__main__":
examples/rl/rlhf_nccl.py
deleted
-196
@@ -1,196 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Demonstrates reinforcement learning using vLLM and Ray,
5
-with native weight syncing APIs at engine instance.
6
-
7
-The script separates training and inference workloads onto distinct GPUs
8
-so that Ray can manage process placement and inter-process communication.
9
-A Hugging Face Transformer model occupies one GPU for training, whereas a
10
-2x tensor-parallel vLLM inference engine occupies two GPUs.
11
-
12
-The example performs the following steps:
13
-* Load the training model on one gpu (scheduled via ray)
14
-* Initialize the inference model with dummy weights across
15
- two gpus using vLLM's tensor parallelism and Ray placement groups.
16
-* Generate gibberish from a list of prompts using the randomly initialized
17
- inference engine.
18
-* Update the weights of the training model and broadcast the updated weights
19
- to the inference engine by using a Ray collective RPC group.
20
-* Generating from the list of prompts after weight sync should result
21
- in sensible outputs.
22
-
23
-This example assumes a single-node cluster with three GPUs, but Ray
24
-supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
25
-workloads. Residual GPU activity interferes with vLLM memory profiling and
26
-causes unexpected behavior.
27
-"""
28
-
29
-import os
30
-
31
-import ray
32
-import torch
33
-from ray.util.placement_group import placement_group
34
-from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
35
-from transformers import AutoModelForCausalLM
36
-
37
-from vllm import LLM, SamplingParams
38
-from vllm.config import WeightTransferConfig
39
-from vllm.distributed.weight_transfer import (
40
- ModuleSource,
41
- RayVLLMWeightSyncClient,
42
- WeightTransferTrainerFactory,
43
-)
44
-from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
45
-from vllm.platforms import current_platform
46
-from vllm.utils.network_utils import get_ip, get_open_port
47
-
48
-MODEL_NAME = "facebook/opt-125m"
49
-# MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128"
50
-
51
-
52
-def get_assigned_gpu():
53
- """This is a temporary workaround for a runtime bug in RCCL on ROCm."""
54
- if not current_platform.is_rocm():
55
- return 0
56
- assigned_gpu = int(ray.get_gpu_ids()[0])
57
- os.environ.pop("CUDA_VISIBLE_DEVICES", None)
58
- os.environ.pop("HIP_VISIBLE_DEVICES", None)
59
- torch.accelerator.set_device_idx(assigned_gpu)
60
- return assigned_gpu
61
-
62
-
63
-class MyLLM(LLM):
64
- """Configure the vLLM worker for Ray placement group execution."""
65
-
66
- def __init__(self, *args, **kwargs):
67
- os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
68
- super().__init__(*args, **kwargs)
69
-
70
-
71
-@ray.remote(num_gpus=1)
72
-class TrainModel:
73
- """Ray actor that wraps the training model on a dedicated GPU."""
74
-
75
- def __init__(self, model_name: str):
76
- assigned_gpu = get_assigned_gpu()
77
-
78
- self.model = AutoModelForCausalLM.from_pretrained(
79
- model_name,
80
- ).to(f"cuda:{assigned_gpu}")
81
-
82
- self.port = get_open_port()
83
- self.master_address = get_ip()
84
-
85
- def init_weight_transfer(self, world_size, llm_handle):
86
- """Build the trainer-side weight-transfer engine.
87
-
88
- `trainer_init` drives the full handshake: it kicks off the inference
89
- side's `init_weight_transfer_engine` (via the Ray client) on a worker
90
- thread while opening the trainer's NCCL endpoint, so both ends
91
- rendezvous together. After this returns, `send_weights()` is callable.
92
- """
93
- self.engine = WeightTransferTrainerFactory.trainer_init(
94
- init_info=NCCLTrainerInitInfo(
95
- master_address=self.master_address,
96
- master_port=self.port,
97
- world_size=world_size,
98
- rank=0, # single-GPU trainer is the sole (sender) rank
99
- packed=True,
100
- ),
101
- client=RayVLLMWeightSyncClient(llm_handle),
102
- source=ModuleSource(self.model),
103
- )
104
-
105
- def broadcast_weights(self):
106
- """Push the current weights to the inference engine.
107
-
108
- Drives start/update/finish on the inference side and the NCCL
109
- broadcast internally — one call.
110
- """
111
- self.engine.send_weights()
112
-
113
-
114
-# Initialize Ray and set the visible devices. The vLLM engine will
115
-# be placed on GPUs 1 and 2.
116
-ray.init()
117
-
118
-# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
119
-# Learn more about Ray placement groups:
120
-# https://docs.ray.io/en/latest/placement-groups.html
121
-# Launch the training model actor. Ray's resource scheduler will allocate
122
-# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
123
-train_model = TrainModel.remote(MODEL_NAME)
124
-
125
-pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
126
-ray.get(pg_inference.ready())
127
-scheduling_inference = PlacementGroupSchedulingStrategy(
128
- placement_group=pg_inference,
129
- placement_group_capture_child_tasks=True,
130
- placement_group_bundle_index=0,
131
-)
132
-
133
-# Launch the vLLM inference engine. The `enforce_eager` flag reduces
134
-# start-up latency.
135
-# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
136
-# are now native to vLLM workers.
137
-llm = ray.remote(
138
- num_cpus=0,
139
- num_gpus=0,
140
- scheduling_strategy=scheduling_inference,
141
-)(MyLLM).remote(
142
- model=MODEL_NAME,
143
- enforce_eager=True,
144
- tensor_parallel_size=2,
145
- data_parallel_size=1,
146
- distributed_executor_backend="ray",
147
- weight_transfer_config=WeightTransferConfig(backend="nccl"),
148
- load_format="dummy",
149
- quantization="fp8",
150
-)
151
-
152
-# Generate text from the prompts.
153
-prompts = [
154
- "Hello, my name is",
155
- "The president of the United States is",
156
- "The capital of France is",
157
- "The future of AI is",
158
-]
159
-
160
-sampling_params = SamplingParams(temperature=0)
161
-
162
-outputs = ray.get(llm.generate.remote(prompts, sampling_params))
163
-
164
-# Generate text with the initial model. The output is expected to be nonsense
165
-# because the weights are randomly initialized.
166
-print("-" * 50)
167
-for output in outputs:
168
- prompt = output.prompt
169
- generated_text = output.outputs[0].text
170
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
171
- print("-" * 50)
172
-
173
-ray.get(llm.sleep.remote(level=0))
174
-
175
-# Set up the communication channel between the training process and the
176
-# inference engine. The trainer engine now owns the full handshake — the
177
-# driver only has to hand it the vLLM actor handle and the world size.
178
-world_size = ray.get(llm.get_world_size.remote()) + 1 # +1 for the trainer
179
-ray.get(train_model.init_weight_transfer.remote(world_size, llm))
180
-
181
-# Synchronize the updated weights to the inference engine. The engine drives
182
-# start_weight_update / update_weights / finish_weight_update on the inference
183
-# side internally, concurrent with the NCCL broadcast.
184
-ray.get(train_model.broadcast_weights.remote())
185
-
186
-ray.get(llm.wake_up.remote(tags=["scheduling"]))
187
-
188
-# Generate text with the updated model. The output is expected to be normal
189
-# because the weights are updated.
190
-outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
191
-print("-" * 50)
192
-for output in outputs_updated:
193
- prompt = output.prompt
194
- generated_text = output.outputs[0].text
195
- print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
196
- print("-" * 50)
tests/entrypoints/weight_transfer/test_weight_transfer_llm.py
-5
@@ -92,11 +92,6 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
92
def shutdown(self) -> None:
93
MockWeightTransferEngine.shutdown_called = True
94
95
- @staticmethod
96
- def trainer_send_weights(*args, **kwargs):
97
- """Mock method to simulate trainer sending weights."""
98
- pass
99
-
95
96
def mock_create_engine(config, vllm_config, device, model):
97
"""Mock factory function that returns our mock engine."""
vllm/distributed/weight_transfer/base.py
+9
-27
@@ -77,7 +77,15 @@ class WeightSource(ABC):
77
78
@abstractmethod
79
def metadata(self) -> list[ParamMeta]:
80
- """Declare what iteration will yield, without transferring anything."""
80
+ """Declare what iteration will yield, without transferring anything.
81
+
82
+ Must agree with iteration element for element: the same parameters, in
83
+ the same order, with the same dtypes and shapes. Backends may read both
84
+ channels and trust that they match (dense NCCL sizes the worker's receive
85
+ buffers and its packed chunk boundaries from this, then sends the bytes
86
+ from iteration), so a source that disagrees between the two splits the
87
+ stream differently on each side.
88
+ """
89
raise NotImplementedError
90
91
@abstractmethod
@@ -342,32 +350,6 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
350
"""
351
raise NotImplementedError
352
345
- @staticmethod
346
- @abstractmethod
347
- def trainer_send_weights(
348
- iterator: Iterator[Any],
349
- trainer_args: dict[str, Any] | Any,
350
- ) -> None:
351
- """
352
- Send weights from trainer to inference workers.
353
-
354
- This is a static method that can be called from the trainer process
355
- to send weights to all inference workers.
356
-
357
- Args:
358
- iterator: Iterator of backend-specific items to send.
359
- trainer_args: Dictionary containing backend-specific arguments needed
360
- to send weights. The structure depends on the backend:
361
- - NCCL: Contains 'group', 'src', 'packed', etc.
362
- - IPC: Contains 'mode' ('http' or 'ray'),
363
- 'llm_handle' (for Ray), 'url' (for HTTP), etc.
364
-
365
- Example:
366
- >>> param_iter = ((n, p) for n, p in model.named_parameters())
367
- >>> engine.trainer_send_weights(param_iter, trainer_args)
368
- """
369
- raise NotImplementedError
370
-
353
354
@runtime_checkable
355
class VLLMWeightSyncClient(Protocol):
vllm/distributed/weight_transfer/ipc_engine.py
-15
@@ -249,21 +249,6 @@ class IPCWeightTransferEngine(
249
def shutdown(self) -> None:
250
self._packed_importer.close()
251
252
- @staticmethod
253
- def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
254
- """Removed. Use the stateful `IPCTrainerWeightTransferEngine` instead.
255
-
256
- Transitional stub kept only to satisfy the (still abstract)
257
- `WeightTransferEngine.trainer_send_weights`; that member is dropped from
258
- the worker ABC once every backend has migrated to the trainer engine.
259
- """
260
- raise NotImplementedError(
261
- "The static IPC trainer path has been replaced by "
262
- "IPCTrainerWeightTransferEngine. Build it via "
263
- "WeightTransferTrainerFactory.trainer_init(IPCTrainerInitInfo(...), "
264
- "client=..., source=...) and drive it with send_weights()."
265
- )
266
-
252
253
class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitInfo]):
254
"""Trainer-side CUDA IPC weight transfer engine.
vllm/distributed/weight_transfer/nccl_engine.py
+1
-16
@@ -5,7 +5,7 @@
5
from collections.abc import Iterator
6
from concurrent.futures import ThreadPoolExecutor
7
from dataclasses import asdict, dataclass
8
-from typing import TYPE_CHECKING, Any, ClassVar
8
+from typing import TYPE_CHECKING, ClassVar
9
10
import torch
11
from typing_extensions import Self
@@ -226,21 +226,6 @@ class NCCLWeightTransferEngine(
226
# Clean up the communicator by removing the reference
227
self.model_update_group = None
228
229
- @staticmethod
230
- def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
231
- """Removed. Use the stateful `NCCLTrainerWeightTransferEngine` instead.
232
-
233
- Transitional stub kept only to satisfy the (still abstract)
234
- `WeightTransferEngine.trainer_send_weights`; that member is dropped from
235
- the worker ABC once every backend has migrated to the trainer engine.
236
- """
237
- raise NotImplementedError(
238
- "The static NCCL trainer path has been replaced by "
239
- "NCCLTrainerWeightTransferEngine. Build it via "
240
- "WeightTransferTrainerFactory.trainer_init(NCCLTrainerInitInfo(...), "
241
- "client=..., source=...) and drive it with send_weights()."
242
- )
243
-
229
230
class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerInitInfo]):
231
"""Trainer-side NCCL weight transfer engine.
vllm/distributed/weight_transfer/sparse_nccl_engine.py
+1
-17
@@ -17,7 +17,7 @@ MVP limitations:
17
from collections.abc import Iterable
18
from concurrent.futures import ThreadPoolExecutor
19
from dataclasses import asdict, dataclass
20
-from typing import TYPE_CHECKING, Any, ClassVar
20
+from typing import TYPE_CHECKING, ClassVar
21
22
import torch
23
from typing_extensions import Self
@@ -227,22 +227,6 @@ class SparseNCCLWeightTransferEngine(
227
if self.model_update_group is not None:
228
self.model_update_group = None
229
230
- @staticmethod
231
- def trainer_send_weights(*args: Any, **kwargs: Any) -> None:
232
- """Removed. Use the stateful `SparseNCCLTrainerWeightTransferEngine`.
233
-
234
- Transitional stub kept only to satisfy the (still abstract)
235
- `WeightTransferEngine.trainer_send_weights`; that member is dropped from
236
- the worker ABC once every backend has migrated to the trainer engine.
237
- """
238
- raise NotImplementedError(
239
- "The static sparse NCCL trainer path has been replaced by "
240
- "SparseNCCLTrainerWeightTransferEngine. Build it via "
241
- "WeightTransferTrainerFactory.trainer_init("
242
- "SparseNCCLTrainerInitInfo(...), client=..., source=...) and drive "
243
- "it with send_weights()."
244
- )
245
-
230
231
class SparseNCCLTrainerWeightTransferEngine(
232
TrainerWeightTransferEngine[SparseNCCLTrainerInitInfo]