refactor(weight-sync): align IPC RPC contract with slime (single-RPC version-with-data)

slime/backends/megatron_utils/sglang.py

@@ -7,12 +7,6 @@
     transform_scale_ue8m0 = None
     should_deepgemm_weight_requant_ue8m0 = None
 
-try:
-    from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
-except ImportError:
-    from sglang.srt.patch_torch import monkey_patch_torch_reductions
-
-
 from sglang.srt.utils import MultiprocessingSerializer
 
 
@@ -25,7 +19,6 @@
     "quant_weight_ue8m0",
     "transform_scale_ue8m0",
     "should_deepgemm_weight_requant_ue8m0",
-    "monkey_patch_torch_reductions",
     "MultiprocessingSerializer",
     "FlattenedTensorBucket",
 ]

slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py

@@ -11,7 +11,6 @@
 from slime.utils.types import ParamInfo
 
 from ..megatron_to_hf import convert_to_hf
-from ..sglang import monkey_patch_torch_reductions
 from .common import all_gather_params_async, named_params_and_buffers
 from .hf_weight_iterator_base import HfWeightIteratorBase
 
@@ -45,7 +44,6 @@ def _get_megatron_full_params(
     megatron_local_param_infos: Sequence[ParamInfo],
     megatron_local_weights,
 ) -> Sequence[torch.Tensor]:
-    monkey_patch_torch_reductions()
     pp_size = mpu.get_pipeline_model_parallel_world_size()
     ep_size = mpu.get_expert_model_parallel_world_size()
     rank = dist.get_rank()

slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py

@@ -8,17 +8,6 @@
 ``--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
@@ -44,13 +33,6 @@
 )
 
 
-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()
-
-
 def _current_gpu_uuid() -> str:
     device_index = torch.cuda.current_device()
     props = torch.cuda.get_device_properties(device_index)
@@ -62,9 +44,17 @@ def _build_ipc_update_info_from_named_tensors(
 ) -> tuple[dict[str, list], list[torch.Tensor]]:
     """Build vLLM IPC ``update_info`` payload from tensors on this rank's GPU.
 
-    Return the contiguous tensor refs with the payload. ``reduce_tensor`` only
-    exports CUDA IPC metadata, so the producer storage must stay alive until the
-    receiver opens the handle.
+    Each handle is keyed by the physical GPU UUID of the producing rank rather
+    than by a local device index. The coordinator gathers all ranks' dicts and
+    merges them; the receiver looks up its own UUID to pick the matching handle,
+    then vLLM unconditionally overwrites ``args[6]`` (device_index) with its own
+    local index before ``rebuild_cuda_tensor``. This UUID-keyed routing makes
+    the path correct under any ``CUDA_VISIBLE_DEVICES`` ordering without
+    relying on a torch reductions monkey-patch.
+
+    Return the contiguous tensor refs alongside the payload. ``reduce_tensor``
+    only exports CUDA IPC metadata, so the producer storage must stay alive
+    until the receiver opens the handle.
     """
     from torch.multiprocessing.reductions import reduce_tensor
 
@@ -133,18 +123,10 @@ def _merge_ipc_update_infos(infos: Sequence[dict[str, list]]) -> dict[str, list]
 
 
 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")``.
-        Each trainer rank sends to the colocated engine on its GPU slot.
-
-    Distributed overflow path (optional):
-        Falls back to NCCL distributed broadcast via
-        ``update_weights_from_distributed`` for engines whose GPUs lie outside
-        the actor GPU range.
+    """Update colocated vLLM engines via CUDA IPC, with NCCL fallback for
+    non-colocated overflow engines. See the module docstring for the
+    high-level design (why we dispatch via ``update_weights_from_tensor``
+    directly instead of vLLM's ``trainer_send_weights``).
 
     Engine lifecycle per ``update_weights`` call::
 
@@ -153,7 +135,7 @@ class UpdateWeightFromTensor:
         init_weight_transfer_engine                      (rank 0, colocated, first call only)
         start_weight_update                              (coordinator rank per engine only)
         [for each HF chunk]
-          trainer_send_weights                           (each rank mapped to _ipc_engine)
+          update_weights_from_tensor                     (per-rank or coordinator-merged)
           update_weights_from_distributed                (src rank, distributed)
           barrier                                        (all ranks)
         finish_weight_update                             (coordinator rank per engine only)
@@ -199,6 +181,8 @@ def __init__(
         # IPC weight transfer engine is initialized once per set of colocated
         # engines (not per update call).
         self._ipc_initialized: bool = False
+        # Per-engine-slot process group for IPC payload gather (created in connect_rollout_engines).
+        self._ipc_slot_group = None
         # vLLM IPC handle payloads may use cloudpickle on the Ray/HTTP bridge.
         os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
 
@@ -248,8 +232,23 @@ def connect_rollout_engines(
         self._ipc_engine_coordinator = False
         self._ipc_engine_slot_start = None
         self._ipc_engine_slot_end = None
+        # Build per-slot process groups so IPC payload gather covers all ranks in the
+        # engine's GPU slot — Megatron TP group does NOT cover the slot when Megatron
+        # TP != rollout-num-gpus-per-engine (e.g. Megatron TP=1 + rollout TP=2 in
+        # parallel-check). Every trainer rank must enter dist.new_group collectively.
+        self._ipc_slot_group = None
+        rank_for_slot = dist.get_rank()
         colocate_gpu_offsets = engine_gpu_offsets[:colocate_engine_nums]
         colocate_gpu_counts = engine_gpu_counts[:colocate_engine_nums]
+        # First pass: create per-slot process groups collectively (every rank must call new_group).
+        for i in range(colocate_engine_nums):
+            slot_start = colocate_gpu_offsets[i]
+            slot_end = slot_start + colocate_gpu_counts[i]
+            slot_ranks = list(range(slot_start, slot_end))
+            grp = dist.new_group(ranks=slot_ranks, backend="gloo")
+            if slot_start <= rank_for_slot < slot_end:
+                self._ipc_slot_group = grp
+        # Second pass: bind this rank to its engine + decide coordinator.
         for i, engine in enumerate(self._colocated_engines):
             start = colocate_gpu_offsets[i]
             end = start + colocate_gpu_counts[i]
@@ -258,8 +257,8 @@ def connect_rollout_engines(
                 self._ipc_engine = engine
                 self._ipc_engine_slot_start = start
                 self._ipc_engine_slot_end = end
-                # TP rank 0 within the engine GPU slot issues start/finish + merged IPC send.
-                if mpu.get_tensor_model_parallel_rank() == 0:
+                # Slot leader (lowest trainer rank in the engine GPU range) issues start/finish.
+                if rank == start:
                     self._ipc_engine_coordinator = True
 
         # Set up NCCL bridge for any overflow (non-colocated) engines.
@@ -329,9 +328,6 @@ def update_weights(self) -> None:
             ray.get(self._ipc_engine.start_weight_update.remote(is_checkpoint_format=True))
         dist.barrier(group=get_gloo_group())
 
-        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):
@@ -353,14 +349,11 @@ def update_weights(self) -> None:
             dist.barrier(group=get_gloo_group())
 
         # ── 5. Signal colocated engines to exit weight-update mode ───────────
-        # Thread the just-incremented weight_version through finish_weight_update
-        # so each colocated engine records it on ``self._weight_version``. The IPC
-        # path otherwise bypasses ``update_weights_from_tensor`` (the normal hook
-        # for setting it) and ci_test's engine-vs-updater check at
-        # slime/backends/megatron_utils/actor.py would mismatch (engine reports
-        # the model path from /v1/models; updater reports the integer version).
+        # State-machine bookend only; ``_weight_version`` is recorded inside
+        # ``update_weights_from_tensor`` (step 4) when the data RPC succeeds —
+        # matches slime's single-RPC version-with-data semantics.
         if self._ipc_engine_coordinator:
-            ray.get(self._ipc_engine.finish_weight_update.remote(weight_version=str(self.weight_version)))
+            ray.get(self._ipc_engine.finish_weight_update.remote())
         dist.barrier(group=get_gloo_group())
 
         # ── 6. Post-process quantization (if needed) and resume ───────────────
@@ -385,55 +378,54 @@ def update_weights(self) -> None:
     def _send_hf_chunk_via_ipc(self, hf_named_tensors: Sequence[tuple[str, torch.Tensor]]) -> None:
         """Send one HF chunk to the colocated vLLM engine via CUDA IPC (Ray → HTTP).
 
-        When ``rollout_num_gpus_per_engine > 1``, every trainer rank in the engine's GPU
-        slot builds an IPC handle on its GPU; the coordinator merges UUIDs and issues
-        a single ``update_weights`` RPC.  vLLM 0.21 ``trainer_send_weights`` alone only
-        ships the calling rank's UUID, which breaks vLLM TP workers on sibling GPUs.
+        ``slot_size == 1``: this rank ships its IPC payload directly.
+        ``slot_size > 1`` (vLLM TP): every rank in the slot builds its handle;
+        the coordinator gathers them, merges UUIDs, and issues one RPC for the
+        slot. Both paths dispatch the same RPC — ``update_weights_from_tensor`` —
+        with ``weight_version`` alongside the data (see module docstring).
         """
         assert self._ipc_engine is not None
         assert self._ipc_engine_slot_start is not None
         assert self._ipc_engine_slot_end is not None
 
         slot_size = self._ipc_engine_slot_end - self._ipc_engine_slot_start
         if slot_size <= 1:
-            from vllm.distributed.weight_transfer.ipc_engine import (  # noqa: PLC0415
-                IPCTrainerSendWeightsArgs,
-                IPCWeightTransferEngine,
-            )
-
-            trainer_args = IPCTrainerSendWeightsArgs(
-                mode="ray",
-                llm_handle=self._ipc_engine,
-            )
-            IPCWeightTransferEngine.trainer_send_weights(
-                iterator=iter(hf_named_tensors),
-                trainer_args=trainer_args,
+            local_info, weight_refs = _build_ipc_update_info_from_named_tensors(hf_named_tensors)
+            ray.get(
+                self._ipc_engine.update_weights_from_tensor.remote(
+                    **local_info,
+                    weight_version=str(self.weight_version),
+                )
             )
+            # Keep CUDA IPC producer tensors alive until ray.get() returns
+            # (the HTTP weight update completes inside the engine actor); then release.
+            del weight_refs
             return
 
         local_info, weight_refs = _build_ipc_update_info_from_named_tensors(hf_named_tensors)
         payload = _serialize_ipc_update_info(local_info)
 
-        tp_group = mpu.get_tensor_model_parallel_group()
-        tp_size = mpu.get_tensor_model_parallel_world_size()
-        tp_ranks = sorted(dist.get_process_group_ranks(tp_group))
+        slot_group = self._ipc_slot_group
+        slot_ranks = list(range(self._ipc_engine_slot_start, self._ipc_engine_slot_end))
 
-        # Use all_gather_object (monkey-patched for ReloadableProcessGroup). gather_object
-        # is not patched and fails after Megatron offload/reload with "Group is not registered".
-        gathered_payloads: list[str | None] = [None] * tp_size
-        dist.all_gather_object(gathered_payloads, payload, group=tp_group)
+        # Gather IPC payloads over the engine slot ranks (NOT Megatron TP group — see
+        # connect_rollout_engines for why). all_gather_object is monkey-patched for
+        # ReloadableProcessGroup; gather_object is not (fails after Megatron reload).
+        gathered_payloads: list[str | None] = [None] * slot_size
+        dist.all_gather_object(gathered_payloads, payload, group=slot_group)
         if self._ipc_engine_coordinator:
             if any(p is None for p in gathered_payloads):
-                raise RuntimeError(
-                    f"Missing IPC payloads on TP group {tp_ranks} (slot "
-                    f"[{self._ipc_engine_slot_start}, {self._ipc_engine_slot_end})); "
-                    f"got {gathered_payloads!r}"
-                )
+                raise RuntimeError(f"Missing IPC payloads on slot ranks {slot_ranks}; " f"got {gathered_payloads!r}")
             slot_infos = [_deserialize_ipc_update_info(p) for p in gathered_payloads]
             merged = _merge_ipc_update_infos(slot_infos)
-            ray.get(self._ipc_engine.update_weights.remote(dict(update_info=merged)))
+            ray.get(
+                self._ipc_engine.update_weights_from_tensor.remote(
+                    **merged,
+                    weight_version=str(self.weight_version),
+                )
+            )
 
-        dist.barrier(group=tp_group)
+        dist.barrier(group=slot_group)
         # Keep CUDA IPC producer tensors alive until every TP worker has opened
         # the handles and the coordinator's HTTP update has completed.
         del weight_refs
@@ -457,7 +449,6 @@ def hijack() -> None:
         _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