A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
[Misc] Clean up mark to fork process in BNB tests (#20692)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Jul 10, 2025 at 21:59 UTC
77f77a951e68fc9564685dd8041208c00fc22bb1
1 file changed
+11
-18
tests/models/quantization/test_bitsandbytes.py
renamed
+11
-18
@@ -13,8 +13,8 @@ from transformers import BitsAndBytesConfig
13
14
from tests.quantization.utils import is_quant_method_supported
15
16
-from ..models.utils import check_embeddings_close
17
-from ..utils import compare_two_settings, create_new_process_for_each_test
16
+from ...utils import compare_two_settings, multi_gpu_test
17
+from ..utils import check_embeddings_close
18
19
models_4bit_to_test = [
20
("facebook/opt-125m", "quantize opt model inflight"),
@@ -42,7 +42,6 @@ models_pre_quant_8bit_to_test = [
42
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
43
reason='bitsandbytes is not supported on this GPU type.')
44
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
45
-@create_new_process_for_each_test()
45
def test_load_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
46
model_name, description) -> None:
47
@@ -56,7 +55,6 @@ def test_load_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
55
reason='bitsandbytes is not supported on this GPU type.')
56
@pytest.mark.parametrize("model_name, description",
57
models_pre_qaunt_4bit_to_test)
59
-@create_new_process_for_each_test()
58
def test_load_pre_quant_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
59
model_name, description) -> None:
60
@@ -68,7 +66,6 @@ def test_load_pre_quant_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
66
reason='bitsandbytes is not supported on this GPU type.')
67
@pytest.mark.parametrize("model_name, description",
68
models_pre_quant_8bit_to_test)
71
-@create_new_process_for_each_test()
69
def test_load_8bit_bnb_model(hf_runner, vllm_runner, example_prompts,
70
model_name, description) -> None:
71
@@ -76,12 +73,10 @@ def test_load_8bit_bnb_model(hf_runner, vllm_runner, example_prompts,
73
model_name, True)
74
75
79
-@pytest.mark.skipif(torch.cuda.device_count() < 2,
80
- reason='Test requires at least 2 GPUs.')
76
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
77
reason='bitsandbytes is not supported on this GPU type.')
78
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
84
-@create_new_process_for_each_test()
79
+@multi_gpu_test(num_gpus=2)
80
def test_load_tp_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
81
model_name, description) -> None:
82
@@ -96,12 +91,10 @@ def test_load_tp_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
91
vllm_tp_size=2)
92
93
99
-@pytest.mark.skipif(torch.cuda.device_count() < 2,
100
- reason='Test requires at least 2 GPUs.')
94
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
95
reason='bitsandbytes is not supported on this GPU type.')
96
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
104
-@create_new_process_for_each_test()
97
+@multi_gpu_test(num_gpus=2)
98
def test_load_pp_4bit_bnb_model(model_name, description) -> None:
99
common_args = [
100
"--disable-log-stats",
@@ -127,7 +120,6 @@ def test_load_pp_4bit_bnb_model(model_name, description) -> None:
120
@pytest.mark.parametrize("model_name, description",
121
models_4bit_to_embedding_test)
122
@pytest.mark.parametrize("dtype", ["half"])
130
-@create_new_process_for_each_test()
123
def test_4bit_bnb_embedding_model(
124
model_name,
125
description,
@@ -146,6 +138,13 @@ def test_4bit_bnb_embedding_model(
138
example_prompts = [str(s).strip() for s in example_prompts]
139
140
# Inflight 4bit quantization
141
+ with vllm_runner(model_name,
142
+ task="embed",
143
+ dtype=dtype,
144
+ gpu_memory_utilization=0.5,
145
+ quantization="bitsandbytes") as vllm_model:
146
+ vllm_outputs = vllm_model.embed(example_prompts)
147
+
148
hf_model_kwargs = dict(quantization_config=BitsAndBytesConfig(
149
load_in_4bit=True))
150
with hf_runner(
@@ -156,12 +155,6 @@ def test_4bit_bnb_embedding_model(
155
) as hf_model:
156
hf_outputs = hf_model.encode(example_prompts)
157
159
- with vllm_runner(model_name,
160
- task="embed",
161
- dtype=dtype,
162
- gpu_memory_utilization=0.5,
163
- quantization="bitsandbytes") as vllm_model:
164
- vllm_outputs = vllm_model.embed(example_prompts)
158
check_embeddings_close(
159
embeddings_0_lst=hf_outputs,
160
embeddings_1_lst=vllm_outputs,