[rl] Stateful Trainer Send: New Abstractions [1/N] (#48042)
Signed-off-by: haoaaron <ahao@anyscale.com> Signed-off-by: Aaron Hao <ahao@anyscale.com> Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>
Aaron Hao committed
Jul 17, 2026 at 00:11 UTC
fb1d8ccaf5bc46935031ef074950eaa8232cd15e
7 files changed
+644
-24
tests/distributed/test_weight_transfer.py
+204
-1
@@ -17,7 +17,19 @@ from torch.multiprocessing.reductions import reduce_tensor
17
18
from vllm.config.parallel import ParallelConfig
19
from vllm.config.weight_transfer import WeightTransferConfig
20
-from vllm.distributed.weight_transfer import WeightTransferEngineFactory
20
+from vllm.distributed.weight_transfer import (
21
+ HTTPVLLMWeightSyncClient,
22
+ ModuleSource,
23
+ RayVLLMWeightSyncClient,
24
+ TrainerWeightTransferEngine,
25
+ VLLMWeightSyncClient,
26
+ WeightTransferEngineFactory,
27
+ WeightTransferTrainerFactory,
28
+)
29
+from vllm.distributed.weight_transfer.base import (
30
+ WeightTransferInitRequest,
31
+ WeightTransferUpdateRequest,
32
+)
33
from vllm.distributed.weight_transfer.ipc_engine import (
34
IPCWeightTransferEngine,
35
IPCWeightTransferInitInfo,
@@ -1214,3 +1226,194 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
1226
1227
with pytest.raises(ValueError, match="IPC handle not found"):
1228
engine.receive_weights(update_info)
1229
+
1230
+
1231
+class RecordingClient:
1232
+ """A fake VLLMWeightSyncClient that records the order of calls."""
1233
+
1234
+ def __init__(self):
1235
+ self.order: list[str] = []
1236
+ self.last_init_info: dict | None = None
1237
+ self.last_update_info: dict | None = None
1238
+
1239
+ def init_weight_transfer_engine(self, init_info: dict) -> None:
1240
+ self.order.append("init")
1241
+ self.last_init_info = init_info
1242
+
1243
+ def start_weight_update(self) -> None:
1244
+ self.order.append("start")
1245
+
1246
+ def update_weights(self, update_info: dict) -> None:
1247
+ self.order.append("update")
1248
+ self.last_update_info = update_info
1249
+
1250
+ def finish_weight_update(self) -> None:
1251
+ self.order.append("finish")
1252
+
1253
+
1254
+def _module_with(*pairs):
1255
+ """A tiny nn.Module exposing the given (name, tensor) pairs as parameters,
1256
+ so trainer tests can build a ModuleSource without a real model."""
1257
+ module = torch.nn.Module()
1258
+ for name, tensor in pairs:
1259
+ module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False))
1260
+ return module
1261
+
1262
+
1263
+class _DummyTrainerEngine(TrainerWeightTransferEngine):
1264
+ """Minimal concrete trainer engine to exercise base-class + factory."""
1265
+
1266
+ @classmethod
1267
+ def trainer_init(cls, config, init_info, *, client, source):
1268
+ return cls(config, client=client, source=source)
1269
+
1270
+ def send_weights(self):
1271
+ pass
1272
+
1273
+
1274
+class TestTrainerClients:
1275
+ """Structural protocol conformance for the built-in clients."""
1276
+
1277
+ def test_recording_client_is_protocol(self):
1278
+ assert isinstance(RecordingClient(), VLLMWeightSyncClient)
1279
+
1280
+ def test_http_client_is_protocol(self):
1281
+ assert isinstance(
1282
+ HTTPVLLMWeightSyncClient("http://localhost:8000"), VLLMWeightSyncClient
1283
+ )
1284
+
1285
+ def test_ray_client_is_protocol(self):
1286
+ assert isinstance(RayVLLMWeightSyncClient(MagicMock()), VLLMWeightSyncClient)
1287
+
1288
+ def test_ray_client_sends_typed_requests(self, monkeypatch):
1289
+ """Ray client must hand the actor typed Request objects, not raw dicts."""
1290
+ import ray
1291
+
1292
+ monkeypatch.setattr(ray, "get", lambda refs: None)
1293
+ handle = MagicMock()
1294
+ client = RayVLLMWeightSyncClient(handle)
1295
+
1296
+ client.init_weight_transfer_engine({"master_addr": "x"})
1297
+ (init_req,), _ = handle.init_weight_transfer_engine.remote.call_args
1298
+ assert isinstance(init_req, WeightTransferInitRequest)
1299
+ assert init_req.init_info == {"master_addr": "x"}
1300
+
1301
+ client.update_weights({"names": ["w"]})
1302
+ (update_req,), _ = handle.update_weights.remote.call_args
1303
+ assert isinstance(update_req, WeightTransferUpdateRequest)
1304
+ assert update_req.update_info == {"names": ["w"]}
1305
+
1306
+ def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch):
1307
+ """HTTP update_weights must encode raw ipc_handles as a base64 pickle."""
1308
+ captured = {}
1309
+
1310
+ def fake_post(self, path, json=None):
1311
+ captured["path"] = path
1312
+ captured["json"] = json
1313
+
1314
+ monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post)
1315
+ client = HTTPVLLMWeightSyncClient("http://localhost:8000")
1316
+ client.update_weights({"names": ["w"], "ipc_handles": [{"gpu": ("args",)}]})
1317
+ sent = captured["json"]["update_info"]
1318
+ assert "ipc_handles" not in sent
1319
+ assert "ipc_handles_pickled" in sent
1320
+ assert pickle.loads(base64.b64decode(sent["ipc_handles_pickled"])) == [
1321
+ {"gpu": ("args",)}
1322
+ ]
1323
+
1324
+ def test_http_client_passes_through_nccl_update_info(self, monkeypatch):
1325
+ """NCCL update_info has only JSON-native fields and passes unchanged."""
1326
+ captured = {}
1327
+
1328
+ def fake_post(self, path, json=None):
1329
+ captured["json"] = json
1330
+
1331
+ monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post)
1332
+ client = HTTPVLLMWeightSyncClient("http://localhost:8000")
1333
+ update_info = {"names": ["w"], "dtype_names": ["float32"], "shapes": [[4]]}
1334
+ client.update_weights(update_info)
1335
+ assert captured["json"]["update_info"] == update_info
1336
+
1337
+
1338
+class TestModuleSource:
1339
+ """`ModuleSource` metadata vs. materialized iteration (dense, no GPU)."""
1340
+
1341
+ def test_metadata_reads_shape_and_dtype(self):
1342
+ source = ModuleSource(
1343
+ _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3)))
1344
+ )
1345
+ meta = source.metadata()
1346
+ assert [m.name for m in meta] == ["w", "b"]
1347
+ assert [m.shape for m in meta] == [(2, 3), (3,)]
1348
+ assert all(m.dtype == torch.float32 for m in meta)
1349
+
1350
+ def test_iteration_yields_materialized_tensors(self):
1351
+ w = torch.arange(6, dtype=torch.float32).reshape(2, 3)
1352
+ source = ModuleSource(_module_with(("w", w)))
1353
+ pairs = list(source)
1354
+ assert [name for name, _ in pairs] == ["w"]
1355
+ assert torch.equal(pairs[0][1], w)
1356
+
1357
+ def test_source_is_reiterable(self):
1358
+ source = ModuleSource(_module_with(("w", torch.zeros(2))))
1359
+ assert [n for n, _ in source] == [n for n, _ in source] == ["w"]
1360
+
1361
+
1362
+class TestTrainerFactory:
1363
+ """WeightTransferTrainerFactory registry mechanics."""
1364
+
1365
+ def test_builtin_registry_has_no_trainer_backends_yet(self):
1366
+ # Concrete backends register in the per-backend migration PRs.
1367
+ assert WeightTransferTrainerFactory._registry == {}
1368
+
1369
+ def test_register_and_dispatch(self):
1370
+ saved = dict(WeightTransferTrainerFactory._registry)
1371
+ try:
1372
+ WeightTransferTrainerFactory.register_engine("dummy", _DummyTrainerEngine)
1373
+ engine = WeightTransferTrainerFactory.trainer_init(
1374
+ "dummy",
1375
+ WeightTransferConfig(backend="dummy"),
1376
+ MagicMock(),
1377
+ client=RecordingClient(),
1378
+ source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1379
+ )
1380
+ assert isinstance(engine, _DummyTrainerEngine)
1381
+ with pytest.raises(ValueError, match="already registered"):
1382
+ WeightTransferTrainerFactory.register_engine(
1383
+ "dummy", _DummyTrainerEngine
1384
+ )
1385
+ finally:
1386
+ WeightTransferTrainerFactory._registry = saved
1387
+
1388
+ def test_unknown_backend_raises(self):
1389
+ with pytest.raises(ValueError, match="Invalid weight transfer backend"):
1390
+ WeightTransferTrainerFactory.trainer_init(
1391
+ "nope",
1392
+ WeightTransferConfig(backend="nope"),
1393
+ MagicMock(),
1394
+ client=RecordingClient(),
1395
+ source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1396
+ )
1397
+
1398
+
1399
+class TestTrainerEngineBase:
1400
+ """Base-class construction (no GPU)."""
1401
+
1402
+ def test_source_stored_and_sender_by_default(self):
1403
+ engine = _DummyTrainerEngine(
1404
+ WeightTransferConfig(backend="nccl"),
1405
+ client=RecordingClient(),
1406
+ source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1407
+ )
1408
+ assert engine.is_sender is True
1409
+ assert [name for name, _ in engine.source] == ["w"]
1410
+
1411
+ def test_shutdown_default_is_noop(self):
1412
+ engine = _DummyTrainerEngine(
1413
+ WeightTransferConfig(backend="nccl"),
1414
+ client=RecordingClient(),
1415
+ source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1416
+ is_sender=False,
1417
+ )
1418
+ assert engine.is_sender is False
1419
+ engine.shutdown() # must not raise
tools/pre_commit/check_forbidden_imports.py
+1
@@ -39,6 +39,7 @@ CHECK_IMPORTS = {
39
"vllm/distributed/device_communicators/shm_broadcast.py",
40
"vllm/distributed/device_communicators/shm_object_storage.py",
41
"vllm/distributed/weight_transfer/ipc_engine.py",
42
+ "vllm/distributed/weight_transfer/clients.py",
43
"tests/distributed/test_weight_transfer.py",
44
"vllm/utils/hashing.py",
45
"tests/multimodal/media/test_base.py",
vllm/distributed/weight_transfer/__init__.py
+24
-2
@@ -5,10 +5,32 @@ Weight transfer engines for syncing model weights from trainers
5
to inference workers.
6
"""
7
8
-from vllm.distributed.weight_transfer.base import WeightTransferEngine
9
-from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
8
+from vllm.distributed.weight_transfer.base import (
9
+ ModuleSource,
10
+ ParamMeta,
11
+ TrainerWeightTransferEngine,
12
+ VLLMWeightSyncClient,
13
+ WeightSource,
14
+ WeightTransferEngine,
15
+)
16
+from vllm.distributed.weight_transfer.clients import (
17
+ HTTPVLLMWeightSyncClient,
18
+ RayVLLMWeightSyncClient,
19
+)
20
+from vllm.distributed.weight_transfer.factory import (
21
+ WeightTransferEngineFactory,
22
+ WeightTransferTrainerFactory,
23
+)
24
25
__all__ = [
26
"WeightTransferEngine",
27
"WeightTransferEngineFactory",
28
+ "TrainerWeightTransferEngine",
29
+ "WeightTransferTrainerFactory",
30
+ "VLLMWeightSyncClient",
31
+ "HTTPVLLMWeightSyncClient",
32
+ "RayVLLMWeightSyncClient",
33
+ "ParamMeta",
34
+ "WeightSource",
35
+ "ModuleSource",
36
]
vllm/distributed/weight_transfer/base.py
+202
-1
@@ -5,9 +5,10 @@
5
from abc import ABC, abstractmethod
6
from collections.abc import Iterator
7
from dataclasses import dataclass, field
8
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
8
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable
9
10
import torch
11
+from typing_extensions import Self
12
13
if TYPE_CHECKING:
14
from vllm.config import VllmConfig
@@ -17,6 +18,85 @@ from vllm.config.weight_transfer import WeightTransferConfig
18
19
TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo")
20
TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo")
21
+TConfig = TypeVar("TConfig", bound="WeightTransferConfig")
22
+
23
+# A trainer supplies its parameters as a `WeightSource` (defined below): a
24
+# re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()`
25
+# channel. The built-in `ModuleSource` uses `materialize_full_tensor`.
26
+
27
+
28
+def materialize_full_tensor(tensor: torch.Tensor) -> torch.Tensor:
29
+ """Return a full, locally-materialized tensor ready to send.
30
+
31
+ FSDP shards (DTensors) expose `full_tensor()`, a collective all-gather;
32
+ regular tensors do not and are returned unchanged. Trainer engines call
33
+ this at send time so the (potentially expensive) gather happens exactly
34
+ once — reading `.shape`/`.dtype` for metadata does not trigger it.
35
+ """
36
+ full_tensor = getattr(tensor, "full_tensor", None)
37
+ return full_tensor() if callable(full_tensor) else tensor
38
+
39
+
40
+@dataclass(frozen=True)
41
+class ParamMeta:
42
+ """Name / wire dtype / full (HF) shape for one output parameter."""
43
+
44
+ name: str
45
+ dtype: torch.dtype
46
+ shape: tuple[int, ...]
47
+
48
+
49
+class WeightSource(ABC):
50
+ """A re-iterable source of the trainer's weights, handed to a trainer engine.
51
+
52
+ Two channels:
53
+
54
+ * `metadata()` — `(name, wire dtype, full shape)` for every parameter,
55
+ *without* transferring. Cheap when shapes are known locally (FSDP
56
+ `DTensor` global shape); may be expensive on first call for backends that
57
+ must materialize to learn shapes (e.g. a Megatron-Bridge export), in which
58
+ case it should cache.
59
+ * iteration — yields fully-materialized `(name, tensor)` pairs, one at a
60
+ time. Materializing is typically a collective (FSDP `full_tensor()`, a
61
+ Megatron export), so every trainer rank must iterate the same source in the
62
+ same order in lockstep, or ranks deadlock. Under pipeline parallelism a
63
+ rank may not own a parameter at all — iterating still drives the collective
64
+ and the yielded tensor is only meaningful on the sender.
65
+
66
+ `iter(source)` must yield a *fresh* pass each round. Backends with custom
67
+ producer logic (Megatron export, RDT plans, MoE re-fusing) subclass this.
68
+ """
69
+
70
+ @abstractmethod
71
+ def metadata(self) -> list[ParamMeta]:
72
+ raise NotImplementedError
73
+
74
+ @abstractmethod
75
+ def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
76
+ raise NotImplementedError
77
+
78
+
79
+class ModuleSource(WeightSource):
80
+ """`WeightSource` over `module.named_parameters()` — the common case.
81
+
82
+ Handles both plain dense modules and FSDP-sharded ones with no special
83
+ casing: iteration all-gathers each `DTensor` via `full_tensor()` (a
84
+ collective) and passes regular tensors through. `metadata()` reads the
85
+ *global* `.shape` / `.dtype`, so it never triggers a gather.
86
+ """
87
+
88
+ def __init__(self, module: torch.nn.Module) -> None:
89
+ self._module = module
90
+
91
+ def metadata(self) -> list[ParamMeta]:
92
+ return [
93
+ ParamMeta(name, p.dtype, tuple(p.shape))
94
+ for name, p in self._module.named_parameters()
95
+ ]
96
+
97
+ def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
98
+ for name, param in self._module.named_parameters():
99
+ yield name, materialize_full_tensor(param)
100
101
102
# Base protocols for backend-specific dataclasses
@@ -27,6 +107,26 @@ class WeightTransferInitInfo(ABC): # noqa: B024
107
pass
108
109
110
+@dataclass
111
+class TrainerInitInfo(WeightTransferInitInfo):
112
+ """Base trainer-side init info: which trainer rank drives the transfer.
113
+
114
+ `rank` is this trainer process's rank, provided **explicitly** by the
115
+ caller — the engine does not read it from a global process group, which is
116
+ ambiguous once several groups (FSDP / TP / PP / EP) exist. Rank 0 is always
117
+ the sender: only it opens the endpoint and drives the inference-side RPCs,
118
+ while every rank still runs the trainer-side collectives. Backend subclasses
119
+ add their own (positional) fields; `rank` is keyword-only so that ordering
120
+ never conflicts.
121
+ """
122
+
123
+ rank: int = field(kw_only=True)
124
+
125
+ @property
126
+ def is_sender(self) -> bool:
127
+ return self.rank == 0
128
+
129
+
130
@dataclass
131
class WeightTransferUpdateInfo(ABC): # noqa: B024
132
"""Base class for backend-specific weight update info."""
@@ -243,3 +343,104 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
343
>>> engine.trainer_send_weights(param_iter, trainer_args)
344
"""
345
raise NotImplementedError
346
+
347
+
348
+@runtime_checkable
349
+class VLLMWeightSyncClient(Protocol):
350
+ """Trainer-side stub for the inference engine's weight-sync control plane.
351
+
352
+ Mirrors the weight-sync methods that the inference engine exposes
353
+ (`EngineClient` / the HTTP RLHF routes / Ray actors). A
354
+ `TrainerWeightTransferEngine` drives the full handshake through this
355
+ protocol so trainer code never has to know the transport.
356
+
357
+ All methods are synchronous and accept plain dicts (matching what the
358
+ inference side already accepts). Concurrency that some backends need
359
+ (e.g. NCCL must run `update_weights` concurrently with the trainer-side
360
+ broadcast) is the engine's responsibility, not the client's, so the
361
+ protocol stays a flat four-method surface that any wrapper can implement.
362
+
363
+ The protocol is structural (PEP 544), so user implementations need only
364
+ define these four methods — no import or subclassing required.
365
+ """
366
+
367
+ def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ...
368
+
369
+ def start_weight_update(self) -> None: ...
370
+
371
+ def update_weights(self, update_info: dict[str, Any]) -> None: ...
372
+
373
+ def finish_weight_update(self) -> None: ...
374
+
375
+
376
+class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):
377
+ """Trainer-side weight transfer engine.
378
+
379
+ Symmetric to `WeightTransferEngine` but lives in the training process.
380
+ Constructed via the `trainer_init` factory classmethod; carries any
381
+ backend-specific state (NCCL communicators, IPC device info, transfer
382
+ plans) on `self`. The `WeightSource` is required at `trainer_init`,
383
+ then replayed each round by the no-argument `send_weights()`.
384
+
385
+ Multi-rank trainers: `trainer_init` and `send_weights` are
386
+ called on *every* trainer rank. Rank 0 is the sender, resolved once at
387
+ `trainer_init` into `is_sender`. Non-sender ranks still run every
388
+ collective (iterating the source, metadata export, IPC handle all-gather) so
389
+ the group stays aligned, but each engine explicitly guards the control-plane
390
+ RPCs and the transmit on `self.is_sender`, so only the sender touches the
391
+ client.
392
+
393
+ Subclasses should define:
394
+ init_info_cls: Type of backend-specific trainer init info
395
+ config_cls: Type of backend-specific config
396
+ """
397
+
398
+ # Subclasses should override these class attributes
399
+ init_info_cls: type[TInitInfo]
400
+ config_cls: type[TConfig]
401
+
402
+ def __init__(
403
+ self,
404
+ config: TConfig,
405
+ *,
406
+ client: "VLLMWeightSyncClient",
407
+ source: "WeightSource",
408
+ is_sender: bool = True,
409
+ ) -> None:
410
+ self.config = config
411
+ self.is_sender = is_sender
412
+ # The real client is held on every rank; each engine only *calls* it when
413
+ # `is_sender`, so non-sender ranks never touch the wire.
414
+ self.client = client
415
+ self.source = source
416
+
417
+ @classmethod
418
+ @abstractmethod
419
+ def trainer_init(
420
+ cls,
421
+ config: TConfig,
422
+ init_info: TInitInfo,
423
+ *,
424
+ client: "VLLMWeightSyncClient",
425
+ source: "WeightSource",
426
+ ) -> Self:
427
+ """Rendezvous with the inference side and return a ready instance.
428
+
429
+ Called on every trainer rank. The sender drives the full handshake via
430
+ `client` (build the worker-side init info, call
431
+ `client.init_weight_transfer_engine`, open the trainer-side endpoint);
432
+ non-sender ranks skip the rendezvous and the RPC. `source` is stored on
433
+ `self.source`; after return, `send_weights()` is callable.
434
+ """
435
+ raise NotImplementedError
436
+
437
+ @abstractmethod
438
+ def send_weights(self) -> None:
439
+ """Push `self.source`'s weights to inference workers and drive the full
440
+ update round trip: `start_weight_update`, `update_weights` (run
441
+ concurrently with the trainer-side broadcast when the backend requires
442
+ it), then `finish_weight_update`. Called on every trainer rank."""
443
+ raise NotImplementedError
444
+
445
+ def shutdown(self) -> None:
446
+ """Tear down communicators / process groups. Default no-op."""
vllm/distributed/weight_transfer/clients.py
new
+114
@@ -0,0 +1,114 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Built-in `VLLMWeightSyncClient` implementations.
4
+
5
+These adapt the inference engine's weight-sync control plane to concrete
6
+transports. A `TrainerWeightTransferEngine` takes one of these (or any object
7
+with the same four methods — the protocol is structural) and drives the full
8
+handshake through it.
9
+
10
+Imports of `ray` / `requests` are deferred to call time so this module is
11
+importable without those packages installed.
12
+"""
13
+
14
+from typing import TYPE_CHECKING, Any
15
+
16
+from vllm.distributed.weight_transfer.base import (
17
+ WeightTransferInitRequest,
18
+ WeightTransferUpdateRequest,
19
+)
20
+
21
+if TYPE_CHECKING:
22
+ from ray.actor import ActorHandle
23
+
24
+
25
+def _json_safe_update_info(update_info: dict[str, Any]) -> dict[str, Any]:
26
+ """Make an update_info dict JSON-serializable for HTTP transport.
27
+
28
+ CUDA IPC handles (`ipc_handles`) are tuples of non-JSON-native objects, so
29
+ over HTTP they are pickled+base64-encoded into `ipc_handles_pickled` (which
30
+ the worker auto-deserializes when `VLLM_ALLOW_INSECURE_SERIALIZATION=1`).
31
+ Other backends (NCCL) carry only JSON-native metadata and pass through
32
+ unchanged. Mirrors the old IPC `_do_send` HTTP branch.
33
+ """
34
+ ipc_handles = update_info.get("ipc_handles")
35
+ if ipc_handles is None:
36
+ return update_info
37
+
38
+ import pickle
39
+
40
+ import pybase64 as base64
41
+
42
+ out = {k: v for k, v in update_info.items() if k != "ipc_handles"}
43
+ out["ipc_handles_pickled"] = base64.b64encode(pickle.dumps(ipc_handles)).decode(
44
+ "utf-8"
45
+ )
46
+ return out
47
+
48
+
49
+class HTTPVLLMWeightSyncClient:
50
+ """Talks to a vLLM server over the RLHF HTTP routes.
51
+
52
+ Mirrors `vllm/entrypoints/serve/dev/rlhf/api_router.py`:
53
+ `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`,
54
+ `/finish_weight_update`.
55
+ """
56
+
57
+ def __init__(self, base_url: str, timeout: float = 300) -> None:
58
+ self.base_url = base_url.rstrip("/")
59
+ self.timeout = timeout
60
+
61
+ def _post(self, path: str, json: dict[str, Any] | None = None) -> None:
62
+ import requests
63
+
64
+ response = requests.post(
65
+ f"{self.base_url}/{path}", json=json, timeout=self.timeout
66
+ )
67
+ response.raise_for_status()
68
+
69
+ def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None:
70
+ self._post("init_weight_transfer_engine", {"init_info": init_info})
71
+
72
+ def start_weight_update(self) -> None:
73
+ self._post("start_weight_update")
74
+
75
+ def update_weights(self, update_info: dict[str, Any]) -> None:
76
+ self._post(
77
+ "update_weights", {"update_info": _json_safe_update_info(update_info)}
78
+ )
79
+
80
+ def finish_weight_update(self) -> None:
81
+ self._post("finish_weight_update")
82
+
83
+
84
+class RayVLLMWeightSyncClient:
85
+ """Talks to one or more vLLM `AsyncLLM`/`LLM` Ray actors.
86
+
87
+ Each call fans out to every handle and blocks on all of them, so a
88
+ multi-actor (e.g. multi-DP) deployment is driven as one unit.
89
+ """
90
+
91
+ def __init__(self, handle: "ActorHandle | list[ActorHandle]") -> None:
92
+ self.handles = handle if isinstance(handle, list) else [handle]
93
+
94
+ def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None:
95
+ import ray
96
+
97
+ request = WeightTransferInitRequest(init_info=init_info)
98
+ ray.get([h.init_weight_transfer_engine.remote(request) for h in self.handles])
99
+
100
+ def start_weight_update(self) -> None:
101
+ import ray
102
+
103
+ ray.get([h.start_weight_update.remote() for h in self.handles])
104
+
105
+ def update_weights(self, update_info: dict[str, Any]) -> None:
106
+ import ray
107
+
108
+ request = WeightTransferUpdateRequest(update_info=update_info)
109
+ ray.get([h.update_weights.remote(request) for h in self.handles])
110
+
111
+ def finish_weight_update(self) -> None:
112
+ import ray
113
+
114
+ ray.get([h.finish_weight_update.remote() for h in self.handles])
vllm/distributed/weight_transfer/factory.py
+97
-1
@@ -6,7 +6,10 @@ import importlib
6
from collections.abc import Callable
7
from typing import TYPE_CHECKING
8
9
-from vllm.distributed.weight_transfer.base import WeightTransferEngine
9
+from vllm.distributed.weight_transfer.base import (
10
+ TrainerWeightTransferEngine,
11
+ WeightTransferEngine,
12
+)
13
from vllm.logger import init_logger
14
15
if TYPE_CHECKING:
@@ -14,6 +17,11 @@ if TYPE_CHECKING:
17
18
from vllm.config import VllmConfig
19
from vllm.config.weight_transfer import WeightTransferConfig
20
+ from vllm.distributed.weight_transfer.base import (
21
+ VLLMWeightSyncClient,
22
+ WeightSource,
23
+ WeightTransferInitInfo,
24
+ )
25
26
logger = init_logger(__name__)
27
@@ -111,6 +119,94 @@ class WeightTransferEngineFactory:
119
return engine_cls(config, vllm_config, device, model)
120
121
122
+class WeightTransferTrainerFactory:
123
+ """Factory for creating trainer-side weight transfer engines.
124
+
125
+ Parallel to `WeightTransferEngineFactory`, with its own lazy-import
126
+ registry. The trainer-side and worker-side registries are kept separate:
127
+ they share backend names by convention, but the trainer process never
128
+ instantiates a worker engine and vice versa, so unifying them would only
129
+ couple the import graphs.
130
+ """
131
+
132
+ _registry: dict[str, Callable[[], type[TrainerWeightTransferEngine]]] = {}
133
+
134
+ @classmethod
135
+ def register_engine(
136
+ cls,
137
+ name: str,
138
+ module_path_or_cls: "str | type[TrainerWeightTransferEngine]",
139
+ class_name: str | None = None,
140
+ ) -> None:
141
+ """Register a trainer engine. Same conventions as
142
+ `WeightTransferEngineFactory.register_engine`."""
143
+ if name in cls._registry:
144
+ raise ValueError(
145
+ f"Weight transfer trainer engine '{name}' is already registered."
146
+ )
147
+
148
+ if isinstance(module_path_or_cls, str):
149
+ module_path = module_path_or_cls
150
+ if class_name is None:
151
+ raise ValueError(
152
+ "class_name is required when registering with module path"
153
+ )
154
+
155
+ def loader() -> type[TrainerWeightTransferEngine]:
156
+ module = importlib.import_module(module_path)
157
+ return getattr(module, class_name)
158
+
159
+ cls._registry[name] = loader
160
+ else:
161
+ engine_cls = module_path_or_cls
162
+ cls._registry[name] = lambda: engine_cls
163
+
164
+ @classmethod
165
+ def trainer_init(
166
+ cls,
167
+ backend: str,
168
+ config: "WeightTransferConfig",
169
+ init_info: "WeightTransferInitInfo",
170
+ *,
171
+ client: "VLLMWeightSyncClient",
172
+ source: "WeightSource",
173
+ ) -> TrainerWeightTransferEngine:
174
+ """Build and rendezvous a ready-to-send trainer engine.
175
+
176
+ Called on every trainer rank (multi-rank trainers construct on all
177
+ ranks; the sender is resolved inside the engine's ``trainer_init``).
178
+
179
+ Args:
180
+ backend: Backend name (must be registered).
181
+ config: Backend-specific weight transfer config.
182
+ init_info: Backend-specific trainer init info.
183
+ client: Inference-side control-plane client.
184
+ source: `WeightSource` of `(name, tensor)` pairs to send each round.
185
+
186
+ Raises:
187
+ ValueError: If the backend is not registered.
188
+ """
189
+ if backend not in cls._registry:
190
+ available = list(cls._registry.keys())
191
+ raise ValueError(
192
+ f"Invalid weight transfer backend: {backend}. "
193
+ f"Available trainer engines: {available}"
194
+ )
195
+ engine_cls = cls._registry[backend]()
196
+
197
+ logger.info(
198
+ "Creating weight transfer trainer engine: %s",
199
+ engine_cls.__name__,
200
+ )
201
+
202
+ return engine_cls.trainer_init(
203
+ config=config,
204
+ init_info=init_info,
205
+ client=client,
206
+ source=source,
207
+ )
208
+
209
+
210
# Register built-in weight transfer engines here.
211
# Registration should be centralized to ensure lazy loading -
212
# engine modules are only imported when actually used.
vllm/v1/engine/async_llm.py
+2
-19
@@ -1067,17 +1067,8 @@ class AsyncLLM(EngineClient):
1067
Args:
1068
request: Weight transfer initialization request with backend-specific info
1069
"""
1070
- from vllm.distributed.weight_transfer.base import (
1071
- WeightTransferInitRequest,
1072
- )
1073
-
1074
- if isinstance(request, WeightTransferInitRequest):
1075
- init_info_dict = request.init_info
1076
- else:
1077
- raise TypeError(f"Expected WeightTransferInitRequest, got {type(request)}")
1078
-
1070
await self.collective_rpc(
1080
- "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
1071
+ "init_weight_transfer_engine", kwargs={"init_info": request.init_info}
1072
)
1073
1074
async def start_weight_update(self) -> None:
@@ -1095,16 +1086,8 @@ class AsyncLLM(EngineClient):
1086
Args:
1087
request: Weight update request with backend-specific update info
1088
"""
1098
-
1099
- if isinstance(request, WeightTransferUpdateRequest):
1100
- update_info_dict = request.update_info
1101
- else:
1102
- raise TypeError(
1103
- f"Expected WeightTransferUpdateRequest, got {type(request)}"
1104
- )
1105
-
1089
await self.collective_rpc(
1107
- "update_weights", kwargs={"update_info": update_info_dict}
1090
+ "update_weights", kwargs={"update_info": request.update_info}
1091
)
1092
1093
async def finish_weight_update(self) -> None: