From 9253015f6699211bfc54ee7d8e6fc1c7f98c3b97 Mon Sep 17 00:00:00 2001 From: kaiyuan xie Date: Fri, 22 May 2026 16:03:58 +0800 Subject: [PATCH] fix: vLLM 0.21 colocate IPC weight sync (llm_handle, wake tags, start/finish sync) --- requirements.txt | 1 + .../update_weight_from_tensor.py | 180 +++++++++----- slime/backends/vllm_utils/vllm_engine.py | 14 +- .../vllm_utils/vllm_worker_extension.py | 51 ---- .../test_update_weight_from_tensor.py | 225 ++++++++++++++++++ 5 files changed, 353 insertions(+), 118 deletions(-) delete mode 100644 slime/backends/vllm_utils/vllm_worker_extension.py create mode 100644 tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py diff --git a/requirements.txt b/requirements.txt index 7e400e7ee..2cdaa711c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ accelerate blobfile +cloudpickle datasets httpx[http2] mcp[cli] diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 3897d6a1b..835407bf2 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/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,6 +46,14 @@ 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). @@ -30,8 +61,7 @@ class UpdateWeightFromTensor: 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,8 +108,8 @@ 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 @@ -86,6 +117,8 @@ def __init__( # 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()) \ No newline at end of file + 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) diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index f28ed2f93..69627a73d 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/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 diff --git a/slime/backends/vllm_utils/vllm_worker_extension.py b/slime/backends/vllm_utils/vllm_worker_extension.py deleted file mode 100644 index b0cc37f4d..000000000 --- a/slime/backends/vllm_utils/vllm_worker_extension.py +++ /dev/null @@ -1,51 +0,0 @@ -"""vLLM worker extension for vime colocated (IPC) mode. - -Passed to ``vllm serve`` via ``--worker-extension-cls`` so that the IPC -engine patch is applied inside every vLLM worker process automatically, -without requiring explicit patching from the trainer side. -""" - -from __future__ import annotations - - -class _VLLMHijack: - """Applies monkey-patches to vLLM internals required for vime colocated IPC weight sync.""" - - @staticmethod - def hijack() -> None: - """Patch ``IPCWeightTransferEngine.receive_weights`` to call - ``monkey_patch_torch_reductions`` before deserialising IPC handles. - - Idempotent – safe to call multiple times. - """ - from vllm.distributed.weight_transfer.ipc_engine import IPCWeightTransferEngine - - if getattr(IPCWeightTransferEngine, "_slime_receive_patched", False): - return - - from slime.backends.megatron_utils.update_weight.torch_patch import monkey_patch_torch_reductions - - _orig = IPCWeightTransferEngine.receive_weights - - def _slime_receive_weights(self, update_info, load_weights, _orig=_orig): - 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 for vime colocated (IPC) weight-sync mode. - - vLLM instantiates this class inside each worker process when - ``--worker-extension-cls`` is supplied. ``__new__`` is the earliest - reliable hook to apply process-wide patches before any weight transfer - occurs. - """ - - def __new__(cls, **kwargs): - # Apply the IPC engine patch so every worker process handles - # CUDA IPC handle deserialisation correctly. - _VLLMHijack.hijack() - return super().__new__(cls) diff --git a/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py b/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py new file mode 100644 index 000000000..481882e97 --- /dev/null +++ b/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py @@ -0,0 +1,225 @@ +"""Unit tests for colocated vLLM IPC weight sync (UpdateWeightFromTensor).""" + +from __future__ import annotations + +import importlib +import sys +import types +from argparse import Namespace +from dataclasses import dataclass, field +from unittest.mock import MagicMock, patch + +import pytest +import torch + +MODULE_PATH = "slime.backends.megatron_utils.update_weight.update_weight_from_tensor" + + +def _install_stubs(): + mpu_stub = MagicMock() + mpu_stub.get_data_parallel_rank.return_value = 0 + mpu_stub.get_tensor_model_parallel_rank.return_value = 0 + mpu_stub.get_pipeline_model_parallel_rank.return_value = 0 + + megatron_core = types.ModuleType("megatron.core") + megatron_core.mpu = mpu_stub + sys.modules.setdefault("megatron", types.ModuleType("megatron")) + sys.modules.setdefault("megatron.core", megatron_core) + + ray_mod = types.ModuleType("ray") + ray_mod.get = lambda refs: refs + ray_mod.actor = types.ModuleType("ray.actor") + ray_mod.actor.ActorHandle = object + sys.modules.setdefault("ray", ray_mod) + sys.modules.setdefault("ray.actor", ray_mod.actor) + + import torch.distributed as _dist + + dist_stub = MagicMock() + dist_stub.get_rank.return_value = 0 + dist_stub.barrier = MagicMock() + _dist.get_rank = dist_stub.get_rank + _dist.barrier = dist_stub.barrier + + slime_utils = types.ModuleType("slime.utils.distributed_utils") + slime_utils.get_gloo_group = MagicMock(return_value="gloo") + sys.modules.setdefault("slime.utils.distributed_utils", slime_utils) + + sglang_mod = types.ModuleType("slime.backends.megatron_utils.sglang") + sglang_mod.monkey_patch_torch_reductions = MagicMock() + sys.modules.setdefault("slime.backends.megatron_utils.sglang", sglang_mod) + + hf_iter_stub = MagicMock() + hf_iter_stub.get_hf_weight_chunks.return_value = iter([]) + + hf_base_mod = types.ModuleType("slime.backends.megatron_utils.update_weight.hf_weight_iterator_base") + hf_base_mod.HfWeightIteratorBase = MagicMock() + hf_base_mod.HfWeightIteratorBase.create.return_value = hf_iter_stub + + upw_dist_mod = types.ModuleType("slime.backends.megatron_utils.update_weight.update_weight_from_distributed") + upw_dist_mod.connect_rollout_engines_from_distributed = MagicMock(return_value="groups") + upw_dist_mod.disconnect_rollout_engines_from_distributed = MagicMock() + upw_dist_mod.post_process_weights = MagicMock() + upw_dist_mod.update_weights_from_distributed = MagicMock(return_value=[]) + + for key, mod in [ + ("slime.backends.megatron_utils.update_weight.hf_weight_iterator_base", hf_base_mod), + ("slime.backends.megatron_utils.update_weight.update_weight_from_distributed", upw_dist_mod), + ]: + sys.modules.setdefault(key, mod) + + return hf_iter_stub, upw_dist_mod + + +_HF_ITER_STUB, _UPW_DIST_MOD = _install_stubs() + + +@pytest.fixture(scope="module") +def upw_vllm(): + sys.modules.pop(MODULE_PATH, None) + return importlib.import_module(MODULE_PATH) + + +@dataclass +class _RemoteCall: + args: tuple + kwargs: dict + + +class RecordingRemoteMethod: + def __init__(self): + self.calls: list[_RemoteCall] = [] + + def remote(self, *args, **kwargs): + self.calls.append(_RemoteCall(args=args, kwargs=kwargs)) + return "ref" + + +@dataclass +class RecordingVLLMEngine: + release_memory_occupation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + resume_memory_occupation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + init_weight_transfer_engine: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + start_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + finish_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + pause_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + flush_cache: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + continue_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + + +def _default_args(**kwargs) -> Namespace: + base = dict( + actor_num_nodes=1, + actor_num_gpus_per_node=4, + rollout_num_gpus_per_engine=2, + megatron_to_hf_mode="raw", + update_weight_buffer_size=1 << 30, + ) + base.update(kwargs) + return Namespace(**base) + + +def _make_instance(upw_vllm, args=None): + obj = object.__new__(upw_vllm.UpdateWeightFromTensor) + obj.args = args or _default_args() + obj.model = [] + obj.weights_getter = lambda: {} + obj.model_name = "test" + obj.quantization_config = None + obj.weight_version = 0 + obj._hf_weight_iterator = _HF_ITER_STUB + obj._colocated_engines = [] + obj._ipc_engine = None + obj._distributed_engines = [] + obj._model_update_groups = None + obj._is_distributed_src_rank = False + obj._group_name = "slime" + obj._ipc_initialized = False + return obj + + +def _chunks(n=1): + return [[(f"p.{i}", torch.zeros(2, 2)) for i in range(2)] for _ in range(n)] + + +def _run_update(obj, *, chunks=None, ipc_engine_cls=None, ipc_args_cls=None) -> int: + chunks = chunks or _chunks(1) + obj._hf_weight_iterator = MagicMock() + obj._hf_weight_iterator.get_hf_weight_chunks.return_value = iter(chunks) + + ipc_engine_cls = ipc_engine_cls or MagicMock() + ipc_args_cls = ipc_args_cls or MagicMock(side_effect=lambda **kw: kw) + + ipc_mod = types.SimpleNamespace( + IPCWeightTransferEngine=ipc_engine_cls, + IPCTrainerSendWeightsArgs=ipc_args_cls, + ) + barrier_calls = {"n": 0} + + def counting_barrier(*args, **kwargs): + barrier_calls["n"] += 1 + + with patch.dict("sys.modules", {"vllm.distributed.weight_transfer.ipc_engine": ipc_mod}): + with patch("torch.distributed.get_rank", return_value=0), patch( + "torch.distributed.barrier", side_effect=counting_barrier + ): + with patch(f"{MODULE_PATH}._apply_monkey_patch_torch_reductions"): + obj.update_weights() + return barrier_calls["n"] + + +@pytest.mark.unit +def test_colocated_lifecycle_uses_vllm_sleep_and_weight_transfer_apis(upw_vllm): + obj = _make_instance(upw_vllm) + engine = RecordingVLLMEngine() + obj._colocated_engines = [engine] + obj._ipc_engine = engine + + barrier_count = _run_update(obj, chunks=_chunks(2)) + + assert len(engine.release_memory_occupation.calls) == 1 + assert engine.release_memory_occupation.calls[0].kwargs.get("level") == 0 + assert len(engine.init_weight_transfer_engine.calls) == 1 + assert engine.init_weight_transfer_engine.calls[0].args[0] == {"init_info": {}} + assert len(engine.start_weight_update.calls) == 1 + assert engine.start_weight_update.calls[0].kwargs.get("is_checkpoint_format") is True + assert len(engine.finish_weight_update.calls) == 1 + assert len(engine.resume_memory_occupation.calls) == 1 + # lifecycle barriers + one per HF chunk + assert barrier_count >= 2 + 2 + assert engine.resume_memory_occupation.calls[0].kwargs.get("tags") == ["weights", "kv_cache"] + + +@pytest.mark.unit +def test_trainer_send_weights_uses_single_llm_handle_per_rank(upw_vllm): + obj = _make_instance(upw_vllm) + engine = RecordingVLLMEngine() + obj._colocated_engines = [engine] + obj._ipc_engine = engine + + captured: list[dict] = [] + + def fake_args(**kw): + captured.append(kw) + return kw + + ipc_engine = MagicMock() + _run_update(obj, chunks=_chunks(2), ipc_engine_cls=ipc_engine, ipc_args_cls=fake_args) + + assert ipc_engine.trainer_send_weights.call_count == 2 + assert captured[0]["mode"] == "ray" + assert captured[0]["llm_handle"] is engine + + +@pytest.mark.unit +def test_ipc_init_runs_once(upw_vllm): + obj = _make_instance(upw_vllm) + engine = RecordingVLLMEngine() + obj._colocated_engines = [engine] + obj._ipc_engine = engine + + _run_update(obj) + _run_update(obj) + + assert len(engine.init_weight_transfer_engine.calls) == 1 + assert obj._ipc_initialized is True