From f3072dbce0f58b8de6cd6622a6a1431813331c79 Mon Sep 17 00:00:00 2001 From: Justin Wood Date: Wed, 24 Jun 2026 05:09:39 +0000 Subject: [PATCH] [BugFix] Fix cumem_tag PP-broadcast wake race with symmetric gloo wake handshake (#45519) With --sleep-mode-backend=cumem_tag on a TP>1/PP>1 deployment, wake_up is dispatched to each worker independently and each worker re-maps its own VMM-backed regions at its own pace. The very next decode step issues a cross-rank torch.distributed.broadcast on pp.device_group (_pp_receive_prev_sampled_token_ids_to_input_batch). If a faster rank reaches that collective before a slower rank has finished re-mapping the regions backing the broadcast buffers, the collective issues against memory the peer MMU still treats as invalid -> CUDA_ERROR_ILLEGAL_ADDRESS, after which the NCCL comm is permanently corrupt: the engine deadlocks while /health keeps returning 200. Fix: gate Worker.wake_up on a cross-rank wake-success handshake after the local allocator wake and before returning to the caller, so no rank can reach a device-group collective until every rank has finished its local wake. The handshake is an all-reduce (ReduceOp.MIN) of a per-rank success flag, NOT a bare barrier: a bare barrier after the local wake would strand peers forever if one rank's allocator.wake_up() raised before reaching it, re-introducing the very full-fleet hang we are fixing. The all-reduce instead lets every rank learn that a peer failed and raise symmetrically -- loud, no hang, no rank silently proceeding into a device-group collective against a peer whose wake never completed. The handshake runs on get_world_group().cpu_group (gloo) deliberately so the synchronization itself never touches the not-yet-resynced NCCL device_group. A local torch.cuda.synchronize() is also added at the end of CuMemAllocator.wake_up so a rank drains its own re-map work before reporting success into the cross-rank all-reduce (a purely local guarantee; the cross-rank ordering lives entirely in Worker.wake_up). This is the correct, fully-Python fix. It supersedes the C++ allocator retry approach in #45565, which addresses a different symptom and (being compiled-code) cannot be carried as a Python-only patch. Tests (GPU-free, fail-pre/pass-post): tests/v1/worker/test_gpu_worker_wake_barrier.py -- handshake fires on multi-rank cumem configs; ordered after local wake, before return; skipped on single-rank / cumem-disabled; routed through cpu_group (gloo) only; plus the adversarial regressions: a failed local wake still participates in the handshake (no peer hang) and a failed peer makes a healthy rank raise symmetrically. tests/device_allocator/test_cumem_wake_synchronize.py -- wake_up calls torch.cuda.synchronize after the remap/restore work. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: terafin --- tests/device_allocator/__init__.py | 0 .../test_cumem_wake_synchronize.py | 144 +++++++ .../v1/worker/test_gpu_worker_wake_barrier.py | 379 ++++++++++++++++++ vllm/device_allocator/cumem.py | 28 ++ vllm/v1/worker/gpu_worker.py | 81 +++- 5 files changed, 631 insertions(+), 1 deletion(-) create mode 100644 tests/device_allocator/__init__.py create mode 100644 tests/device_allocator/test_cumem_wake_synchronize.py create mode 100644 tests/v1/worker/test_gpu_worker_wake_barrier.py diff --git a/tests/device_allocator/__init__.py b/tests/device_allocator/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/device_allocator/test_cumem_wake_synchronize.py b/tests/device_allocator/test_cumem_wake_synchronize.py new file mode 100644 index 000000000000..e8c881062960 --- /dev/null +++ b/tests/device_allocator/test_cumem_wake_synchronize.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test for the device-side wake synchronization in +``CuMemAllocator.wake_up`` (part of the fix for issue #45519). + +The fix adds a ``torch.cuda.synchronize()`` at the *end* of +``CuMemAllocator.wake_up`` — after the per-region ``create_and_map`` re-map +loop and the ``cudaMemcpy`` restores, and before the method returns to the +caller (``Worker.wake_up``). + +This synchronize is a purely LOCAL guarantee: it drains *this* rank's device +work (the freshly re-mapped VMM regions) so the rank does not report wake +success into the cross-rank handshake while its own remaps are still pending on +the device. It provides NO cross-rank ordering on its own — the cross-rank +guarantee for the #45519 PP-broadcast race lives in ``Worker.wake_up`` (the +gloo all-reduce of a wake-success flag). + +Two accuracy notes (the original rationale for this synchronize was wrong on +both, and a maintainer would bounce it): + * The ``cudaMemcpy`` restores use the *synchronous*, host-blocking + ``cudaMemcpy`` (see ``CudaRTLibrary``), NOT ``cudaMemcpyAsync`` — so the + copies have already completed w.r.t. the host by the time the loop exits; + they are not what the synchronize guards. + * The ``sleep()`` offload loop does NOT ``torch.cuda.synchronize`` per handle + between its D2H ``cudaMemcpy`` and the following ``unmap_and_release`` — the + only per-handle sync lives in ``_python_free_callback`` (a different path). + So this is not a "wake-side analogue" of a sleep-side per-handle sync. + +LIMITATION (documented, per the task): a *true* unit test of ``wake_up`` is not +feasible GPU-free — ``create_and_map`` and ``libcudart.cudaMemcpy`` drive real +CUDA driver / VMM calls, and ``torch.cuda.synchronize`` requires a CUDA context. +This is the closest mockable test: we construct a real ``CuMemAllocator`` (its +``__init__`` is pure-Python — only dict assignment), populate +``pointer_to_data`` with one real ``AllocationData`` carrying a CPU backup +tensor (so the ``cudaMemcpy`` restore branch is exercised), and mock the three +CUDA-touching surfaces (``create_and_map``, the module ``libcudart``, and +``torch.cuda.synchronize``). We then assert ``synchronize`` was called and that +it was called AFTER the remap/restore work (call-order), which is the property +the fix guarantees. +""" + +from unittest import mock + +import pytest + +from vllm.device_allocator import AllocationData +from vllm.device_allocator import cumem as cumem_mod +from vllm.device_allocator.cumem import CuMemAllocator + + +def _make_allocator_with_one_region(calls: list[str]): + """A real CuMemAllocator with one offloaded region queued for wake. + + `calls` records the load-bearing event order: each per-region remap + (`create_and_map`) / restore (`cudaMemcpy`) appends, and the final + `torch.cuda.synchronize` appends `"synchronize"`. + """ + allocator = CuMemAllocator.__new__(CuMemAllocator) + allocator.pointer_to_data = {} + allocator.current_tag = CuMemAllocator.default_tag + allocator.allocator_and_pools = {} + + # One region: handle = (device, size, ptr, vmm_handle). Give it a CPU + # backup tensor so the cudaMemcpy restore branch runs (the host-blocking + # restore the synchronize sequences the VMM remap against). + cpu_backup = mock.MagicMock() + cpu_backup.numel.return_value = 4 + cpu_backup.element_size.return_value = 2 + cpu_backup.data_ptr.return_value = 0xCAFE + ptr = 0x1000 + handle = (0, 8, ptr, 0xABCD) + allocator.pointer_to_data[ptr] = AllocationData( + handle=handle, tag=CuMemAllocator.default_tag, cpu_backup_tensor=cpu_backup + ) + return allocator + + +def _run_wake(allocator, calls): + fake_libcudart = mock.MagicMock() + fake_libcudart.cudaMemcpy.side_effect = lambda *a, **k: calls.append("cudaMemcpy") + + with mock.patch.object( + cumem_mod, "create_and_map", + side_effect=lambda *a, **k: calls.append("create_and_map"), + ), mock.patch.object(cumem_mod, "libcudart", fake_libcudart), mock.patch( + "torch.cuda.synchronize", + side_effect=lambda *a, **k: calls.append("synchronize"), + ) as sync: + allocator.wake_up() + return sync + + +def test_wake_up_calls_cuda_synchronize(): + """`CuMemAllocator.wake_up` must call `torch.cuda.synchronize()`.""" + calls: list[str] = [] + allocator = _make_allocator_with_one_region(calls) + sync = _run_wake(allocator, calls) + + sync.assert_called_once() + assert "synchronize" in calls + + +def test_wake_up_synchronize_runs_after_remap_and_restore(): + """The synchronize must come AFTER the remap (`create_and_map`) and the + `cudaMemcpy` restore, and is the last device op before return. + + A synchronize placed *inside* / *before* the loop — or omitted — would + leave this rank reporting wake success while its VMM remaps are still + pending on the device. We assert ordering explicitly. + """ + calls: list[str] = [] + allocator = _make_allocator_with_one_region(calls) + _run_wake(allocator, calls) + + assert calls == ["create_and_map", "cudaMemcpy", "synchronize"], ( + "expected remap -> cudaMemcpy restore -> device synchronize, got " + f"{calls}" + ) + # Synchronize is the final device action before wake_up returns. + assert calls[-1] == "synchronize" + + +def test_wake_up_synchronizes_even_with_no_backup_tensor(): + """Even a region with no CPU backup (re-map only, no cudaMemcpy) must + still be followed by the device synchronize before returning.""" + calls: list[str] = [] + allocator = CuMemAllocator.__new__(CuMemAllocator) + allocator.pointer_to_data = {} + allocator.current_tag = CuMemAllocator.default_tag + allocator.allocator_and_pools = {} + ptr = 0x2000 + allocator.pointer_to_data[ptr] = AllocationData( + handle=(0, 8, ptr, 0xBEEF), + tag=CuMemAllocator.default_tag, + cpu_backup_tensor=None, + ) + sync = _run_wake(allocator, calls) + + sync.assert_called_once() + assert calls == ["create_and_map", "synchronize"] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/v1/worker/test_gpu_worker_wake_barrier.py b/tests/v1/worker/test_gpu_worker_wake_barrier.py new file mode 100644 index 000000000000..b3bd8391b844 --- /dev/null +++ b/tests/v1/worker/test_gpu_worker_wake_barrier.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the cross-rank wake synchronization in Worker.wake_up. + +Context (issue #45519): with ``--sleep-mode-backend=cumem_tag`` on a TP>1/PP>1 +deployment, ``wake_up`` is dispatched to each worker independently. Each worker +re-maps its own VMM-backed regions at its own pace. The very next decode step +issues a cross-rank ``torch.distributed.broadcast`` on ``pp.device_group`` +(``_pp_receive_prev_sampled_token_ids_to_input_batch``). If a fast rank reaches +that collective before a slow rank has finished re-mapping the regions backing +the broadcast buffers, the collective issues against memory the peer MMU still +treats as invalid -> ``cudaErrorIllegalAddress`` and the NCCL comm is corrupted +permanently (engine deadlocks, /health lies 200). + +The fix gates ``Worker.wake_up`` on a cross-rank wake-success handshake over the +CPU (gloo) group *after* the local allocator wake and *before* returning to the +caller, so no rank can reach a device-group collective until every rank has +finished its local wake. Crucially the handshake is an all-reduce of a per-rank +success flag (``ReduceOp.MIN``), NOT a bare barrier: if one rank's local wake +raises, a bare barrier would strand its peers forever (re-introducing the very +full-fleet hang we are fixing). The all-reduce instead lets every rank learn +that a peer failed and raise *symmetrically* -- loud, no hang, no rank silently +proceeding into a device-group collective against a peer whose wake never +completed. The gloo group is used deliberately so the synchronization itself +never touches the not-yet-resynced NCCL device_group. + +These tests are GPU-free: they drive ``Worker.wake_up`` against mocked allocator +and distributed groups, and assert the handshake is (a) performed on multi-rank +cumem configs, (b) ordered after the allocator wake and before return, +(c) skipped on single-rank configs, (d) routed through the gloo cpu_group only, +and -- the adversarial-finding regression -- (e) does NOT hang and DOES raise +symmetrically when one simulated rank's local wake fails. +""" + +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from vllm.v1.worker.gpu_worker import Worker + + +def _make_worker(*, enable_cumem: bool = True) -> Worker: + """Build a Worker shell sufficient to exercise wake_up without a GPU. + + We bypass __init__ entirely and set only the attributes wake_up touches. + """ + worker = Worker.__new__(Worker) + worker._sleep_saved_buffers = {} + worker.model_runner = mock.MagicMock() + worker.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(enable_cumem_allocator=enable_cumem) + ) + return worker + + +def _patch_groups( + calls: list[str], + *, + tp_world_size: int, + pp_world_size: int, + allocator_raises: bool = False, + # Simulated *peer* outcome folded into the all-reduce result. The real + # all-reduce uses ReduceOp.MIN over int32 flags, so the post-reduce value is + # 0 if THIS rank failed OR any peer failed, else 1. `peer_ok=False` models + # "some other rank failed even though this rank's local wake succeeded". + peer_ok: bool = True, +): + """Patch the distributed group getters, allocator, and the gloo all-reduce + used by wake_up, recording the order of the load-bearing events into + ``calls``. + + The all-reduce is patched to emulate the real ReduceOp.MIN cross-rank + reduction *in process*: the post-reduce flag becomes + ``min(local_flag, 1 if peer_ok else 0)``. This is what lets a GPU-free test + exercise the symmetric-failure semantics deterministically. + """ + tp = mock.MagicMock() + tp.world_size = tp_world_size + pp = mock.MagicMock() + pp.world_size = pp_world_size + + world = mock.MagicMock() + # `barrier` must NOT be used by the fix anymore; record it if it ever is so + # a regression to a bare barrier is caught. + world.barrier.side_effect = lambda: calls.append("barrier") + + allocator = mock.MagicMock() + + def _wake(tags=None): + calls.append("alloc_wake") + if allocator_raises: + raise RuntimeError("simulated local cumem wake failure on this rank") + + allocator.wake_up.side_effect = _wake + + platform = mock.MagicMock() + platform.is_cuda_alike.return_value = True + + def _all_reduce(tensor, op=None, group=None): + # Emulate ReduceOp.MIN across this rank + a simulated peer. + calls.append("all_reduce") + assert op == torch.distributed.ReduceOp.MIN, ( + "wake handshake must use ReduceOp.MIN so any single-rank failure " + "(flag=0) forces the global result to 0" + ) + # The fix must route the handshake through the gloo cpu_group, never the + # NCCL device_group. + assert group is world.cpu_group, ( + "wake handshake must run on get_world_group().cpu_group (gloo)" + ) + peer_flag = 1 if peer_ok else 0 + tensor[0] = min(int(tensor[0].item()), peer_flag) + return tensor + + patcher = mock.patch.multiple( + "vllm.v1.worker.gpu_worker", + get_tp_group=mock.MagicMock(return_value=tp), + get_pp_group=mock.MagicMock(return_value=pp), + get_world_group=mock.MagicMock(return_value=world), + get_mem_allocator_instance=mock.MagicMock(return_value=allocator), + current_platform=platform, + ) + dist_patcher = mock.patch( + "vllm.v1.worker.gpu_worker.torch.distributed.all_reduce", + side_effect=_all_reduce, + ) + return patcher, dist_patcher, world, allocator + + +@pytest.mark.parametrize( + "tp,pp", + [(2, 2), (2, 1), (1, 2)], +) +def test_wake_up_handshakes_on_multi_rank_cumem(tp, pp): + """On any multi-rank (TP>1 or PP>1) cumem config the wake handshake fires.""" + worker = _make_worker() + calls: list[str] = [] + patcher, dist_patcher, world, allocator = _patch_groups( + calls, tp_world_size=tp, pp_world_size=pp + ) + with patcher, dist_patcher: + worker.wake_up() + + # Exactly one cross-rank all-reduce handshake; no bare barrier. + assert calls.count("all_reduce") == 1 + assert "barrier" not in calls + # The post_kv_cache_wake_up hook should still run (tags=None path). + worker.model_runner.post_kv_cache_wake_up.assert_called_once() + + +@pytest.mark.parametrize( + "tp,pp", + [(2, 2), (2, 1), (1, 2)], +) +def test_wake_up_handshake_ordered_after_alloc_before_return(tp, pp): + """The handshake must run AFTER the local allocator wake and BEFORE the + method returns control to the caller (which then issues the PP broadcast). + + This is the load-bearing ordering: a handshake that ran before the local + re-map, or one skipped entirely, would leave the #45519 race open. + """ + worker = _make_worker() + calls: list[str] = [] + patcher, dist_patcher, world, allocator = _patch_groups( + calls, tp_world_size=tp, pp_world_size=pp + ) + with patcher, dist_patcher: + worker.wake_up() + + assert calls == ["alloc_wake", "all_reduce"], ( + f"expected local wake then cross-rank handshake, got {calls}" + ) + + +def test_wake_up_no_handshake_on_single_rank(): + """TP=1 PP=1 is unaffected by the PP-broadcast race; no handshake overhead.""" + worker = _make_worker() + calls: list[str] = [] + patcher, dist_patcher, world, allocator = _patch_groups( + calls, tp_world_size=1, pp_world_size=1 + ) + with patcher, dist_patcher: + worker.wake_up() + + assert "all_reduce" not in calls + assert "barrier" not in calls + assert calls == ["alloc_wake"] + + +def test_wake_up_no_handshake_when_cumem_disabled(): + """Without the cumem allocator there is no VMM remap to race against.""" + worker = _make_worker(enable_cumem=False) + calls: list[str] = [] + patcher, dist_patcher, world, allocator = _patch_groups( + calls, tp_world_size=2, pp_world_size=2 + ) + with patcher, dist_patcher: + worker.wake_up() + + assert "all_reduce" not in calls + assert "barrier" not in calls + assert calls == ["alloc_wake"] + + +def test_wake_up_handshake_routes_through_cpu_group_not_device_group(): + """The wake handshake MUST go through the world GroupCoordinator's + CPU-group (gloo), NOT a device_group (NCCL) collective. + + Synchronizing the ranks via the not-yet-resynced NCCL ``device_group`` + would itself issue against the corrupt communicator. We assert ``wake_up`` + drives synchronization *only* via a ``torch.distributed.all_reduce`` on + ``get_world_group().cpu_group`` and never reaches for any group's + ``device_group`` or a GroupCoordinator device collective (``.broadcast``, + ``.all_reduce``, etc., which dispatch to the NCCL device communicator). + """ + worker = _make_worker() + + forbidden_hits: list[str] = [] + FORBIDDEN_METHODS = ( + "broadcast", + "all_reduce", + "all_gather", + "send", + "recv", + "barrier", + ) + + def _make_group(name: str, world_size: int): + grp = mock.MagicMock() + grp.world_size = world_size + for meth in FORBIDDEN_METHODS: + getattr(grp, meth).side_effect = ( + lambda *a, n=name, m=meth, **k: forbidden_hits.append(f"{n}.{m}()") + ) + return grp + + tp = _make_group("tp", 2) + pp = _make_group("pp", 2) + world = _make_group("world", 4) + + allocator = mock.MagicMock() + + cpu_group_used: list[object] = [] + + def _all_reduce(tensor, op=None, group=None): + cpu_group_used.append(group) + tensor[0] = min(int(tensor[0].item()), 1) + return tensor + + with mock.patch.multiple( + "vllm.v1.worker.gpu_worker", + get_tp_group=mock.MagicMock(return_value=tp), + get_pp_group=mock.MagicMock(return_value=pp), + get_world_group=mock.MagicMock(return_value=world), + get_mem_allocator_instance=mock.MagicMock(return_value=allocator), + current_platform=mock.MagicMock( + is_cuda_alike=mock.MagicMock(return_value=True) + ), + ), mock.patch( + "vllm.v1.worker.gpu_worker.torch.distributed.all_reduce", + side_effect=_all_reduce, + ): + worker.wake_up() + + # The raw gloo all-reduce ran exactly once, on the world group's cpu_group. + assert cpu_group_used == [world.cpu_group] + # ...and NO GroupCoordinator device-group collective (incl. NCCL barrier) + # was invoked on any group. + assert forbidden_hits == [], ( + "wake handshake must synchronize via a raw torch.distributed.all_reduce " + "on get_world_group().cpu_group (gloo) only; these GroupCoordinator " + f"device-group collectives were invoked: {forbidden_hits}" + ) + + +# --------------------------------------------------------------------------- +# Adversarial-finding regression: a single-rank local wake failure must NOT +# hang peers, and must surface symmetrically (every rank raises). +# --------------------------------------------------------------------------- + + +def test_wake_up_failed_local_wake_still_participates_in_handshake(): + """If THIS rank's local ``allocator.wake_up`` raises, the rank must STILL + reach the cross-rank handshake (so peers are never stranded waiting) and + then raise loudly. + + PRE-FIX BEHAVIOR (bare barrier *after* the wake, no exception handling): + the exception escaped ``allocator.wake_up`` before ``barrier()`` was + reached, so the failing rank never entered the collective -> every surviving + rank would block forever on the gloo barrier. That is the full-fleet hang + this test guards against. Emulated here as: 'all_reduce' would be ABSENT + from ``calls`` and the surviving peers (in a real cluster) deadlock. + + POST-FIX: the local wake is wrapped in try/except so the all-reduce + ('all_reduce' in ``calls``) runs even on local failure; the rank then raises + a unified abort (chaining the original cause). + """ + worker = _make_worker() + calls: list[str] = [] + patcher, dist_patcher, world, allocator = _patch_groups( + calls, tp_world_size=2, pp_world_size=2, allocator_raises=True + ) + with patcher, dist_patcher: + with pytest.raises( + RuntimeError, match="wake_up failed on at least one rank" + ) as excinfo: + worker.wake_up() + + # The original local failure is chained as the cause (debuggability) on the + # rank that actually failed. + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert "simulated local cumem wake failure" in str(excinfo.value.__cause__) + + # The crux: even though the local wake raised, the rank still entered the + # cross-rank handshake. Absence here == peers hang in a real cluster. + assert calls == ["alloc_wake", "all_reduce"], ( + "a rank whose local wake failed MUST still participate in the handshake " + f"so peers do not hang; got {calls}" + ) + # It must NOT have proceeded to restore buffers / post-kv-cache hooks with + # corrupt state. + worker.model_runner.post_kv_cache_wake_up.assert_not_called() + + +def test_wake_up_raises_symmetrically_when_a_peer_failed(): + """If a PEER rank failed its local wake (but THIS rank's local wake + succeeded), this rank must learn of the failure via the all-reduce and + raise too -- it must NOT silently proceed into a device-group collective + against the peer whose wake never completed. + + PRE-FIX BEHAVIOR (bare barrier): a barrier conveys no success/failure + payload, so a healthy rank sails past it and issues the PP broadcast + against the failed peer -> CUDA_ERROR_ILLEGAL_ADDRESS / NCCL wedge (the + #45519 wedge class, just relocated). Asymmetric: this rank proceeds, peer + is broken. + + POST-FIX: the ReduceOp.MIN all-reduce yields 0 when any peer failed, so + this rank raises symmetrically. + """ + worker = _make_worker() + calls: list[str] = [] + # This rank's local wake SUCCEEDS (allocator_raises=False) but a peer failed + # (peer_ok=False). + patcher, dist_patcher, world, allocator = _patch_groups( + calls, + tp_world_size=2, + pp_world_size=2, + allocator_raises=False, + peer_ok=False, + ) + with patcher, dist_patcher: + with pytest.raises(RuntimeError, match="wake_up failed on at least one rank"): + worker.wake_up() + + # This rank's local wake succeeded, the handshake ran, and it still raised + # (because a peer's flag was 0). No silent proceed. + assert calls == ["alloc_wake", "all_reduce"] + worker.model_runner.post_kv_cache_wake_up.assert_not_called() + + +def test_wake_up_all_ranks_ok_proceeds_normally(): + """Happy path: this rank and all peers succeeded -> all-reduce yields 1, + wake_up returns normally and runs the post-kv-cache hook.""" + worker = _make_worker() + calls: list[str] = [] + patcher, dist_patcher, world, allocator = _patch_groups( + calls, + tp_world_size=2, + pp_world_size=2, + allocator_raises=False, + peer_ok=True, + ) + with patcher, dist_patcher: + worker.wake_up() # must not raise + + assert calls == ["alloc_wake", "all_reduce"] + worker.model_runner.post_kv_cache_wake_up.assert_called_once() diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index c30790df9cac..d449f091a9f2 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -240,6 +240,34 @@ def wake_up(self, tags: list[str] | None = None) -> None: libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None + # Ensure every re-map (`create_and_map`) issued above has completed on + # the device before the caller returns. `create_and_map` performs the + # cuMemCreate/cuMemMap/cuMemSetAccess VMM dance; the freshly re-mapped + # regions are not guaranteed to be coherent for an arbitrary downstream + # stream the instant the loop exits. + # + # NOTE on the `cudaMemcpy` restores above: those use the *synchronous*, + # host-blocking `cudaMemcpy` (NOT `cudaMemcpyAsync`), so each restore has + # already completed w.r.t. the host by the time we get here -- the + # restores are not what this synchronize is guarding. + # + # This `torch.cuda.synchronize()` is a purely LOCAL guarantee: it drains + # this rank's device work, nothing more. It does NOT provide any + # cross-rank ordering. The cross-rank guarantee for the #45519 + # PP-broadcast race lives in `Worker.wake_up`, which gates every rank on + # a CPU (gloo) all-reduce of a wake-success flag before any rank may + # reach a device-group collective. Keeping this local sync here ensures + # a rank's local remaps are actually done before it reports success into + # that cross-rank all-reduce. + # + # (Note: the matching `sleep()` offload loop does NOT synchronize per + # handle between its `cudaMemcpy` D2H backup and the following + # `unmap_and_release` -- the per-handle `torch.cuda.synchronize` lives in + # the pluggable-allocator free callback (`_python_free_callback`), a + # different path. So this is not a "wake-side analogue" of an existing + # sleep-side per-handle sync; it stands on its own.) + torch.cuda.synchronize() + @contextmanager def use_memory_pool(self, tag: str | None = None): """ diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 052e1fe76f43..f458deda9f12 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -13,6 +13,7 @@ import numpy as np import regex as re import torch +import torch.distributed import torch.nn as nn import vllm.envs as envs @@ -42,6 +43,7 @@ Handle, get_pp_group, get_tp_group, + get_world_group, ) from vllm.distributed.weight_transfer import ( WeightTransferEngine, @@ -185,8 +187,85 @@ def sleep(self, level: int = 1) -> None: ) def wake_up(self, tags: list[str] | None = None) -> None: + # Cross-rank wake synchronization (fix for the cumem_tag PP-broadcast + # race, issue #45519). + # + # `wake_up` is dispatched to every worker independently via + # `collective_rpc`, and each worker re-maps its own VMM-backed regions + # at its own pace. With pipeline/tensor parallelism the very next + # decode step issues `torch.distributed.broadcast(..., group=pp.device_ + # group)` (see `_pp_receive_prev_sampled_token_ids_to_input_batch`). + # If a faster rank reaches that collective before a slower rank has + # finished re-mapping the regions backing the broadcast buffers, the + # collective issues against memory the peer's MMU still treats as + # invalid -> CUDA_ERROR_ILLEGAL_ADDRESS, after which the NCCL comm is + # permanently corrupt and the engine deadlocks while /health lies 200. + # + # We therefore must not let any rank proceed to a device-group + # collective until *all* ranks have completed their local wake. A plain + # barrier after the local wake would do that on the happy path -- but if + # one rank's `allocator.wake_up()` raises *before* reaching the barrier, + # the surviving ranks would block forever on the barrier, converting a + # single-rank wake failure into the exact full-fleet hang we are trying + # to fix. + # + # Instead we all-reduce a per-rank success flag over the CPU (gloo) + # group and require unanimity: either every rank succeeded (all proceed) + # or at least one failed (every rank raises). This fails LOUD and + # SYMMETRIC -- no hang, and no rank silently proceeding into a + # device-group collective against a peer whose wake never completed. + # The all-reduce is itself a synchronization point, so it subsumes the + # barrier. We run it on the CPU/gloo group (the same group + # `GroupCoordinator.barrier()` uses) precisely because it never touches + # the not-yet-resynced NCCL device_group. + needs_wake_sync = ( + current_platform.is_cuda_alike() + and self.vllm_config.model_config.enable_cumem_allocator + and (get_tp_group().world_size > 1 or get_pp_group().world_size > 1) + ) + allocator = get_mem_allocator_instance() - allocator.wake_up(tags) + if needs_wake_sync: + # Catch (do not let escape) the local wake failure so this rank + # ALWAYS reaches the cross-rank all-reduce below. If the exception + # were allowed to propagate here, this rank would skip the collective + # and strand every peer waiting in it forever -- the exact full-fleet + # hang we are fixing. We re-raise (chaining the original cause) only + # AFTER the handshake, via the unanimity check. + local_exc: BaseException | None = None + try: + allocator.wake_up(tags) + except BaseException as e: # noqa: BLE001 - re-raised after handshake + local_exc = e + + local_ok = 0 if local_exc is not None else 1 + world_group = get_world_group() + flag = torch.ones(1, dtype=torch.int32) * local_ok + # Raw torch.distributed.all_reduce on the gloo cpu_group (the same + # group GroupCoordinator.barrier() uses), NOT + # GroupCoordinator.all_reduce (which dispatches to the NCCL device + # communicator). ReduceOp.MIN => the global flag is 0 iff ANY rank + # contributed 0, so every rank observes the same all_ok verdict. + torch.distributed.all_reduce( + flag, + op=torch.distributed.ReduceOp.MIN, + group=world_group.cpu_group, + ) + all_ok = bool(flag.item()) + + if not all_ok: + msg = ( + "Worker.wake_up failed on at least one rank; aborting on " + "all ranks to avoid issuing a device-group collective " + "against a peer whose cumem wake did not complete " + f"(issue #45519). This rank's local wake " + f"{'succeeded' if local_ok else 'FAILED'}." + ) + # Preserve the original traceback/cause on the rank that + # actually failed; surviving ranks raise the symmetric abort. + raise RuntimeError(msg) from local_exc + else: + allocator.wake_up(tags) # Restore the buffers after level 2 sleep if len(self._sleep_saved_buffers):