diff --git a/tests/models/language/pooling/test_reward.py b/tests/models/language/pooling/test_reward.py index 1872ca4ae09..f981266d0c0 100644 --- a/tests/models/language/pooling/test_reward.py +++ b/tests/models/language/pooling/test_reward.py @@ -9,6 +9,7 @@ import torch.nn.functional as F from transformers import AutoModel from vllm.platforms import current_platform +from vllm.utils.mem_constants import MiB_bytes from ....conftest import HfRunner from ....utils import VLLM_PATH @@ -106,7 +107,12 @@ def test_prm_models( if current_platform.is_cpu(): pytest.skip("CPU only supports V1") - with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model: + with vllm_runner( + model, + max_model_len=1024, + dtype=dtype, + kv_cache_memory_bytes=64 * MiB_bytes, + ) as vllm_model: vllm_outputs = vllm_model.token_classify(math_step_prompts) with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model: @@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs( if not FIXTURE_REWARD_RESULT.get(model): pytest.skip(f"No available golden outputs for {model}.") - with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model: + with vllm_runner( + model, + max_model_len=1024, + dtype=dtype, + kv_cache_memory_bytes=64 * MiB_bytes, + ) as vllm_model: vllm_outputs = vllm_model.token_classify(math_step_prompts) golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model]) diff --git a/tests/utils_/test_mem_utils.py b/tests/utils_/test_mem_utils.py index 421aec3e9b1..19690d1f34f 100644 --- a/tests/utils_/test_mem_utils.py +++ b/tests/utils_/test_mem_utils.py @@ -67,6 +67,23 @@ def test_memory_profiling(): non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa assert abs(non_torch_ratio - 1) <= 0.05 assert result.torch_peak_increase == 1024 * 1024 * 1024 + + expected_total_consumed = (256 + 512) * 1024 * 1024 + total_consumed_ratio = result.total_consumed / expected_total_consumed + assert abs(total_consumed_ratio - 1) <= 0.05, ( + f"total_consumed={result.total_consumed}, " + f"expected={expected_total_consumed}, " + f"ratio={total_consumed_ratio}" + ) + + expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024 + non_kv_ratio = result.non_kv_cache_memory / expected_non_kv + assert abs(non_kv_ratio - 1) <= 0.05, ( + f"non_kv_cache_memory={result.non_kv_cache_memory}, " + f"expected={expected_non_kv}, " + f"ratio={non_kv_ratio}" + ) + del weights lib.cudaFree(handle1) lib.cudaFree(handle2) diff --git a/vllm/utils/mem_utils.py b/vllm/utils/mem_utils.py index b0ac4b16e47..6a522a150d3 100644 --- a/vllm/utils/mem_utils.py +++ b/vllm/utils/mem_utils.py @@ -110,6 +110,7 @@ class MemorySnapshot: """Memory snapshot.""" torch_peak: int = 0 + torch_allocated: int = 0 free_memory: int = 0 total_memory: int = 0 cuda_memory: int = 0 @@ -139,9 +140,9 @@ class MemorySnapshot: # After `torch.accelerator.reset_peak_memory_stats()`, # `torch.accelerator.memory_reserved()` will keep growing, and only shrink # when we call `torch.accelerator.empty_cache()` or OOM happens. - self.torch_peak = torch.accelerator.memory_stats(device).get( - "allocated_bytes.all.peak", 0 - ) + stats = torch.accelerator.memory_stats(device) + self.torch_peak = stats.get("allocated_bytes.all.peak", 0) + self.torch_allocated = stats.get("allocated_bytes.all.current", 0) self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device) if current_platform.is_integrated_gpu(device.index): @@ -172,6 +173,7 @@ class MemorySnapshot: return MemorySnapshot( torch_peak=self.torch_peak - other.torch_peak, + torch_allocated=self.torch_allocated - other.torch_allocated, free_memory=self.free_memory - other.free_memory, total_memory=self.total_memory - other.total_memory, cuda_memory=self.cuda_memory - other.cuda_memory, @@ -185,6 +187,7 @@ class MemorySnapshot: def __repr__(self) -> str: return ( f"torch_peak={format_gib(self.torch_peak)}GiB, " + f"torch_allocated={format_gib(self.torch_allocated)}GiB, " f"free_memory={format_gib(self.free_memory)}GiB, " f"total_memory={format_gib(self.total_memory)}GiB, " f"{current_platform.device_name}_memory={format_gib(self.cuda_memory)}GiB, " @@ -202,6 +205,8 @@ class MemoryProfilingResult: non_kv_cache_memory: int = 0 torch_peak_increase: int = 0 non_torch_increase: int = 0 + total_consumed: int = 0 + transient_peak_headroom: int = 0 weights_memory: int = 0 before_create: MemorySnapshot = field(default_factory=MemorySnapshot) profile_time: float = 0.0 @@ -219,8 +224,8 @@ class MemoryProfilingResult: f"{format_gib(self.non_kv_cache_memory)}GiB; " f"torch peak memory increase: " f"{format_gib(self.torch_peak_increase)}GiB; " - f"non-torch forward increase memory: " - f"{format_gib(self.non_torch_increase)}GiB; " + f"total consumed (from mem_get_info): " + f"{format_gib(self.total_consumed)}GiB; " f"weights memory: {format_gib(self.weights_memory)}GiB." ) @@ -306,8 +311,16 @@ def memory_profiling( result.non_torch_increase = diff_from_create.non_torch_memory result.profile_time = diff_profile.timestamp - non_torch_memory = result.non_torch_increase - peak_activation_memory = result.torch_peak_increase - result.non_kv_cache_memory = ( - non_torch_memory + peak_activation_memory + result.weights_memory + # Measure total consumption via mem_get_info() instead of + # memory_reserved(), which goes negative when pluggable allocators + # (e.g. cumem) bypass PyTorch's tracking. + result.total_consumed = ( + result.before_create.free_memory - result.after_profile.free_memory ) + + # total_consumed already covers persistent torch allocations; add only the + # transient peak headroom to avoid double-counting. + result.transient_peak_headroom = ( + result.after_profile.torch_peak - result.after_profile.torch_allocated + ) + result.non_kv_cache_memory = result.total_consumed + result.transient_peak_headroom diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 7ac35ad2329..0c89b2fb13e 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -488,33 +488,18 @@ class Worker(WorkerBase): ) as profile_result: self.model_runner.profile_run() - profile_torch_peak = torch.accelerator.memory_stats(self.device).get( - "allocated_bytes.all.peak", 0 - ) - - # Profile CUDA graph memory if graphs will be captured. - # ROCm is included: #44825 moved the profiler to - # torch.accelerator.get_memory_info (reliable on ROCm, as used by - # the AMD-CI mem tests), and graph_pool_handle resolves to the same - # torch.cuda handle the live capture path already uses on ROCm. - # XPU stays excluded (see #39977). - cudagraph_memory_estimate = 0 - if ( - current_platform.is_cuda_alike() - and self.vllm_config.compilation_config.cudagraph_mode - != CUDAGraphMode.NONE - ): - cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory() - - # Use the pre-cudagraph torch peak to avoid double-counting. - profile_result.torch_peak_increase = ( - profile_torch_peak - profile_result.before_profile.torch_peak - ) - profile_result.non_kv_cache_memory = ( - profile_result.non_torch_increase - + profile_result.torch_peak_increase - + profile_result.weights_memory - ) + # Profile CUDA graph memory if graphs will be captured. + # ROCm is included: #44825 moved the profiler to + # torch.accelerator.get_memory_info (reliable on ROCm, as used by + # the AMD-CI mem tests), and graph_pool_handle resolves to the same + # torch.cuda handle the live capture path already uses on ROCm. + # XPU stays excluded (see #39977). + cudagraph_memory_estimate = 0 + if ( + current_platform.is_cuda_alike() + and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + ): + cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory() # Respect the opt-in flag as originally designed. cudagraph_memory_estimate_applied = ( @@ -523,8 +508,10 @@ class Worker(WorkerBase): else 0 ) - self.non_torch_memory = profile_result.non_torch_increase - self.peak_activation_memory = profile_result.torch_peak_increase + 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 @@ -817,9 +804,8 @@ class Worker(WorkerBase): redundancy_buffer_memory = 150 * (1 << 20) non_kv_cache_memory = ( - self.model_runner.model_memory_usage + self.total_consumed + self.peak_activation_memory - + self.non_torch_memory + cuda_graph_memory_bytes ) kv_cache_memory_bytes_to_gpu_limit = ( @@ -840,10 +826,10 @@ class Worker(WorkerBase): f"Desired GPU memory utilization is " f"({self.cache_config.gpu_memory_utilization}, " f"{format_gib(self.requested_memory)} GiB). " - f"Actual usage is {format_gib(self.model_runner.model_memory_usage)} " - f"GiB for weight, {format_gib(self.peak_activation_memory)} GiB " - f"for peak activation, {format_gib(self.non_torch_memory)} GiB " - f"for non-torch memory, and {format_gib(cuda_graph_memory_bytes)} " + f"Actual usage is {format_gib(self.total_consumed)} " + f"GiB for consumed memory (weights + non-torch), " + f"{format_gib(self.peak_activation_memory)} GiB " + f"for peak activation, and {format_gib(cuda_graph_memory_bytes)} " f"GiB for CUDAGraph memory. Replace gpu_memory_utilization " f"config with `--kv-cache-memory=" f"{kv_cache_memory_bytes_to_requested_limit}` "