diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 43232ca6be..0b9cc5a79c 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -7,7 +7,8 @@ from unittest.mock import patch import pytest from vllm.utils.mem_constants import GiB_bytes -from vllm.v1.worker import startup_plan +from vllm.v1.worker import gpu_worker, startup_plan +from vllm.v1.worker.gpu_worker import maybe_rocm_profiling_fallback from vllm.v1.worker.startup_plan import ( maybe_apply_startup_plan, maybe_save_startup_plan, @@ -77,3 +78,77 @@ def test_startup_plan_apply_gate(plan_env): explicit = _plan_worker(kv_bytes=7 * GiB_bytes) maybe_apply_startup_plan(explicit) assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes + + +# Memory accounting of the profiling run (Worker.determine_available_memory). + +# The fallback reads only the sign of the measured drop and this process's torch +# reservation; free memory is only logged, so no amount here is a device size. +ANY_FREE_MEMORY = 8 * GiB_bytes +MEASURED_DROP = 4 * GiB_bytes +TORCH_RESERVED = 3 * GiB_bytes +RELEASED_BY_OTHERS = 2 * GiB_bytes + + +def _snapshot(free_memory, torch_memory=0): + return SimpleNamespace(free_memory=free_memory, torch_memory=torch_memory) + + +def _profile_result(consumed, reserved_before=0, reserved_after=0): + """A result whose free-memory readings agree with `consumed`, which + `memory_profiling` derives as the drop in free memory, negative when it grew.""" + return SimpleNamespace( + total_consumed=consumed, + transient_peak_headroom=0, + before_create=_snapshot(ANY_FREE_MEMORY, reserved_before), + after_profile=_snapshot(ANY_FREE_MEMORY - consumed, reserved_after), + ) + + +@pytest.fixture +def rocm(request): + with patch.object( + gpu_worker, "current_platform", SimpleNamespace(is_rocm=lambda: request.param) + ): + yield request.param + + +@pytest.mark.parametrize("rocm", [True, False], indirect=True) +def test_profiling_fallback_declines_when_free_memory_dropped(rocm): + """The profiling measurement is kept as-is whenever free memory dropped.""" + result = _profile_result(consumed=MEASURED_DROP) + + assert maybe_rocm_profiling_fallback(result) is None + + +@pytest.mark.parametrize("rocm", [True], indirect=True) +def test_profiling_fallback_replaces_a_released_measurement(rocm): + """A negative measurement describes the rest of the device, so it is replaced + by this process's reservation, which the rest of the device cannot move.""" + result = _profile_result( + consumed=-RELEASED_BY_OTHERS, + reserved_after=TORCH_RESERVED, + ) + + assert maybe_rocm_profiling_fallback(result) == TORCH_RESERVED + + +@pytest.mark.parametrize("rocm", [True], indirect=True) +def test_profiling_fallback_never_returns_a_negative_amount(rocm): + """A reservation that shrank across the run cannot become negative usage.""" + result = _profile_result( + consumed=-RELEASED_BY_OTHERS, + reserved_before=TORCH_RESERVED, + reserved_after=0, + ) + + assert maybe_rocm_profiling_fallback(result) == 0 + + +@pytest.mark.parametrize("rocm", [False], indirect=True) +def test_profiling_fallback_declines_off_rocm(rocm): + """Platforms that account frees eagerly keep reporting the error, so the + caller's assertion stays reachable there.""" + result = _profile_result(consumed=-RELEASED_BY_OTHERS) + + assert maybe_rocm_profiling_fallback(result) is None diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 65797bc2a1..15f194ffe5 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -66,7 +66,12 @@ from vllm.tracing import instrument from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes -from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling +from vllm.utils.mem_utils import ( + MemoryProfilingResult, + MemorySnapshot, + format_gib, + memory_profiling, +) from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime from vllm.v1.attention.backends.utils import record_kv_cache_layout from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -102,6 +107,33 @@ def _num_workspace_lanes(vllm_config: VllmConfig, use_v2_model_runner: bool) -> ) +def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int | None: + """Memory in bytes to size the KV cache from when profiling measured a + release on ROCm, or None to keep what profiling measured. + + The measurement is the drop in free device memory across startup, so it comes + out negative when anything else on the device released memory meanwhile. Torch + reserved is process local, so it stands in as a lower bound instead of refusing + to start. Kept to ROCm, where the AMD CI groups hit this, following #45490. + """ + if profile_result.total_consumed >= 0 or not current_platform.is_rocm(): + return None + + baseline = profile_result.before_create + after_profile = profile_result.after_profile + torch_reserved = max(after_profile.torch_memory - baseline.torch_memory, 0) + logger.warning( + "Free memory grew by %s GiB during profiling (initial %s GiB, current " + "%s GiB). Sizing the KV cache from the %s GiB torch reserved instead of " + "the profiling difference.", + format_gib(-profile_result.total_consumed), + format_gib(baseline.free_memory), + format_gib(after_profile.free_memory), + format_gib(torch_reserved), + ) + return torch_reserved + + if TYPE_CHECKING: from vllm.device_allocator.sleep_mode_backend import SleepModeBackend from vllm.model_executor.model_loader.tensorizer import TensorizerConfig @@ -543,24 +575,33 @@ class Worker(WorkerBase): else 0 ) + init_free_memory = self.init_snapshot.free_memory + free_gpu_memory = profile_result.after_profile.free_memory + rocm_fallback = maybe_rocm_profiling_fallback(profile_result) + if rocm_fallback is None: + # NOTE(woosuk): Here we assume that the other processes using the same + # GPU did not change their memory usage during the profiling. + assert init_free_memory >= free_gpu_memory, ( + "Error in memory profiling. " + f"Initial free memory {format_gib(init_free_memory)} GiB, " + f"current free memory {format_gib(free_gpu_memory)} GiB. " + "This happens when other processes sharing the same container " + "release GPU memory while vLLM is profiling during initialization. " + "To fix this, ensure consistent GPU memory allocation or " + "isolate vLLM in its own container." + ) + else: + profile_result.total_consumed = rocm_fallback + profile_result.non_kv_cache_memory = ( + profile_result.total_consumed + profile_result.transient_peak_headroom + ) + self.total_consumed = profile_result.total_consumed self.peak_activation_memory = ( profile_result.transient_peak_headroom + cudagraph_memory_estimate_applied ) self.cudagraph_memory_estimate = cudagraph_memory_estimate - free_gpu_memory = profile_result.after_profile.free_memory - # NOTE(woosuk): Here we assume that the other processes using the same - # GPU did not change their memory usage during the profiling. - assert self.init_snapshot.free_memory >= free_gpu_memory, ( - "Error in memory profiling. " - f"Initial free memory {format_gib(self.init_snapshot.free_memory)} GiB, " - f"current free memory {format_gib(free_gpu_memory)} GiB. " - "This happens when other processes sharing the same container " - "release GPU memory while vLLM is profiling during initialization. " - "To fix this, ensure consistent GPU memory allocation or " - "isolate vLLM in its own container." - ) self.available_kv_cache_memory_bytes = ( self.requested_memory - profile_result.non_kv_cache_memory