[RL] Multi IPC handles example for rlhf colocated

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

examples/offline_inference/rlhf_colocate.py

@@ -28,6 +28,7 @@ Learn more about Ray placement groups:
 https://docs.ray.io/en/latest/placement-groups.html
 """
 
+import argparse
 import gc
 import os
 
@@ -78,7 +79,9 @@ class RayTrainingActor:
         # Ray sets CUDA_VISIBLE_DEVICES to the GPUs assigned to this actor.
         from transformers import AutoModelForCausalLM
 
-        self.model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            "/mnt/nvme3n1/models/qwen2.5_7B"
+        )
         self.model.to("cuda:0")
         # Zero out all the parameters.
         for name, p in self.model.named_parameters():
@@ -156,7 +159,105 @@ class RayTrainingActor:
         gc.collect()
         torch.cuda.empty_cache()
 
+    def update_weights_async(self, num_handles: int):
+        align_size = 256
 
+        def get_size(p: torch.Tensor) -> int:
+            return (p.nbytes + align_size - 1) // align_size * align_size
+
+        named_parameters: dict[str, torch.nn.Parameter] = dict(
+            self.model.named_parameters()
+        )
+        max_tensor_size = max(get_size(p) for p in named_parameters.values())
+        # use max_tensor_size * 2 as buffer size
+        buffer_capacity = max_tensor_size * 2
+        buffers = [
+            torch.empty(buffer_capacity, dtype=torch.uint8, device="cuda:0")
+            for _ in range(num_handles)
+        ]
+        handles = [reduce_tensor(b) for b in buffers]
+
+        # Establish ZMQ connection
+        s = self.zmq_context.socket(zmq.DEALER)
+        s.connect(self.zmq_handle)
+        s.send_pyobj(handles)
+        _ = s.recv_pyobj()
+
+        # === Partition tensors into buffers ===
+        offset = 0
+        buckets: list[tuple[list[dict], list[torch.Tensor]]] = []
+        named_tensors: list[dict] = []
+        real_tensors: list[torch.Tensor] = []
+
+        for name, p in named_parameters.items():
+            size = get_size(p)
+            if offset + size > buffer_capacity:
+                buckets.append((named_tensors, real_tensors))
+                named_tensors, real_tensors = [], []
+                offset = 0
+            # assume tensors are contiguous
+            named_tensors.append(
+                {"name": name, "dtype": p.dtype, "shape": p.shape, "offset": offset}
+            )
+            real_tensors.append(p)
+            offset += size
+        if named_tensors:
+            buckets.append((named_tensors, real_tensors))
+
+        poller = zmq.Poller()
+        poller.register(s, zmq.POLLOUT)
+        free_buffers = list(range(num_handles))
+        inflight = 0
+        idx = 0
+        total = len(buckets)
+        print(f"[Training] Total {total} buckets to send.")
+
+        # === Send loop ===
+        while idx < total or inflight > 0:
+            events = dict(poller.poll(timeout=50))
+            if (
+                s in events
+                and (events[s] & zmq.POLLOUT)
+                and idx < total
+                and free_buffers
+            ):
+                buf_id = free_buffers.pop(0)
+                meta, tensors = buckets[idx]
+                buffer = buffers[buf_id]
+                offset = 0
+                for item, t in zip(meta, tensors):
+                    size = get_size(t)
+                    buffer[offset : offset + t.nbytes].copy_(
+                        t.contiguous().view(torch.uint8).flatten()
+                    )
+                    offset += size
+                torch.cuda.synchronize()
+                s.send_pyobj((buf_id, meta))
+                inflight += 1
+                print(f"[Training] Sent bucket {idx} using buffer {buf_id}.")
+                idx += 1
+
+            # Receive buffer-free ACKs
+            try:
+                while s.getsockopt(zmq.EVENTS) & zmq.POLLIN:
+                    ack = s.recv_pyobj(flags=zmq.NOBLOCK)
+                    if isinstance(ack, int):
+                        free_buffers.append(ack)
+                        inflight -= 1
+                        print(f"[Training] Ack received: buffer {ack} now free.")
+            except zmq.Again:
+                pass
+
+        # Signal done
+        s.send_pyobj((None, None))
+        _ = s.recv_pyobj()
+        s.close()
+        del buffers
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def setup_train_cluster():
     # Ray manages four GPUs.
 
     os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
@@ -205,12 +306,12 @@ for i, bundle_indices in enumerate([[0, 1], [2, 3]]):
                 placement_group_capture_child_tasks=True,
             ),
         )(MyLLM).remote(
-        model="facebook/opt-125m",
+            model="/mnt/nvme3n1/models/qwen2.5_7B",
             enforce_eager=True,
             worker_extension_cls="rlhf_utils.ColocateWorkerExtension",
             tensor_parallel_size=2,
             distributed_executor_backend="ray",
-        gpu_memory_utilization=0.4,
+            gpu_memory_utilization=0.3,
             bundle_indices=bundle_indices,
         )
         inference_engines.append(llm)
@@ -236,16 +337,57 @@ for actor in training_actors:
         zmq_handles.update(ray.get(actor.get_zmq_handles.remote()))
 
     print(f"ZMQ handles: {zmq_handles}")
+    return training_actors, inference_engines, zmq_handles
 
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Update model weights across training and inference actors."
+    )
+    parser.add_argument(
+        "--num-ipc-handles",
+        type=int,
+        default=1,
+        help="Number of IPC handles. If 1, use synchronous update; \
+              if >1, use asynchronous update.",
+    )
+    args = parser.parse_args()
+
+    num_handles = args.num_ipc_handles
+    training_actors, inference_engines, zmq_handles = setup_train_cluster()
     print("Update the weights of the inference engines.")
+    if num_handles == 1:
+        # Synchronous update
         ray.get(
             [actor.update_weights.remote() for actor in training_actors]
             + [
-        llm.collective_rpc.remote("update_weights_from_ipc", args=(zmq_handles,))
+                llm.collective_rpc.remote(
+                    "update_weights_from_ipc", args=(zmq_handles,)
+                )
+                for llm in inference_engines
+            ]
+        )
+    else:
+        # Asynchronous update
+        ray.get(
+            [
+                actor.update_weights_async.remote(num_handles=num_handles)
+                for actor in training_actors
+            ]
+            + [
+                llm.collective_rpc.remote(
+                    "update_weights_from_ipc_async", args=(zmq_handles,)
+                )
                 for llm in inference_engines
             ]
         )
 
     print("Check if the weights are updated.")
     for llm in inference_engines:
-    assert ray.get(llm.collective_rpc.remote("check_weights_changed", args=tuple()))
+        assert ray.get(
+            llm.collective_rpc.remote("check_weights_changed", args=tuple())
+        ), "Weights were not updated properly!"
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/rlhf_utils.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import gc
+import pickle
+from ast import Dict, Tuple
 from collections.abc import Callable
+from enum import Enum
 from typing import TypedDict
 
 import torch
@@ -92,6 +95,15 @@ class FlattenedTensorMetadata(TypedDict):
     offset: int
 
 
+class PayloadType(Enum):
+    """Enumerates possible payload types in IPC protocol."""
+
+    HANDLES = "handles"
+    BUFFER_UPDATE = "buffer_update"
+    DONE = "done"
+    UNKNOWN = "unknown"
+
+
 class ColocateWorkerExtension:
     """
     The class for vLLM's worker to inherit from, in the colocate setting.
@@ -152,12 +164,90 @@ class ColocateWorkerExtension:
         gc.collect()
         torch.cuda.empty_cache()
 
+    def update_weights_from_ipc_async(self, zmq_handles: dict[str, str]):
+        assert self.device is not None
+        if not hasattr(self, "_zmq_ctx") or self._zmq_ctx is None:
+            self._zmq_ctx = zmq.Context()
+        socket = self._zmq_ctx.socket(zmq.ROUTER)
+        socket.bind(zmq_handles[self.report_device_id()])
+        poller = zmq.Poller()
+        poller.register(socket, zmq.POLLIN)
+
+        buffers: Dict[int, torch.Tensor] = {}
+        while True:
+            events = dict(poller.poll(timeout=100))
+            if socket in events and (events[socket] & zmq.POLLIN):
+                # Router identity
+                identity = socket.recv()
+
+                payload: (
+                    list[tuple[Callable, tuple]]
+                    | tuple[int, list[FlattenedTensorMetadata]]
+                    | None
+                ) = socket.recv_pyobj()
+
+                payload_type = self._identify_payload_type(payload)
+
+                # === HANDLE LIST OF SHARED MEMORY HANDLES ===
+                if payload_type == PayloadType.HANDLES:
+                    handles: list[tuple[Callable, tuple]] = payload
+                    for i, h in enumerate(handles):
+                        buffers[i] = rebuild_ipc(h, self.device.index)
+                    socket.send_multipart([identity, pickle.dumps("ACK_HANDLES")])
+                    continue
+
+                # === HANDLE BUFFERED MODEL UPDATES ===
+                if payload_type == PayloadType.BUFFER_UPDATE:
+                    buf_id, items = payload
+                    buffer = buffers.get(buf_id)
+                    if buffer is None:
+                        continue
+
+                    weights: list[Tuple[str, torch.Tensor]] = []
+                    for item in items:
+                        assert isinstance(item, dict)
+                        shape = torch.Size(item["shape"])
+                        dtype, offset = item["dtype"], item["offset"]
+                        size = dtype.itemsize * shape.numel()
+                        tensor = (
+                            buffer[offset : offset + size].view(dtype=dtype).view(shape)
+                        )
+                        weights.append((item["name"], tensor))
+
+                    self.model_runner.model.load_weights(weights=weights)
+                    torch.cuda.synchronize()
+                    # # --- Added verify step here ---
+                    # if self.verify_weights_enabled:
+                    #     self.verify_weights(weights)
+                    socket.send_multipart([identity, pickle.dumps(buf_id)])
+
+                # === DONE SIGNAL ===
+                elif payload_type == PayloadType.DONE:
+                    socket.send_multipart([identity, pickle.dumps("DONE")])
+                    break
+                else:
+                    continue
+
+        socket.close()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def report_device_id(self) -> str:
         from vllm.platforms import current_platform
 
         self.device_uuid = current_platform.get_device_uuid(self.device.index)
         return self.device_uuid
 
+    def _identify_payload_type(self, payload) -> PayloadType:
+        if isinstance(payload, list):
+            return PayloadType.HANDLES
+        elif isinstance(payload, tuple):
+            buf_id, _ = payload
+            if buf_id is None:
+                return PayloadType.DONE
+            return PayloadType.BUFFER_UPDATE
+        return PayloadType.UNKNOWN
+
     def check_weights_changed(self):
         """
         Check if the weights are updated to 0.