[CI Test] Add specific unit test for mrv2 offloading (#51440)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 7, 2026 at 18:49 UTC
c39076feff40b5638a2b911414bc09b8e6e953bc
2 files changed
+58
-10
tests/basic_correctness/test_cpu_offload.py
+57
-4
@@ -3,15 +3,22 @@
3
4
import pytest
5
6
+import vllm.envs as envs
7
+from vllm.model_executor.offloader import (
8
+ PrefetchOffloader,
9
+ UVAOffloader,
10
+ get_offloader,
11
+ set_offloader,
12
+)
13
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
14
+
15
from ..utils import compare_two_settings
16
17
18
@pytest.mark.parametrize("disable_pin_memory", [False, True])
19
@pytest.mark.parametrize("disable_uva", [False, True])
11
-@pytest.mark.parametrize("use_v2_model_runner", [False, True])
12
-def test_cpu_offload(disable_pin_memory, disable_uva, use_v2_model_runner):
20
+def test_cpu_offload(disable_pin_memory, disable_uva):
21
env_vars = {
14
- "VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner)),
22
"VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": str(int(disable_pin_memory)),
23
"VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": str(int(disable_uva)),
24
}
@@ -26,6 +33,52 @@ def test_cpu_offload(disable_pin_memory, disable_uva, use_v2_model_runner):
33
model="hmellor/tiny-random-LlamaForCausalLM",
34
arg1=[],
35
arg2=args,
29
- env1=env_vars,
36
+ env1=None,
37
env2=env_vars,
38
)
39
+
40
+
41
+@pytest.mark.parametrize(
42
+ ("offload_kwargs", "offloader_type"),
43
+ [
44
+ ({"cpu_offload_gb": 1}, UVAOffloader),
45
+ (
46
+ {
47
+ "offload_group_size": 1,
48
+ "offload_num_in_group": 1,
49
+ "offload_prefetch_step": 1,
50
+ },
51
+ PrefetchOffloader,
52
+ ),
53
+ ],
54
+)
55
+def test_mrv2_weight_offloading(
56
+ vllm_runner, monkeypatch, offload_kwargs, offloader_type
57
+):
58
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
59
+ envs.disable_envs_cache()
60
+ original_offloader = get_offloader()
61
+
62
+ try:
63
+ with vllm_runner(
64
+ "hmellor/tiny-random-LlamaForCausalLM",
65
+ enforce_eager=True,
66
+ gpu_memory_utilization=0.02,
67
+ max_model_len=128,
68
+ max_num_seqs=1,
69
+ **offload_kwargs,
70
+ ) as vllm_model:
71
+ engine_core = vllm_model.llm.llm_engine.engine_core.engine_core
72
+ model_runner = engine_core.model_executor.driver_worker.worker.model_runner
73
+ assert isinstance(model_runner, GPUModelRunner)
74
+
75
+ offloader = get_offloader()
76
+ assert isinstance(offloader, offloader_type)
77
+ if isinstance(offloader, UVAOffloader):
78
+ assert offloader.cpu_offload_bytes > 0
79
+ else:
80
+ assert offloader.total_offloaded_bytes > 0
81
+ assert offloader.buffer_pool is not None
82
+ finally:
83
+ set_offloader(original_offloader)
84
+ envs.disable_envs_cache()
tests/basic_correctness/test_prefetch_offload.py
+1
-6
@@ -2,13 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Test prefetch offloading correctness with Llama model."""
4
5
-import pytest
6
-
5
from ..utils import compare_two_settings
6
7
10
-@pytest.mark.parametrize("use_v2_model_runner", [False, True])
11
-def test_prefetch_offload_llama(use_v2_model_runner):
8
+def test_prefetch_offload_llama():
9
"""Test prefetch CPU offloading with Llama-3.2-1B-Instruct.
10
11
Compares outputs between:
@@ -33,6 +30,4 @@ def test_prefetch_offload_llama(use_v2_model_runner):
30
"down_proj",
31
],
32
[], # Baseline: no offloading
36
- env1={"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))},
37
- env2={"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))},
33
)