diff --git a/slime/backends/megatron_utils/sglang.py b/slime/backends/megatron_utils/sglang.py index 97c82a31c..44cd66fe0 100644 --- a/slime/backends/megatron_utils/sglang.py +++ b/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", ] diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index cd1976f90..3e3737cdd 100644 --- a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/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() 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 918e1079f..80bfd849e 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 @@ -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,10 +378,11 @@ 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 @@ -396,44 +390,42 @@ def _send_hf_chunk_via_ipc(self, hf_named_tensors: Sequence[tuple[str, torch.Ten 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 diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index 92bf0cc6a..9a2fe83be 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -1,5 +1,6 @@ from __future__ import annotations +import base64 import ipaddress import logging import multiprocessing @@ -7,6 +8,7 @@ import time from urllib.parse import quote +import cloudpickle import requests from slime.ray.ray_actor import RayActor @@ -589,36 +591,6 @@ def _post_vllm_update_weights_http(self, update_info: dict) -> dict: except Exception: return {"ok": True, "raw": response.text} - def _run_vllm_weight_update(self, update_info: dict, *, is_checkpoint_format: bool = False): - """Backward-compatible alias for non-NCCL ``update_info`` shapes (e.g. tensor/IPC path).""" - del is_checkpoint_format - return self._post_vllm_update_weights_http(update_info) - - def update_weights(self, update_info: dict) -> dict: - """Public Ray-callable entry point used by IPCWeightTransferEngine (Ray mode). - - ``IPCWeightTransferEngine.trainer_send_weights`` calls - ``llm_handle.update_weights.remote(dict(update_info=update_info))``, - with ``update_info`` containing raw ``ipc_handles`` (Python callables from - ``reduce_tensor``). Since vime communicates with vLLM over HTTP those - callables cannot be JSON-serialised; convert them to ``ipc_handles_pickled`` - (base64-encoded pickle) which the vLLM server's ``IPCWeightTransferUpdateInfo`` - accepts when ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` is set. - """ - import base64 - - import cloudpickle - - inner = dict(update_info["update_info"]) # shallow copy — do not mutate caller - if inner.get("ipc_handles") is not None: - # ipc_handles contain local closures from monkey_patch_torch_reductions - # (_rebuild_cuda_tensor_modified) that standard pickle cannot serialize. - # cloudpickle handles local functions and closures correctly. - inner["ipc_handles_pickled"] = base64.b64encode(cloudpickle.dumps(inner.pop("ipc_handles"))).decode( - "utf-8" - ) - return self._post_vllm_update_weights_http(inner) - def health_generate(self, timeout: float = 5.0) -> bool: """Return True if ``GET /health`` succeeds (SGLang uses ``GET /health_generate`` for the same role).""" if self.node_rank != 0: @@ -629,37 +601,31 @@ def health_generate(self, timeout: float = 5.0) -> bool: def update_weights_from_tensor( self, - serialized_named_tensors: list[str], - load_format: str | None = None, - flush_cache: bool = False, + *, + names: list[str], + dtype_names: list[str], + shapes: list[list[int]], + ipc_handles: list[dict] | None = None, weight_version: str | None = None, - ): - """Post tensor metadata via ``/update_weights`` (vLLM RLHF native protocol). - - Contrasts with SGLang, which posts to ``update_weights_from_tensor`` with a different payload shape. - - If the POST fails this raises — there is intentionally no "fallback to reload" - 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 ``UpdateWeightFromTensor`` (vllm-native IPC) is - ported — see PR #12 review. + flush_cache: bool = False, + ) -> dict | None: + """POST ``IPCWeightTransferUpdateInfo`` (names / dtype_names / shapes / + ipc_handles) to ``/update_weights``; record ``weight_version`` only on + success. ``ipc_handles`` are base64-cloudpickle'd (rebuild_fn closures). """ - del load_format if self.node_rank != 0: - return + return None - if weight_version is not None: - self._weight_version = str(weight_version) + payload: dict = {"names": names, "dtype_names": dtype_names, "shapes": shapes} + if ipc_handles is not None: + payload["ipc_handles_pickled"] = base64.b64encode(cloudpickle.dumps(ipc_handles)).decode("utf-8") if flush_cache: self.flush_cache() - update_info = { - "serialized_named_tensors": serialized_named_tensors, - "format": "serialized_named_tensors", - "weight_version": self._weight_version, - } - return self._run_vllm_weight_update(update_info, is_checkpoint_format=False) + response = self._post_vllm_update_weights_http(payload) + if weight_version is not None: + self._weight_version = str(weight_version) + return response def flush_cache(self): """Clear prefix cache via ``POST /reset_prefix_cache`` (SGLang uses ``GET /flush_cache``).""" @@ -714,25 +680,25 @@ def shutdown(self): pass self.process = None - def get_weight_version(self): - """ - Prefer ``_weight_version`` if weight sync already set it; else try ``GET /v1/models`` for a stable id string. + def get_weight_version(self) -> str | None: + """Return the version recorded by the last successful weight transfer. - SGLang exposes ``GET /get_weight_version``; vLLM has no name-equivalent route, so semantics differ from that endpoint. + Raises ``RuntimeError`` if no weight transfer has recorded a version + yet — we don't fall back to a ``/v1/models`` lookup, which would + return the model path string and never match the trainer's integer + counter (i.e. produce a misleading "mismatch" downstream). + Worker ranks (``node_rank != 0``) short-circuit per the class idiom. """ if self.node_rank != 0: - return - if self._weight_version is not None: - return self._weight_version - try: - r = requests.get(f"{self._http_base()}/v1/models", timeout=10) - r.raise_for_status() - data = r.json().get("data") or [] - if data and isinstance(data[0], dict) and "id" in data[0]: - return str(data[0]["id"]) - except requests.RequestException as e: - logger.info("get_weight_version: /v1/models failed (%s)", e) - return None + return None + if self._weight_version is None: + raise RuntimeError( + "VLLMEngine.get_weight_version called before any successful " + "weight transfer recorded a version (update_weights_from_tensor " + "/ update_weights_from_distributed never reached their " + "post-POST version write)." + ) + return self._weight_version def release_memory_occupation(self, level: int = 1): """``POST /sleep?level={level}`` when sleep mode is enabled (SGLang: ``POST /release_memory_occupation``). @@ -812,29 +778,15 @@ def start_weight_update(self, is_checkpoint_format: bool = False) -> dict: except Exception: return {"ok": True, "raw": response.text} - def finish_weight_update(self, weight_version: str | None = None) -> dict: + def finish_weight_update(self) -> dict: """``POST /finish_weight_update`` — signals vLLM to exit IPC weight-update mode. - ``weight_version`` records the version the trainer just transferred via IPC. - The IPC path bypasses ``update_weights_from_tensor`` (which is the normal - place ``_weight_version`` gets recorded for the distributed/NCCL path), so - callers must thread the new version through here for ``get_weight_version`` - to report it back. Otherwise ci_test's engine-vs-updater version check - (slime/backends/megatron_utils/actor.py) fails the first time IPC weight - sync runs — ``_weight_version`` stays ``None`` and ``get_weight_version`` - falls back to ``GET /v1/models``, which returns the model path string - (e.g. ``/root/models/Qwen2.5-0.5B-Instruct``), never matching the - updater's integer version (``"1"``, ``"2"``, …). + Purely a state-machine bookend now; ``_weight_version`` is recorded by + ``update_weights_from_tensor`` (the IPC data-carrying RPC), matching slime's + single-RPC version-with-data semantics. """ response = self._post_json("finish_weight_update", {}, timeout=self._weight_transfer_http_timeout()) response.raise_for_status() - # Record the new version only after the POST succeeded — if the engine - # never actually exited weight-update mode, ``_weight_version`` must not - # advance, otherwise a retry would skip the resync. (Defensive: per the - # current call sites, any exception above propagates out of - # ``update_weights`` and the ci_test check below it would not run.) - if weight_version is not None: - self._weight_version = str(weight_version) try: return response.json() except Exception: 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 index ec0ca4049..21abd6196 100644 --- 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 @@ -25,7 +25,9 @@ def _install_stubs(): megatron_core = types.ModuleType("megatron.core") megatron_core.mpu = mpu_stub - sys.modules.setdefault("megatron", types.ModuleType("megatron")) + megatron_mod = types.ModuleType("megatron") + megatron_mod.core = megatron_core + sys.modules.setdefault("megatron", megatron_mod) sys.modules.setdefault("megatron.core", megatron_core) ray_mod = types.ModuleType("ray") @@ -110,6 +112,7 @@ class RecordingVLLMEngine: 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) + update_weights_from_tensor: 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) @@ -141,6 +144,7 @@ def _make_instance(upw_vllm, args=None): obj._ipc_engine_coordinator = False obj._ipc_engine_slot_start = None obj._ipc_engine_slot_end = None + obj._ipc_slot_group = None obj._distributed_engines = [] obj._model_update_groups = None obj._is_distributed_src_rank = False @@ -174,8 +178,7 @@ def counting_barrier(*args, **kwargs): 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() + obj.update_weights() return barrier_calls["n"] @@ -205,7 +208,11 @@ def test_colocated_lifecycle_uses_vllm_sleep_and_weight_transfer_apis(upw_vllm): @pytest.mark.unit -def test_trainer_send_weights_uses_single_llm_handle_per_rank(upw_vllm): +def test_send_via_ipc_dispatches_update_weights_from_tensor_with_version(upw_vllm): + """slot_size=1: every HF chunk fires + ``engine.update_weights_from_tensor.remote(**fields, weight_version=...)``. + Mirrors slime's IPC RPC contract — same name, parameterized fields, + version travels with data (no piggyback onto ``finish_weight_update``).""" obj = _make_instance(upw_vllm) engine = RecordingVLLMEngine() obj._colocated_engines = [engine] @@ -214,18 +221,79 @@ def test_trainer_send_weights_uses_single_llm_handle_per_rank(upw_vllm): obj._ipc_engine_slot_start = 0 obj._ipc_engine_slot_end = 1 - captured: list[dict] = [] + dummy_info = {"names": ["w"], "dtype_names": ["bfloat16"], "shapes": [[2, 2]], "ipc_handles": [{"u": ("f", ())}]} + with patch( + f"{MODULE_PATH}._build_ipc_update_info_from_named_tensors", + return_value=(dummy_info, []), + ): + _run_update(obj, chunks=_chunks(2)) + + # 2 HF chunks → 2 IPC RPCs + assert len(engine.update_weights_from_tensor.calls) == 2 + kwargs = engine.update_weights_from_tensor.calls[0].kwargs + # fields are passed as explicit kwargs (** expanded from local_info) + assert kwargs["names"] == dummy_info["names"] + assert kwargs["dtype_names"] == dummy_info["dtype_names"] + assert kwargs["shapes"] == dummy_info["shapes"] + assert kwargs["ipc_handles"] is dummy_info["ipc_handles"] + # weight_version is the trainer's post-increment version (0 + 1 = 1) as a str + assert kwargs["weight_version"] == "1" + # finish_weight_update is a stateless bookend now — no kwargs + assert len(engine.finish_weight_update.calls) == 1 + assert engine.finish_weight_update.calls[0].kwargs == {} + + +@pytest.mark.unit +def test_send_via_ipc_dispatches_update_weights_from_tensor_coordinator_multi_gpu(upw_vllm): + """slot_size > 1: coordinator gathers payloads from all slot ranks, merges them, + and fires a single engine.update_weights_from_tensor.remote() RPC per chunk.""" + obj = _make_instance(upw_vllm) + engine = RecordingVLLMEngine() + obj._colocated_engines = [engine] + obj._ipc_engine = engine + obj._ipc_engine_coordinator = True + obj._ipc_engine_slot_start = 0 + obj._ipc_engine_slot_end = 2 + + dummy_info_0 = { + "names": ["w"], + "dtype_names": ["bfloat16"], + "shapes": [[2, 2]], + "ipc_handles": [{"uuid-gpu0": ("f", ())}], + } + dummy_info_1 = { + "names": ["w"], + "dtype_names": ["bfloat16"], + "shapes": [[2, 2]], + "ipc_handles": [{"uuid-gpu1": ("f", ())}], + } - def fake_args(**kw): - captured.append(kw) - return kw + def fake_all_gather_object(gathered_payloads, payload, group=None): + gathered_payloads[0] = "payload0" + gathered_payloads[1] = "payload1" - ipc_engine = MagicMock() - _run_update(obj, chunks=_chunks(2), ipc_engine_cls=ipc_engine, ipc_args_cls=fake_args) + with patch("torch.distributed.get_rank", return_value=0), patch( + "megatron.core.mpu.get_tensor_model_parallel_rank", return_value=0 + ), patch( + f"{MODULE_PATH}._build_ipc_update_info_from_named_tensors", + return_value=(dummy_info_0, []), + ), patch( + f"{MODULE_PATH}._serialize_ipc_update_info", return_value="payload0" + ), patch( + f"{MODULE_PATH}._deserialize_ipc_update_info", side_effect=[dummy_info_0, dummy_info_1] * 2 + ), patch( + "torch.distributed.all_gather_object", side_effect=fake_all_gather_object + ): + _run_update(obj, chunks=_chunks(2)) - assert ipc_engine.trainer_send_weights.call_count == 2 - assert captured[0]["mode"] == "ray" - assert captured[0]["llm_handle"] is engine + assert len(engine.update_weights_from_tensor.calls) == 2 + kwargs = engine.update_weights_from_tensor.calls[0].kwargs + assert kwargs["names"] == dummy_info_0["names"] + assert kwargs["dtype_names"] == dummy_info_0["dtype_names"] + assert kwargs["shapes"] == dummy_info_0["shapes"] + assert len(kwargs["ipc_handles"]) == 1 + assert set(kwargs["ipc_handles"][0].keys()) == {"uuid-gpu0", "uuid-gpu1"} + assert kwargs["weight_version"] == "1" @pytest.mark.unit @@ -262,7 +330,7 @@ def test_connect_marks_one_coordinator_per_engine_gpu_slot(upw_vllm): ) with patch("torch.distributed.get_rank", return_value=rank), patch( "megatron.core.mpu.get_tensor_model_parallel_rank", return_value=tp_rank - ): + ), patch("torch.distributed.new_group", return_value="slot_group"): obj.connect_rollout_engines( engines, rollout_engine_lock=MagicMock(), diff --git a/tests/unit/backends/vllm_utils/test_vllm_engine.py b/tests/unit/backends/vllm_utils/test_vllm_engine.py index c699834ea..2de45b961 100644 --- a/tests/unit/backends/vllm_utils/test_vllm_engine.py +++ b/tests/unit/backends/vllm_utils/test_vllm_engine.py @@ -103,16 +103,71 @@ def fake_post(endpoint: str, payload: dict, timeout: float): @pytest.mark.unit -def test_finish_weight_update_records_weight_version_after_success(vllm_engine, monkeypatch): - def fake_post(endpoint: str, payload: dict, timeout: float): - assert endpoint == "finish_weight_update" - return _MockResponse(json_data={"done": True}) +def test_update_weights_from_tensor_posts_ipc_payload_and_records_version(vllm_engine, monkeypatch): + posted: list[dict] = [] + monkeypatch.setattr( + vllm_engine, + "_post_vllm_update_weights_http", + lambda payload: (posted.append(payload), {"ok": True})[1], + ) + assert vllm_engine._weight_version is None + + vllm_engine.update_weights_from_tensor( + names=["layer.0.weight"], + dtype_names=["float32"], + shapes=[[2, 2]], + ipc_handles=[{"uuid-gpu0": ("rebuild_fn", (1, 2, 3))}], + weight_version="42", + ) - monkeypatch.setattr(vllm_engine, "_post_json", fake_post) + assert len(posted) == 1 + sent = posted[0] + # ipc_handles got cloudpickle'd into ipc_handles_pickled + assert "ipc_handles" not in sent + assert isinstance(sent["ipc_handles_pickled"], str) + assert sent["names"] == ["layer.0.weight"] + assert sent["shapes"] == [[2, 2]] + # version recorded after POST success + assert vllm_engine._weight_version == "42" + + +@pytest.mark.unit +def test_update_weights_from_tensor_does_not_advance_version_on_failure(vllm_engine, monkeypatch): + """POST failure must not advance _weight_version (else a retry would skip the resync).""" + + def fake_post_vllm_fail(payload: dict) -> dict: + raise RuntimeError("simulated POST failure") - vllm_engine.finish_weight_update(weight_version="3") + monkeypatch.setattr(vllm_engine, "_post_vllm_update_weights_http", fake_post_vllm_fail) - assert vllm_engine.get_weight_version() == "3" + vllm_engine._weight_version = "old" + with pytest.raises(RuntimeError, match="simulated POST failure"): + vllm_engine.update_weights_from_tensor( + names=[], dtype_names=[], shapes=[], ipc_handles=[], weight_version="new" + ) + assert vllm_engine._weight_version == "old" + + +@pytest.mark.unit +def test_get_weight_version_returns_recorded_version(vllm_engine): + vllm_engine._weight_version = "7" + assert vllm_engine.get_weight_version() == "7" + + +@pytest.mark.unit +def test_get_weight_version_raises_when_unset(vllm_engine): + """Unrecorded version is a hard error — no silent /v1/models fallback.""" + assert vllm_engine._weight_version is None + with pytest.raises(RuntimeError, match="before any successful weight transfer"): + vllm_engine.get_weight_version() + + +@pytest.mark.unit +def test_get_weight_version_worker_rank_returns_none_without_raise(vllm_engine): + """Worker ranks short-circuit (matches the class-wide idiom).""" + vllm_engine.node_rank = 1 + vllm_engine._weight_version = None + assert vllm_engine.get_weight_version() is None @pytest.mark.unit