[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>

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

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 = {}

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

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: ...
 

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)

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()