diff --git a/tests/distributed/test_cumem_vmm_barrier.py b/tests/distributed/test_cumem_vmm_barrier.py new file mode 100644 index 000000000000..27efd9f5986f --- /dev/null +++ b/tests/distributed/test_cumem_vmm_barrier.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests for the cross-rank quiesce barrier the cumem allocator issues at +sleep entry / wake_up exit (fix for vllm-project/vllm#45519). + +These tests deliberately avoid the GPU path: they exercise the +``_quiesce_distributed_before_vmm_mutation`` helper directly under several +mock world-states, asserting that the right primitives are called (or +correctly skipped) on each path. The real PP-broadcast crash that motivates +the fix requires multi-rank NCCL and a sleep-mode-capable backend +(``cumem_tag`` from #45398) — that integration smoke is left to the +``tests/basic_correctness/test_mem.py`` suite running under the +hardware-gated multi-GPU CI. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +# Module under test. Import lazily so the test file is collectable even if +# the cumem C extension isn't built (e.g. on macOS dev machines): the +# fallback path in vllm.device_allocator.cumem leaves the API surface +# defined even when ``cumem_available`` is False. +from vllm.device_allocator import cumem as cumem_module + + +class _Recorder: + """Records sequence of calls so tests can assert ordering.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + + def record(self, name: str) -> Any: + def _inner(*args, **kwargs): + self.calls.append(name) + + return _inner + + +@pytest.fixture +def recorder() -> _Recorder: + return _Recorder() + + +def _make_allocator() -> cumem_module.CuMemAllocator: + """Construct an allocator without invoking the C-extension singleton. + + The class is a singleton via ``get_instance``; for unit tests we want + a fresh instance each time without poisoning module-level state, so we + instantiate directly. ``__init__`` does not touch the C extension — + only the malloc / free callbacks invoke it lazily. + """ + return cumem_module.CuMemAllocator() + + +def test_quiesce_noop_when_torch_distributed_unavailable( + monkeypatch: pytest.MonkeyPatch, recorder: _Recorder +) -> None: + """When ``torch.distributed.is_available()`` returns False (e.g. CPU-only + PyTorch builds), the quiesce helper must not raise and must not attempt + to look up a world group.""" + allocator = _make_allocator() + + # Force the cuda.synchronize call to record but no-op. + monkeypatch.setattr( + cumem_module.torch.cuda, "is_available", lambda: False + ) + # If we ever called barrier, it would error in this CPU-only env. + monkeypatch.setattr( + cumem_module.torch.distributed, "is_available", lambda: False + ) + + # Should return cleanly. + allocator._quiesce_distributed_before_vmm_mutation() + + +def test_quiesce_noop_when_torch_distributed_uninitialized( + monkeypatch: pytest.MonkeyPatch, recorder: _Recorder +) -> None: + """When ``torch.distributed`` is available but not initialized (no + ``init_process_group`` call), the helper must skip the barrier.""" + allocator = _make_allocator() + + monkeypatch.setattr( + cumem_module.torch.cuda, "is_available", lambda: False + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_available", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_initialized", lambda: False + ) + + # If the helper attempted to call barrier in this state torch would raise; + # passing here proves the early-return fired. + allocator._quiesce_distributed_before_vmm_mutation() + + +def test_quiesce_noop_when_kill_switch_disabled( + monkeypatch: pytest.MonkeyPatch, recorder: _Recorder +) -> None: + """When ``_ENABLE_BARRIER_FOR_VMM_MUTATION`` is set False, the helper + must short-circuit before reaching the barrier — even with a fully- + initialized distributed environment.""" + allocator = _make_allocator() + + # Stub out the parallel_state module import path. + fake_parallel_state = SimpleNamespace( + _ENABLE_BARRIER_FOR_VMM_MUTATION=False, + get_world_group=lambda: pytest.fail( + "get_world_group must not be called when the kill switch is off" + ), + ) + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.parallel_state", + fake_parallel_state, + ) + monkeypatch.setattr( + cumem_module.torch.cuda, "is_available", lambda: False + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_available", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_initialized", lambda: True + ) + + allocator._quiesce_distributed_before_vmm_mutation() + + +def test_quiesce_swallows_world_group_lookup_failure( + monkeypatch: pytest.MonkeyPatch, recorder: _Recorder +) -> None: + """If ``get_world_group`` raises (e.g. allocator used outside an inference + engine), the helper must log and continue rather than crash the caller.""" + allocator = _make_allocator() + + def _raise() -> Any: + raise AssertionError("world group not initialized") + + fake_parallel_state = SimpleNamespace( + _ENABLE_BARRIER_FOR_VMM_MUTATION=True, + get_world_group=_raise, + ) + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.parallel_state", + fake_parallel_state, + ) + monkeypatch.setattr( + cumem_module.torch.cuda, "is_available", lambda: False + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_available", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_initialized", lambda: True + ) + + # Must NOT raise. + allocator._quiesce_distributed_before_vmm_mutation() + + +def test_quiesce_calls_barrier_on_world_cpu_group( + monkeypatch: pytest.MonkeyPatch, recorder: _Recorder +) -> None: + """Happy path: with everything wired, the helper must (a) sync cuda and + (b) issue a CPU-side ``barrier`` on the world group's ``cpu_group``, + in that order. The CPU group (not the device group) is deliberate — + NCCL is the subsystem whose buffer registrations are in flux during a + cumem_tag sleep/wake, so we keep our coordination primitive off it.""" + allocator = _make_allocator() + + sentinel_cpu_group = object() + + def _record_synchronize(*args, **kwargs) -> None: + recorder.calls.append("cuda.synchronize") + + def _record_barrier(*args, **kwargs) -> None: + # Assert we got the world's CPU group, not its device group. + assert kwargs.get("group") is sentinel_cpu_group, ( + "quiesce must barrier on CPU group, not device group" + ) + recorder.calls.append("dist.barrier") + + fake_world = SimpleNamespace( + cpu_group=sentinel_cpu_group, + device_group=object(), # distinct sentinel — must NOT be used + ) + fake_parallel_state = SimpleNamespace( + _ENABLE_BARRIER_FOR_VMM_MUTATION=True, + get_world_group=lambda: fake_world, + ) + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.parallel_state", + fake_parallel_state, + ) + monkeypatch.setattr( + cumem_module.torch.cuda, "is_available", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.cuda, "synchronize", _record_synchronize + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_available", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_initialized", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "barrier", _record_barrier + ) + + allocator._quiesce_distributed_before_vmm_mutation() + + # Order matters: drain device-side work before issuing the cross-rank + # ordering primitive, so any in-flight NCCL kernels have completed + # before other ranks proceed past the barrier. + assert recorder.calls == ["cuda.synchronize", "dist.barrier"], ( + f"Expected synchronize then barrier; got {recorder.calls}" + ) + + +def test_quiesce_swallows_barrier_failure( + monkeypatch: pytest.MonkeyPatch, recorder: _Recorder +) -> None: + """Barrier failures must not mask the originating sleep/wake call — + they degrade coordination but should never become the user-visible + exception over the actual VMM mutation.""" + allocator = _make_allocator() + + sentinel_cpu_group = object() + fake_world = SimpleNamespace(cpu_group=sentinel_cpu_group) + fake_parallel_state = SimpleNamespace( + _ENABLE_BARRIER_FOR_VMM_MUTATION=True, + get_world_group=lambda: fake_world, + ) + + def _raising_barrier(*args, **kwargs) -> None: + raise RuntimeError("simulated NCCL/gloo failure") + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.parallel_state", + fake_parallel_state, + ) + monkeypatch.setattr( + cumem_module.torch.cuda, "is_available", lambda: False + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_available", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_initialized", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "barrier", _raising_barrier + ) + + # Must NOT raise. + allocator._quiesce_distributed_before_vmm_mutation() + + +def test_quiesce_swallows_synchronize_failure( + monkeypatch: pytest.MonkeyPatch, recorder: _Recorder +) -> None: + """``torch.cuda.synchronize`` errors must also degrade-not-fail. Symmetric + to the barrier-failure case above.""" + allocator = _make_allocator() + + def _raising_synchronize(*args, **kwargs) -> None: + raise RuntimeError("simulated CUDA context fault") + + monkeypatch.setattr( + cumem_module.torch.cuda, "is_available", lambda: True + ) + monkeypatch.setattr( + cumem_module.torch.cuda, "synchronize", _raising_synchronize + ) + monkeypatch.setattr( + cumem_module.torch.distributed, "is_available", lambda: False + ) + + # Must NOT raise. + allocator._quiesce_distributed_before_vmm_mutation() + + +def test_kill_switch_setter_round_trips(monkeypatch: pytest.MonkeyPatch) -> None: + """``set_enable_cumem_vmm_barrier`` toggles the module-level flag in + both directions and returns the value the getter sees.""" + from vllm.distributed import parallel_state + + original = parallel_state._ENABLE_BARRIER_FOR_VMM_MUTATION + try: + parallel_state.set_enable_cumem_vmm_barrier(False) + assert parallel_state._ENABLE_BARRIER_FOR_VMM_MUTATION is False + + parallel_state.set_enable_cumem_vmm_barrier(True) + assert parallel_state._ENABLE_BARRIER_FOR_VMM_MUTATION is True + finally: + parallel_state._ENABLE_BARRIER_FOR_VMM_MUTATION = original diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index c30790df9cac..a0c14556f402 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -164,6 +164,112 @@ def _python_free_callback(self, ptr: int) -> HandleType: ) return data.handle + def _quiesce_distributed_before_vmm_mutation(self) -> None: + """ + Barrier on every CPU process group and synchronize the device before + we mutate VMM mappings via cuMemUnmap / cuMemCreate+Map. + + Background — fix for vllm-project/vllm#45519: + + ``CuMemAllocator.sleep/wake_up`` walks ``pointer_to_data`` and + unmaps / remaps the VMM range for every cumem-backed allocation. + NCCL communicators bound to ``torch.distributed`` process groups + (e.g. ``get_pp_group().device_group``) hold persistent registrations + of GPU buffers in that same VMM space — rendezvous buffers, ring + topology slots, and any user-registered buffers cached via + NCCL_REGISTER. With the default (``CuMemBackend``) every offloaded + tag is CPU-backed up and the post-wake content is byte-identical, so + any NCCL registration that happens to point at those VAs still sees + plausibly-valid data. With selective-offload backends (notably + ``CuMemTagBackend`` from #45398) some tags are discarded — the VA is + unmapped and remapped to fresh physical pages with garbage — and any + NCCL registration into that range now points at undefined memory. + + Compounding the per-rank state corruption: ``sleep`` / ``wake_up`` + run independently on each rank with no cross-rank ordering. If + rank 0 has already started ``cuMemUnmap`` while rank 1 is still + mid-``torch.distributed.broadcast`` against the soon-to-be-unmapped + VA, the broadcast's P2P send/recv hits a region the peer's GPU MMU + treats as illegal → ``cudaErrorIllegalAddress``. This was the cycle-5 + crash signature reported in #45519 (TP=2 PP=2, varied prompts) and + is also one suspected upstream cause of #45094-class deadlocks. + + Fix shape: at sleep entry and at wake_up exit, (a) drain all + in-flight kernels on the affected device via ``torch.cuda.synchronize()`` + so no NCCL P2P call is mid-flight against the VAs we're about to + mutate, and (b) issue a CPU-side ``barrier()`` on every known + process group so all participating ranks reach the + sleep/wake boundary together. Doing this on the CPU group + (rather than the device group) deliberately avoids round-tripping + through NCCL — which is exactly the subsystem whose buffer + registrations may be in flux. Both calls are no-ops in + single-process / single-rank setups. + + Cost is one ``cudaDeviceSynchronize`` plus one ``gloo`` barrier per + registered process group per sleep/wake call. On a TP=2 PP=2 setup + with a single PP group + a single TP group on the CPU side, that's + ~2 small CPU barriers per call — measured below the noise floor of + a typical wake (~1-2 s on a 27 B AWQ-INT4 model with cumem_tag). + """ + # Drain all in-flight CUDA work on every visible device. Required so + # NCCL kernel launches that were queued behind a recent collective + # call complete before we unmap their buffers, and so post-wake we + # don't return to the caller before remap is visible on the device. + try: + if torch.cuda.is_available(): + # Synchronize the current device first (cheap, common case). + torch.cuda.synchronize() + except Exception as exc: + # synchronize should never raise on a healthy context, but if it + # does, prefer logging over masking the user's actual bug. + logger.warning( + "CuMemAllocator: torch.cuda.synchronize() raised during VMM " + "quiesce: %s. Continuing — sleep/wake may race with in-flight " + "NCCL work.", + exc, + ) + + # CPU-side barrier on every known process group. We deliberately + # use the CPU (gloo) group, not the device (NCCL) group: NCCL is + # the subsystem whose buffer registrations are about to become + # invalid for selective-offload backends. Going CPU-only keeps the + # ordering primitive independent of the corrupted state. + if not torch.distributed.is_available(): + return + if not torch.distributed.is_initialized(): + return + + try: + from vllm.distributed.parallel_state import ( + _ENABLE_BARRIER_FOR_VMM_MUTATION, + get_world_group, + ) + except Exception: + # Importing at module load time risks circular imports with + # vllm.distributed; defer until first use. + return + + if not _ENABLE_BARRIER_FOR_VMM_MUTATION: + return + + try: + world = get_world_group() + except (AssertionError, AttributeError): + # World group not yet initialized (e.g. cumem_allocator used + # outside an inference engine, like in standalone tests). Nothing + # to coordinate. + return + + try: + torch.distributed.barrier(group=world.cpu_group) + except Exception as exc: + # Barrier failures shouldn't mask the originating sleep/wake. + logger.warning( + "CuMemAllocator: CPU-group barrier raised during VMM " + "quiesce: %s. Multi-rank coordination may be degraded.", + exc, + ) + def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: """ Put the allocator in sleep mode. @@ -183,6 +289,11 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: assert isinstance(offload_tags, tuple) + # Drain in-flight NCCL/CUDA work and align all ranks at the sleep + # boundary BEFORE we mutate any VMM mappings. See + # ``_quiesce_distributed_before_vmm_mutation`` for rationale (#45519). + self._quiesce_distributed_before_vmm_mutation() + total_bytes = 0 backup_bytes = 0 @@ -240,6 +351,12 @@ def wake_up(self, tags: list[str] | None = None) -> None: libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None + # Drain restore-memcpy work and align all ranks at the wake boundary + # BEFORE returning control to the caller (who will immediately drive + # NCCL P2P traffic against the freshly-remapped VAs). See + # ``_quiesce_distributed_before_vmm_mutation`` for rationale (#45519). + self._quiesce_distributed_before_vmm_mutation() + @contextmanager def use_memory_pool(self, tag: str | None = None): """ diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 8bd6e92157ab..728bee246510 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1416,12 +1416,29 @@ def graph_capture(device: torch.device): _ENABLE_CUSTOM_ALL_REDUCE = True +# Toggle for the cross-rank quiesce barrier the cumem allocator issues at +# sleep entry / wake_up exit. Default ON: the barrier is cheap (one CPU-side +# gloo barrier on the world group) and required to avoid #45519-class +# ``cudaErrorIllegalAddress`` in NCCL P2P after selective-offload sleep/wake +# (TP×PP with cumem_tag). Provided as a kill switch in case a future setup +# can demonstrate the barrier is unwanted (e.g. embedded use without a +# coordinating CPU group). Toggle via ``set_enable_cumem_vmm_barrier``. +_ENABLE_BARRIER_FOR_VMM_MUTATION = True + def set_custom_all_reduce(enable: bool): global _ENABLE_CUSTOM_ALL_REDUCE _ENABLE_CUSTOM_ALL_REDUCE = enable +def set_enable_cumem_vmm_barrier(enable: bool) -> None: + """Enable/disable the cross-rank quiesce barrier the cumem allocator + issues at sleep entry / wake_up exit. See ``_ENABLE_BARRIER_FOR_VMM_MUTATION`` + for the rationale (vllm-project/vllm#45519).""" + global _ENABLE_BARRIER_FOR_VMM_MUTATION + _ENABLE_BARRIER_FOR_VMM_MUTATION = enable + + def _init_process_group_for_split_group( *, backend: str,