@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 150 lines 5.54 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from unittest.mock import MagicMock, patch
4
5 import torch
6 from vllm_test_utils.monitor import monitor
7
8 from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
9
10 from ..utils import create_new_process_for_each_test
11
12
13 @create_new_process_for_each_test()
14 def test_memory_profiling():
15 # Fake out some model loading + inference memory usage to test profiling
16 # Memory used by other processes will show up as cuda usage outside of torch
17 from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
18
19 lib = CudaRTLibrary()
20 # 512 MiB allocation outside of this instance
21 handle1 = lib.cudaMalloc(512 * 1024 * 1024)
22
23 # Warm up PyTorch's CUDA/ROCm context so that its internal initialization
24 # overhead (streams, cuBLAS handles, etc.) is included in the baseline and
25 # does not inflate non-torch increase which is larger on ROCm than on CUDA
26 _warmup = torch.zeros(1, device="cuda")
27 del _warmup
28 torch.accelerator.empty_cache()
29
30 baseline_snapshot = MemorySnapshot()
31
32 # load weights
33
34 weights = torch.randn(128, 1024, 1024, device="cuda", dtype=torch.float32)
35
36 weights_memory = 128 * 1024 * 1024 * 4 # 512 MiB
37
38 def measure_current_non_torch():
39 free, total = torch.accelerator.get_memory_info()
40 current_used = total - free
41 current_torch = torch.accelerator.memory_reserved()
42 current_non_torch = current_used - current_torch
43 return current_non_torch
44
45 with (
46 memory_profiling(
47 baseline_snapshot=baseline_snapshot, weights_memory=weights_memory
48 ) as result,
49 monitor(measure_current_non_torch) as monitored_values,
50 ):
51 # make a memory spike, 1 GiB
52 spike = torch.randn(256, 1024, 1024, device="cuda", dtype=torch.float32)
53 del spike
54
55 # Add some extra non-torch memory 256 MiB (simulate NCCL)
56 handle2 = lib.cudaMalloc(256 * 1024 * 1024)
57
58 # this is an analytic value, it is exact,
59 # we only have 256 MiB non-torch memory increase
60 measured_diff = monitored_values.values[-1] - monitored_values.values[0]
61 assert measured_diff == 256 * 1024 * 1024
62
63 # Check that the memory usage is within 5% of the expected values
64 # 5% tolerance is caused by cuda runtime.
65 # we cannot control cuda runtime in the granularity of bytes,
66 # which causes a small error (<10 MiB in practice)
67 non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa
68 assert abs(non_torch_ratio - 1) <= 0.05
69 assert result.torch_peak_increase == 1024 * 1024 * 1024
70
71 expected_total_consumed = (256 + 512) * 1024 * 1024
72 total_consumed_ratio = result.total_consumed / expected_total_consumed
73 assert abs(total_consumed_ratio - 1) <= 0.05, (
74 f"total_consumed={result.total_consumed}, "
75 f"expected={expected_total_consumed}, "
76 f"ratio={total_consumed_ratio}"
77 )
78
79 expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
80 non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
81 assert abs(non_kv_ratio - 1) <= 0.05, (
82 f"non_kv_cache_memory={result.non_kv_cache_memory}, "
83 f"expected={expected_non_kv}, "
84 f"ratio={non_kv_ratio}"
85 )
86
87 del weights
88 lib.cudaFree(handle1)
89 lib.cudaFree(handle2)
90
91
92 def test_memory_snapshot_uses_psutil_on_integrated_gpu():
93 """On integrated (UMA) GPUs, free_memory should come from psutil."""
94 mock_cuda_free = 40 * 1024**3
95 mock_cuda_total = 120 * 1024**3
96 mock_psutil_available = 100 * 1024**3
97
98 with (
99 patch("vllm.utils.mem_utils.current_platform") as mock_platform,
100 patch("vllm.utils.mem_utils.psutil") as mock_psutil,
101 patch("torch.accelerator") as mock_accelerator,
102 ):
103 mock_accelerator.get_memory_info.return_value = (
104 mock_cuda_free,
105 mock_cuda_total,
106 )
107 mock_platform.is_integrated_gpu.return_value = True
108 mock_platform.memory_stats.return_value = {
109 "allocated_bytes.all.peak": 0,
110 }
111 mock_accelerator.memory_reserved.return_value = 0
112 mock_accelerator.current_device = lambda: "cuda:0"
113
114 mock_vmem = MagicMock()
115 mock_vmem.available = mock_psutil_available
116 mock_psutil.virtual_memory.return_value = mock_vmem
117
118 snapshot = MemorySnapshot(device="cuda:0")
119
120 assert snapshot.free_memory == mock_psutil_available
121 assert snapshot.total_memory == mock_cuda_total
122 mock_psutil.virtual_memory.assert_called_once()
123
124
125 def test_memory_snapshot_uses_cuda_on_discrete_gpu():
126 """On discrete GPUs, free_memory should come from accelerator get_memory_info."""
127 mock_cuda_free = 70 * 1024**3
128 mock_cuda_total = 80 * 1024**3
129
130 with (
131 patch("vllm.utils.mem_utils.current_platform") as mock_platform,
132 patch("vllm.utils.mem_utils.psutil") as mock_psutil,
133 patch("torch.accelerator") as mock_accelerator,
134 ):
135 mock_accelerator.get_memory_info.return_value = (
136 mock_cuda_free,
137 mock_cuda_total,
138 )
139 mock_platform.is_integrated_gpu.return_value = False
140 mock_accelerator.memory_stats.return_value = {
141 "allocated_bytes.all.peak": 0,
142 }
143 mock_accelerator.memory_reserved.return_value = 0
144 mock_accelerator.current_device = lambda: "cuda:0"
145
146 snapshot = MemorySnapshot(device="cuda:0")
147
148 assert snapshot.free_memory == mock_cuda_free
149 assert snapshot.total_memory == mock_cuda_total
150 mock_psutil.virtual_memory.assert_not_called()