[RL] Add rank-local IPC weight updates (#52497)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Roger Wang <hey@rogerw.io>
aoshen02 committed
Aug 26, 2026 at 15:36 UTC
ee5ae2ae70e03b328c7a197681b7b8485617f4ed
6 files changed
+80
-8
docs/training/weight_transfer/ipc.md
+6
@@ -110,6 +110,12 @@ mode carries a per-chunk barrier across ranks: without it, a rank could overwrit
110
its buffer while its colocated worker is still reading the current chunk. This is
111
handled inside `send_weights()`.
112
113
+## Rank-Local Updates
114
+
115
+For different parameter subsets per worker, `update_info` may be a list indexed
116
+by worker rank. This form is IPC-only because collective backends require every
117
+worker to participate.
118
+
119
## Examples
120
121
- [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
tests/distributed/test_weight_transfer.py
+7
@@ -1361,6 +1361,13 @@ class TestTrainerClients:
1361
{"gpu": ("args",)}
1362
]
1363
1364
+ client.update_weights([{"ipc_handles": {"gpu": ("args",)}}, {"names": []}])
1365
+ sent = captured["json"]["update_info"]
1366
+ assert sent[1] == {"names": []}
1367
+ assert pickle.loads(base64.b64decode(sent[0]["ipc_handles_pickled"])) == {
1368
+ "gpu": ("args",)
1369
+ }
1370
+
1371
def test_http_client_passes_through_nccl_update_info(self, monkeypatch):
1372
"""NCCL update_info has only JSON-native fields and passes unchanged."""
1373
captured = {}
tests/v1/worker/test_gpu_worker_weight_transfer.py
+40
@@ -118,6 +118,46 @@ def test_start_update_finish_delegates_to_engine():
118
assert worker.model_runner.reset_lora_calls == 1
119
120
121
+@pytest.mark.parametrize(
122
+ ("rank", "expected"),
123
+ [(1, {"names": ["rank-1"]}), (2, {"names": []})],
124
+)
125
+def test_rank_local_update_selects_worker_payload(rank, expected):
126
+ engine = _RecordingEngine()
127
+ worker = _make_worker(engine)
128
+ worker.rank = rank
129
+ Worker.start_weight_update(worker)
130
+
131
+ Worker.update_weights(
132
+ worker, [{"names": ["rank-0"]}, {"names": ["rank-1"]}, {"names": []}]
133
+ )
134
+
135
+ assert engine.update_calls == [expected]
136
+ assert worker._weight_update_active is True
137
+
138
+
139
+def test_rank_local_update_includes_data_parallel_rank():
140
+ engine = _RecordingEngine()
141
+ worker = _make_worker(engine)
142
+ worker.rank = 0
143
+ worker.vllm_config.parallel_config.data_parallel_size = 4
144
+ worker.vllm_config.parallel_config.data_parallel_rank = 2
145
+ Worker.start_weight_update(worker)
146
+
147
+ Worker.update_weights(
148
+ worker,
149
+ [
150
+ {"names": ["dp-0"]},
151
+ {"names": ["dp-1"]},
152
+ {"names": ["dp-2"]},
153
+ {"names": ["dp-3"]},
154
+ ],
155
+ )
156
+
157
+ assert engine.update_calls == [{"names": ["dp-2"]}]
158
+ assert worker._weight_update_active is True
159
+
160
+
161
def test_finish_draft_session_keeps_lora_state():
162
engine = _RecordingEngine()
163
engine.supports_draft_weight_update = True
vllm/distributed/weight_transfer/base.py
+3
-2
@@ -27,6 +27,7 @@ from vllm.config.weight_transfer import WeightTransferConfig
27
TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo")
28
TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo")
29
TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo")
30
+WeightTransferUpdatePayload = dict[str, Any] | list[dict[str, Any]]
31
32
# A trainer supplies its parameters as a `WeightSource` (defined below): a
33
# re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()`
@@ -324,7 +325,7 @@ class WeightTransferInitRequest:
325
class WeightTransferUpdateRequest:
326
"""API-level weight update request."""
327
327
- update_info: dict[str, Any] = field(default_factory=dict)
328
+ update_info: WeightTransferUpdatePayload = field(default_factory=dict)
329
330
331
class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
@@ -545,7 +546,7 @@ class VLLMWeightSyncClient(Protocol):
546
547
def start_weight_update(self) -> None: ...
548
548
- def update_weights(self, update_info: dict[str, Any]) -> None: ...
549
+ def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: ...
550
551
def finish_weight_update(self, weight_version: str | None = None) -> None: ...
552
vllm/distributed/weight_transfer/clients.py
+12
-3
@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any
15
16
from vllm.distributed.weight_transfer.base import (
17
WeightTransferInitRequest,
18
+ WeightTransferUpdatePayload,
19
WeightTransferUpdateRequest,
20
)
21
@@ -46,6 +47,14 @@ def _json_safe_update_info(update_info: dict[str, Any]) -> dict[str, Any]:
47
return out
48
49
50
+def _json_safe_update_payload(
51
+ update_info: WeightTransferUpdatePayload,
52
+) -> WeightTransferUpdatePayload:
53
+ if isinstance(update_info, list):
54
+ return [_json_safe_update_info(info) for info in update_info]
55
+ return _json_safe_update_info(update_info)
56
+
57
+
58
class HTTPVLLMWeightSyncClient:
59
"""Talks to a vLLM server over the RLHF HTTP routes.
60
@@ -72,9 +81,9 @@ class HTTPVLLMWeightSyncClient:
81
def start_weight_update(self) -> None:
82
self._post("start_weight_update")
83
75
- def update_weights(self, update_info: dict[str, Any]) -> None:
84
+ def update_weights(self, update_info: WeightTransferUpdatePayload) -> None:
85
self._post(
77
- "update_weights", {"update_info": _json_safe_update_info(update_info)}
86
+ "update_weights", {"update_info": _json_safe_update_payload(update_info)}
87
)
88
89
def finish_weight_update(self, weight_version: str | None = None) -> None:
@@ -105,7 +114,7 @@ class RayVLLMWeightSyncClient:
114
115
ray.get([h.start_weight_update.remote() for h in self.handles])
116
108
- def update_weights(self, update_info: dict[str, Any]) -> None:
117
+ def update_weights(self, update_info: WeightTransferUpdatePayload) -> None:
118
import ray
119
120
request = WeightTransferUpdateRequest(update_info=update_info)
vllm/v1/worker/gpu_worker.py
+12
-3
@@ -1326,7 +1326,7 @@ class Worker(WorkerBase):
1326
self._weight_update_active = True
1327
self._weight_update_is_draft = is_draft
1328
1329
- def update_weights(self, update_info: dict) -> None:
1329
+ def update_weights(self, update_info: dict | list[dict]) -> None:
1330
"""
1331
Receive one weight update chunk from the trainer.
1332
@@ -1336,7 +1336,8 @@ class Worker(WorkerBase):
1336
/ start_draft_weight_update call selected.
1337
1338
Args:
1339
- update_info: Dictionary containing backend-specific update info
1339
+ update_info: Backend-specific update info, or a list indexed by
1340
+ global worker rank across data parallel replicas.
1341
"""
1342
self._check_weight_transfer_engine()
1343
assert self.weight_transfer_engine is not None
@@ -1348,7 +1349,15 @@ class Worker(WorkerBase):
1349
1350
with set_current_vllm_config(self.vllm_config):
1351
try:
1351
- self.weight_transfer_engine.update_weights(update_info)
1352
+ if isinstance(update_info, list):
1353
+ parallel_config = self.vllm_config.parallel_config
1354
+ local_update_info = update_info[
1355
+ parallel_config.data_parallel_rank * parallel_config.world_size
1356
+ + self.rank
1357
+ ]
1358
+ else:
1359
+ local_update_info = update_info
1360
+ self.weight_transfer_engine.update_weights(local_update_info)
1361
except BaseException:
1362
self._weight_update_active = False
1363
self.weight_transfer_engine.reset_weight_update_target()