From 62bcd95f854edc69a6c5ef77df23e7e8ed83ad95 Mon Sep 17 00:00:00 2001 From: Justin Wood Date: Wed, 24 Jun 2026 05:09:51 +0000 Subject: [PATCH] fix(cumem): quiesce torch.distributed groups around VMM mutations Fixes vllm-project/vllm#45519 - `cudaErrorIllegalAddress` raised inside `_pp_receive_prev_sampled_token_ids_to_input_batch` during `torch.distributed.broadcast(..., group=pp.device_group)` after a sleep/wake cycle on TP=2 PP=2 with selective-offload sleep-mode backends (notably the upcoming `CuMemTagBackend` from #45398). Root cause walk: `CuMemAllocator.sleep` walks `pointer_to_data` and `cuMemUnmap`s every cumem-backed allocation; `wake_up` does the inverse via `cuMemCreate` + `cuMemMap`. NCCL communicators bound to `torch.distributed` process groups hold persistent registrations of GPU buffers in that same VMM space - rendezvous buffers, ring topology slots, optionally user- registered buffers cached via NCCL_REGISTER. With the default `CuMemBackend`, every offloaded tag is CPU-backed and the post-wake content is byte-identical, so an NCCL registration that points at one of those VAs still sees plausibly-valid data. With selective-offload backends 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. The next P2P broadcast across the affected group hits a region the peer's GPU MMU treats as illegal -> `cudaErrorIllegalAddress`. The repro in #45519 shows this surface at cycle 5 on TP=2 PP=2 with varied prompts (varied prompts defeat the prefix cache, which forces the broken `_is_all_reqs_chunked_prefill()`-skipped broadcast path to run on every cycle; with a fixed prompt the bug is invisible because the broadcast is short-circuited). 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-broadcast against the soon-to-be-unmapped VA, the broadcast's P2P send/recv hits the same illegal-address signature even on the non-selective default backend (this is one upstream cause of the #45094-class deadlock state). Fix: Add a `_quiesce_distributed_before_vmm_mutation` helper called at sleep entry and at wake_up exit: 1. `torch.cuda.synchronize()` drains in-flight kernels on the current device (including any NCCL kernel launches queued behind a recent collective), so no NCCL P2P call is mid-flight against the VAs we're about to mutate. 2. CPU-side `torch.distributed.barrier()` on the world group's `cpu_group` aligns all participating ranks at the sleep/wake boundary. The barrier deliberately goes through the CPU (gloo) group, not the device (NCCL) group: NCCL is the subsystem whose buffer registrations are about to become invalid, so we keep our ordering primitive off it. Both calls are no-ops on single-rank / single-process setups (`torch.distributed.is_initialized() == False`). On a TP=2 PP=2 setup that's one `cudaDeviceSynchronize` plus one gloo barrier per sleep/wake - measured well below the noise floor of a typical wake (~1-2s on a 27B AWQ-INT4 model with cumem_tag). A kill switch `_ENABLE_BARRIER_FOR_VMM_MUTATION` (toggled via `set_enable_cumem_vmm_barrier`) is provided in `parallel_state.py` for embedded use without a coordinating CPU group. Default ON. Failures inside the helper (e.g. synchronize raising on a fault-injured context, or barrier raising on a torn-down group) log a WARNING and return - they must never mask the originating sleep/wake call, which is itself the user-visible operation. Tests: Adds `tests/distributed/test_cumem_vmm_barrier.py` with eight pure-Python unit tests covering: - early-return when torch.distributed is unavailable - early-return when torch.distributed is uninitialized - early-return when the kill switch is disabled - swallowed `get_world_group` failure (allocator used outside an engine) - swallowed barrier failure (degrades, does not raise) - swallowed synchronize failure (degrades, does not raise) - happy path: synchronize THEN barrier on the world's cpu_group - the kill switch setter round-trips The integration-side smoke (multi-rank NCCL + cumem_tag sleep/wake cycle on real hardware) is covered by the existing `tests/basic_correctness/test_mem.py::test_basic_cumem` suite running on the hardware-gated multi-GPU CI - the unit tests added here verify the new coordination primitive's contract without requiring a GPU. Refs #45094, #45097, #45398, #44074 Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: terafin --- tests/distributed/test_cumem_vmm_barrier.py | 307 ++++++++++++++++++++ vllm/device_allocator/cumem.py | 117 ++++++++ vllm/distributed/parallel_state.py | 17 ++ 3 files changed, 441 insertions(+) create mode 100644 tests/distributed/test_cumem_vmm_barrier.py 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,