fix: vLLM 0.21 colocate IPC weight sync

requirements.txt

@@ -1,5 +1,6 @@
 accelerate
 blobfile
+cloudpickle
 datasets
 httpx[http2]
 mcp[cli]

slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py

@@ -1,3 +1,26 @@
+"""
+Colocated vLLM weight sync (trainer + worker)
+=============================================
+
+Trainer: ``UpdateWeightFromTensor`` — Megatron → HF chunks → CUDA IPC (Ray).
+
+Worker: ``vLLMColocateWorkerExtension`` — passed to ``vllm serve`` via
+``--worker-extension-cls``; patches IPC receive before handle deserialisation.
+
+https://docs.vllm.ai/en/stable/examples/rl/rlhf_ipc/
+
+The flow for colocated engines:
+1. Megatron params → HF conversion (via HfWeightIteratorBase)
+2. All trainer ranks call ``IPCWeightTransferEngine.trainer_send_weights()``
+   with ``send_mode="ray"`` pointing at the colocated vLLM engine actor on the
+   same GPU slot.  Each rank creates a CUDA IPC handle for its GPU; the engine
+   collects all handles via ``_all_gather_and_merge_handles`` so every vLLM
+   worker can pick the handle belonging to its physical GPU UUID.
+
+For non-colocated overflow engines the existing NCCL distributed broadcast
+(``update_weights_from_distributed``) is used unchanged.
+"""
+
 from __future__ import annotations
 
 import logging
@@ -23,15 +46,22 @@
 
 logger = logging.getLogger(__name__)
 
+
+def _apply_monkey_patch_torch_reductions() -> None:
+    """CUDA IPC tensor rebuild uses GPU UUIDs; patch torch reductions before IPC."""
+    from slime.backends.megatron_utils.sglang import monkey_patch_torch_reductions
+
+    monkey_patch_torch_reductions()
+
+
 class UpdateWeightFromTensor:
     """
     Update colocated vLLM engines from tensors via CUDA IPC (Ray send mode).
 
     Colocated path:
         Megatron weights → HF conversion → CUDA IPC to vLLM engine actors via
         ``IPCWeightTransferEngine.trainer_send_weights(send_mode="ray")``.
-        All trainer ranks participate in the IPC handle all-gather; only rank 0
-        actually delivers the merged payload to the vLLM actors.
+        Each trainer rank sends to the colocated engine on its GPU slot.
 
     Distributed overflow path (optional):
         Falls back to NCCL distributed broadcast via
@@ -43,12 +73,13 @@ class UpdateWeightFromTensor:
         colocated:   release_memory_occupation(level=0) (rank 0)
         distributed: pause_generation / flush_cache      (rank 0)
         init_weight_transfer_engine                      (rank 0, colocated, first call only)
-        start_weight_update                              (rank 0, colocated)
+        start_weight_update                              (each rank, its colocated engine)
         [for each HF chunk]
-          trainer_send_weights                           (all ranks, colocated)
+          trainer_send_weights                           (rank with _ipc_engine)
           update_weights_from_distributed                (src rank, distributed)
-        finish_weight_update                             (rank 0, colocated)
-        colocated:   resume_memory_occupation(tags=["scheduling"]) (rank 0)
+          barrier                                        (all ranks)
+        finish_weight_update                             (each rank, its colocated engine)
+        colocated:   resume_memory_occupation(tags=["weights", "kv_cache"]) (rank 0)
         distributed: continue_generation                           (rank 0)
     """
 
@@ -77,15 +108,17 @@ def __init__(
 
         # Populated by connect_rollout_engines
         self._colocated_engines: list[ActorHandle] = []
-        self._colocated_engine_gpu_offsets: list[int] = []
-        self._colocated_engine_gpu_counts: list[int] = []
+        # vLLM 0.21 IPC (mode=ray): one Ray actor per GPU slot; this rank's engine.
+        self._ipc_engine: ActorHandle | None = None
         self._distributed_engines: list[ActorHandle] = []
         self._model_update_groups = None
         self._is_distributed_src_rank: bool = False
         self._group_name = "slime"
         # IPC weight transfer engine is initialized once per set of colocated
         # engines (not per update call).
         self._ipc_initialized: bool = False
+        # vLLM IPC handle payloads may use cloudpickle on the Ray/HTTP bridge.
+        os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
 
     # ------------------------------------------------------------------
     # connect / disconnect
@@ -104,10 +137,6 @@ def connect_rollout_engines(
         Colocated engines are those whose GPU range fits entirely within the
         trainer actor GPU range.  The remainder are treated as distributed and
         receive weights via NCCL broadcast.
-
-        The NCCL bridge for distributed engines is (re-)created whenever the
-        engine set changes, matching the behaviour of
-        ``UpdateWeightFromTensor.connect_rollout_engines``.
         """
         self.rollout_engine_lock = rollout_engine_lock
 
@@ -128,10 +157,20 @@ def connect_rollout_engines(
             colocate_engine_nums += 1
 
         self._colocated_engines = list(rollout_engines[:colocate_engine_nums])
-        self._colocated_engine_gpu_offsets = list(engine_gpu_offsets[:colocate_engine_nums])
-        self._colocated_engine_gpu_counts = list(engine_gpu_counts[:colocate_engine_nums])
         self._distributed_engines = list(rollout_engines[colocate_engine_nums:])
 
+        # Map this trainer rank to the colocated vLLM engine on the same GPU slot.
+        # vLLM 0.21 ``trainer_send_weights(mode="ray")`` expects a single ``llm_handle``,
+        # not a list (list fan-out is only in newer vLLM with ``send_mode="ray"``).
+        self._ipc_engine = None
+        colocate_gpu_offsets = engine_gpu_offsets[:colocate_engine_nums]
+        colocate_gpu_counts = engine_gpu_counts[:colocate_engine_nums]
+        for i, engine in enumerate(self._colocated_engines):
+            start = colocate_gpu_offsets[i]
+            end = start + colocate_gpu_counts[i]
+            if start <= dist.get_rank() < end:
+                self._ipc_engine = engine
+
         # Set up NCCL bridge for any overflow (non-colocated) engines.
         if self._distributed_engines:
             distributed_gpu_counts = engine_gpu_counts[colocate_engine_nums:]
@@ -164,18 +203,14 @@ def update_weights(self) -> None:
         """
         Transfer updated Megatron weights to all rollout engines.
 
-        Colocated engines receive weights via CUDA IPC (all trainer ranks
-        participate).  Distributed overflow engines receive weights via NCCL
-        broadcast (source rank only).
+        Colocated engines receive weights via CUDA IPC (per-rank engine RPC).
+        Distributed overflow engines receive weights via NCCL broadcast (source rank only).
         """
         self.weight_version += 1
         rank = dist.get_rank()
         all_engines = self._colocated_engines + self._distributed_engines
 
         # ── 1. Pause generation and flush KV cache (rank 0 only) ────────────
-        # vLLM colocated engines: release_memory_occupation(level=0) suspends generation
-        # and frees both KV cache and model weights (required for IPC tensor injection).
-        # Distributed (non-vLLM) engines keep the sglang-style pause+flush API.
         if rank == 0:
             if self._colocated_engines:
                 ray.get([engine.release_memory_occupation.remote(level=0) for engine in self._colocated_engines])
@@ -192,54 +227,38 @@ def update_weights(self) -> None:
 
         # ── 2. One-time IPC weight transfer engine init (rank 0 only) ───────
         if rank == 0 and self._colocated_engines and not self._ipc_initialized:
-            for engine in self._colocated_engines:
-                ray.get(engine.init_weight_transfer_engine.remote(dict(init_info=dict())))
-            self._ipc_initialized = True
-        dist.barrier(group=get_gloo_group())
-
-        # ── 3. Signal colocated vLLM engines to enter weight-update mode ─────
-        if rank == 0 and self._colocated_engines:
             ray.get(
-                [engine.start_weight_update.remote(is_checkpoint_format=True) for engine in self._colocated_engines]
+                [engine.init_weight_transfer_engine.remote({"init_info": {}}) for engine in self._colocated_engines]
             )
+            self._ipc_initialized = True
         dist.barrier(group=get_gloo_group())
 
-        # Required so vLLM can deserialize CUDA IPC handle payloads.
-        os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+        # ── 3. Enter weight-update mode (vLLM #39212: /start_weight_update) ───
+        if self._ipc_engine is not None:
+            ray.get(self._ipc_engine.start_weight_update.remote(is_checkpoint_format=True))
+        dist.barrier(group=get_gloo_group())
 
         from vllm.distributed.weight_transfer.ipc_engine import (  # noqa: PLC0415
             IPCTrainerSendWeightsArgs,
             IPCWeightTransferEngine,
         )
 
+        if self._colocated_engines:
+            _apply_monkey_patch_torch_reductions()
+
         # ── 4. Iterate HF weight chunks and send ─────────────────────────────
         megatron_local_weights = self.weights_getter()
         for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
-            # Colocated path: each trainer rank sends weights only to the engine
-            # that is colocated on the SAME physical GPU.  vLLM's
-            # trainer_send_weights (Ray mode) creates an IPC handle for the
-            # *current* GPU only — sending it to a different-GPU engine causes
-            # a UUID mismatch.  The matching engine is found by comparing
-            # torch.cuda.current_device() against the stored GPU offsets.
-            if self._colocated_engines:
-                current_device = torch.cuda.current_device()
-                for engine, offset, count in zip(
-                    self._colocated_engines,
-                    self._colocated_engine_gpu_offsets,
-                    self._colocated_engine_gpu_counts,
-                ):
-                    if offset <= current_device < offset + count:
-                        trainer_args = IPCTrainerSendWeightsArgs(
-                            mode="ray",
-                            llm_handle=engine,
-                        )
-                        IPCWeightTransferEngine.trainer_send_weights(
-                            iterator=iter(hf_named_tensors),
-                            trainer_args=trainer_args,
-                        )
-                        break
-
-            # Distributed overflow path (only the designated src rank).
+            if self._ipc_engine is not None:
+                trainer_args = IPCTrainerSendWeightsArgs(
+                    mode="ray",
+                    llm_handle=self._ipc_engine,
+                )
+                IPCWeightTransferEngine.trainer_send_weights(
+                    iterator=iter(hf_named_tensors),
+                    trainer_args=trainer_args,
+                )
+
             if self._distributed_engines and self._is_distributed_src_rank:
                 refs = update_weights_from_distributed(
                     self._group_name,
@@ -252,15 +271,14 @@ def update_weights(self) -> None:
                 if refs:
                     ray.get(refs)
 
+            dist.barrier(group=get_gloo_group())
+
         # ── 5. Signal colocated engines to exit weight-update mode ───────────
-        if rank == 0 and self._colocated_engines:
-            ray.get([engine.finish_weight_update.remote() for engine in self._colocated_engines])
+        if self._ipc_engine is not None:
+            ray.get(self._ipc_engine.finish_weight_update.remote())
         dist.barrier(group=get_gloo_group())
 
         # ── 6. Post-process quantization (if needed) and resume ───────────────
-        # vLLM colocated engines: resume_memory_occupation(tags=["scheduling"]) restores
-        # scheduling only (weights were just injected via IPC).
-        # Distributed engines use the sglang-style continue_generation.
         if rank == 0:
             if self.quantization_config and self.quantization_config.get("quant_method") in ["compressed-tensors"]:
                 post_process_weights(
@@ -269,7 +287,45 @@ def update_weights(self) -> None:
                     rollout_engines=all_engines,
                 )
             if self._colocated_engines:
-                ray.get([engine.resume_memory_occupation.remote(tags=["scheduling"]) for engine in self._colocated_engines])
+                ray.get(
+                    [
+                        engine.resume_memory_occupation.remote(tags=["weights", "kv_cache"])
+                        for engine in self._colocated_engines
+                    ]
+                )
             if self._distributed_engines:
                 ray.get([engine.continue_generation.remote() for engine in self._distributed_engines])
-        dist.barrier(group=get_gloo_group())
+        dist.barrier(group=get_gloo_group())
+
+
+# ---------------------------------------------------------------------------
+# vLLM worker extension (loaded by ``--worker-extension-cls`` in colocate mode)
+# ---------------------------------------------------------------------------
+
+
+class _VLLMHijack:
+    """Monkey-patch vLLM IPC receive so CUDA IPC handles deserialize on the correct GPU."""
+
+    @staticmethod
+    def hijack() -> None:
+        from vllm.distributed.weight_transfer.ipc_engine import IPCWeightTransferEngine
+
+        if getattr(IPCWeightTransferEngine, "_slime_receive_patched", False):
+            return
+
+        _orig = IPCWeightTransferEngine.receive_weights
+
+        def _slime_receive_weights(self, update_info, load_weights, _orig=_orig):
+            _apply_monkey_patch_torch_reductions()
+            _orig(self, update_info, load_weights)
+
+        IPCWeightTransferEngine.receive_weights = _slime_receive_weights
+        IPCWeightTransferEngine._slime_receive_patched = True  # type: ignore[attr-defined]
+
+
+class vLLMColocateWorkerExtension:
+    """vLLM ``--worker-extension-cls`` entry for colocated IPC weight sync."""
+
+    def __new__(cls, **kwargs):
+        _VLLMHijack.hijack()
+        return super().__new__(cls)

slime/backends/vllm_utils/vllm_engine.py

@@ -288,9 +288,13 @@ def launch_server_process(
     ]
     if getattr(args, "fp16", False):
         cmd += ["--dtype", "float16"]
-    # offload_rollout (vime top-level flag) implies sleep mode.
-    if getattr(args, "offload_rollout", False) and not getattr(args, "vllm_enable_sleep_mode", False):
+    # Colocated IPC weight sync releases model weights via POST /sleep?level=0.
+    # offload_rollout also needs sleep/wake for memory handoff.
+    if (getattr(args, "offload_rollout", False) or getattr(args, "colocate", False)) and not getattr(
+        args, "vllm_enable_sleep_mode", False
+    ):
         cmd += ["--enable-sleep-mode"]
+        args.vllm_enable_sleep_mode = True
     # rollout_max_context_len (vime top-level flag) maps to --max-model-len when set,
     # unless the user already passed --vllm-max-model-len explicitly.
     if args.rollout_max_context_len is not None and getattr(args, "vllm_max_model_len", None) is None:
@@ -330,7 +334,7 @@ def _user_overrode(dest: str) -> bool:
 
     # 2) weight_transfer_config: vllm default None disables /init_weight_transfer_engine,
     #    so vime's weight sync would fail.
-    #    - Colocated mode: use IPC backend. UpdateVLLMWeightFromTensor calls
+    #    - Colocated mode: use IPC backend. UpdateWeightFromTensor calls
     #      IPCWeightTransferEngine.trainer_send_weights and passes an empty init_info
     #      dict, which is the correct signature for the IPC backend.
     #    - Non-colocated mode: use NCCL backend. Weight sync goes through
@@ -353,7 +357,7 @@ def _user_overrode(dest: str) -> bool:
     if getattr(args, "colocate", False) and "--worker-extension-cls" not in cmd:
         cmd += [
             "--worker-extension-cls",
-            "slime.backends.vllm_utils.vllm_worker_extension.vLLMColocateWorkerExtension",
+            "slime.backends.megatron_utils.update_weight.update_weight_from_tensor.vLLMColocateWorkerExtension",
         ]
 
     # Auto-forward all other args.vllm_* that differ from their vllm-side default.
@@ -622,7 +626,7 @@ def update_weights_from_tensor(
         path. The previous fallback restarted vllm from ``self.model_path``, which is
         the original HF checkpoint (not the just-trained weights), so silently using
         it would let training continue with stale rollout weights. Failing fast keeps
-        the bug visible until ``UpdateVLLMWeightFromTensor`` (vllm-native IPC) is
+        the bug visible until ``UpdateWeightFromTensor`` (vllm-native IPC) is
         ported — see PR #12 review.
         """
         del load_format