From ee5ae2ae70e03b328c7a197681b7b8485617f4ed Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 26 Aug 2026 15:36:57 +0800 Subject: [PATCH] [RL] Add rank-local IPC weight updates (#52497) Signed-off-by: aoshen02 Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: Roger Wang --- docs/training/weight_transfer/ipc.md | 6 +++ tests/distributed/test_weight_transfer.py | 7 ++++ .../worker/test_gpu_worker_weight_transfer.py | 40 +++++++++++++++++++ vllm/distributed/weight_transfer/base.py | 5 ++- vllm/distributed/weight_transfer/clients.py | 15 +++++-- vllm/v1/worker/gpu_worker.py | 15 +++++-- 6 files changed, 80 insertions(+), 8 deletions(-) diff --git a/docs/training/weight_transfer/ipc.md b/docs/training/weight_transfer/ipc.md index 70edefd64d..b71568c3df 100644 --- a/docs/training/weight_transfer/ipc.md +++ b/docs/training/weight_transfer/ipc.md @@ -110,6 +110,12 @@ mode carries a per-chunk barrier across ranks: without it, a rank could overwrit its buffer while its colocated worker is still reading the current chunk. This is handled inside `send_weights()`. +## Rank-Local Updates + +For different parameter subsets per worker, `update_info` may be a list indexed +by worker rank. This form is IPC-only because collective backends require every +worker to participate. + ## Examples - [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 diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 44eafe0123..5f9c6ff531 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -1361,6 +1361,13 @@ class TestTrainerClients: {"gpu": ("args",)} ] + client.update_weights([{"ipc_handles": {"gpu": ("args",)}}, {"names": []}]) + sent = captured["json"]["update_info"] + assert sent[1] == {"names": []} + assert pickle.loads(base64.b64decode(sent[0]["ipc_handles_pickled"])) == { + "gpu": ("args",) + } + def test_http_client_passes_through_nccl_update_info(self, monkeypatch): """NCCL update_info has only JSON-native fields and passes unchanged.""" captured = {} diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index 23c0b4bf61..e576ebd031 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -118,6 +118,46 @@ def test_start_update_finish_delegates_to_engine(): assert worker.model_runner.reset_lora_calls == 1 +@pytest.mark.parametrize( + ("rank", "expected"), + [(1, {"names": ["rank-1"]}), (2, {"names": []})], +) +def test_rank_local_update_selects_worker_payload(rank, expected): + engine = _RecordingEngine() + worker = _make_worker(engine) + worker.rank = rank + Worker.start_weight_update(worker) + + Worker.update_weights( + worker, [{"names": ["rank-0"]}, {"names": ["rank-1"]}, {"names": []}] + ) + + assert engine.update_calls == [expected] + assert worker._weight_update_active is True + + +def test_rank_local_update_includes_data_parallel_rank(): + engine = _RecordingEngine() + worker = _make_worker(engine) + worker.rank = 0 + worker.vllm_config.parallel_config.data_parallel_size = 4 + worker.vllm_config.parallel_config.data_parallel_rank = 2 + Worker.start_weight_update(worker) + + Worker.update_weights( + worker, + [ + {"names": ["dp-0"]}, + {"names": ["dp-1"]}, + {"names": ["dp-2"]}, + {"names": ["dp-3"]}, + ], + ) + + assert engine.update_calls == [{"names": ["dp-2"]}] + assert worker._weight_update_active is True + + def test_finish_draft_session_keeps_lora_state(): engine = _RecordingEngine() engine.supports_draft_weight_update = True diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index ed997a4529..2328829fc0 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -27,6 +27,7 @@ from vllm.config.weight_transfer import WeightTransferConfig TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo") TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo") TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo") +WeightTransferUpdatePayload = dict[str, Any] | list[dict[str, Any]] # A trainer supplies its parameters as a `WeightSource` (defined below): a # re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()` @@ -324,7 +325,7 @@ class WeightTransferInitRequest: class WeightTransferUpdateRequest: """API-level weight update request.""" - update_info: dict[str, Any] = field(default_factory=dict) + update_info: WeightTransferUpdatePayload = field(default_factory=dict) class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]): @@ -545,7 +546,7 @@ class VLLMWeightSyncClient(Protocol): def start_weight_update(self) -> None: ... - def update_weights(self, update_info: dict[str, Any]) -> None: ... + def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: ... def finish_weight_update(self, weight_version: str | None = None) -> None: ... diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 12dd0c9eac..1bd2f98937 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any from vllm.distributed.weight_transfer.base import ( WeightTransferInitRequest, + WeightTransferUpdatePayload, WeightTransferUpdateRequest, ) @@ -46,6 +47,14 @@ def _json_safe_update_info(update_info: dict[str, Any]) -> dict[str, Any]: return out +def _json_safe_update_payload( + update_info: WeightTransferUpdatePayload, +) -> WeightTransferUpdatePayload: + if isinstance(update_info, list): + return [_json_safe_update_info(info) for info in update_info] + return _json_safe_update_info(update_info) + + class HTTPVLLMWeightSyncClient: """Talks to a vLLM server over the RLHF HTTP routes. @@ -72,9 +81,9 @@ class HTTPVLLMWeightSyncClient: def start_weight_update(self) -> None: self._post("start_weight_update") - def update_weights(self, update_info: dict[str, Any]) -> None: + def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: self._post( - "update_weights", {"update_info": _json_safe_update_info(update_info)} + "update_weights", {"update_info": _json_safe_update_payload(update_info)} ) def finish_weight_update(self, weight_version: str | None = None) -> None: @@ -105,7 +114,7 @@ class RayVLLMWeightSyncClient: ray.get([h.start_weight_update.remote() for h in self.handles]) - def update_weights(self, update_info: dict[str, Any]) -> None: + def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: import ray request = WeightTransferUpdateRequest(update_info=update_info) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 50417ce6eb..7ac9a2c6c9 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1326,7 +1326,7 @@ class Worker(WorkerBase): self._weight_update_active = True self._weight_update_is_draft = is_draft - def update_weights(self, update_info: dict) -> None: + def update_weights(self, update_info: dict | list[dict]) -> None: """ Receive one weight update chunk from the trainer. @@ -1336,7 +1336,8 @@ class Worker(WorkerBase): / start_draft_weight_update call selected. Args: - update_info: Dictionary containing backend-specific update info + update_info: Backend-specific update info, or a list indexed by + global worker rank across data parallel replicas. """ self._check_weight_transfer_engine() assert self.weight_transfer_engine is not None @@ -1348,7 +1349,15 @@ class Worker(WorkerBase): with set_current_vllm_config(self.vllm_config): try: - self.weight_transfer_engine.update_weights(update_info) + if isinstance(update_info, list): + parallel_config = self.vllm_config.parallel_config + local_update_info = update_info[ + parallel_config.data_parallel_rank * parallel_config.world_size + + self.rank + ] + else: + local_update_info = update_info + self.weight_transfer_engine.update_weights(local_update_info) except BaseException: self._weight_update_active = False self.weight_transfer_engine.reset_weight_update_target()