Skip to content

Commit 77f77a9

Browse files
authored
[Misc] Clean up mark to fork process in BNB tests (#20692)
Signed-off-by: Isotr0py <2037008807@qq.com>
1 parent 1a4f35e commit 77f77a9

File tree

1 file changed

+11
-18
lines changed

1 file changed

+11
-18
lines changed

tests/quantization/test_bitsandbytes.py renamed to tests/models/quantization/test_bitsandbytes.py

Lines changed: 11 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313

1414
from tests.quantization.utils import is_quant_method_supported
1515

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
1818

1919
models_4bit_to_test = [
2020
("facebook/opt-125m", "quantize opt model inflight"),
@@ -42,7 +42,6 @@
4242
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
4343
reason='bitsandbytes is not supported on this GPU type.')
4444
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
45-
@create_new_process_for_each_test()
4645
def test_load_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
4746
model_name, description) -> None:
4847

@@ -56,7 +55,6 @@ def test_load_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
5655
reason='bitsandbytes is not supported on this GPU type.')
5756
@pytest.mark.parametrize("model_name, description",
5857
models_pre_qaunt_4bit_to_test)
59-
@create_new_process_for_each_test()
6058
def test_load_pre_quant_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
6159
model_name, description) -> None:
6260

@@ -68,20 +66,17 @@ def test_load_pre_quant_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
6866
reason='bitsandbytes is not supported on this GPU type.')
6967
@pytest.mark.parametrize("model_name, description",
7068
models_pre_quant_8bit_to_test)
71-
@create_new_process_for_each_test()
7269
def test_load_8bit_bnb_model(hf_runner, vllm_runner, example_prompts,
7370
model_name, description) -> None:
7471

7572
validate_generated_texts(hf_runner, vllm_runner, example_prompts[:1],
7673
model_name, True)
7774

7875

79-
@pytest.mark.skipif(torch.cuda.device_count() < 2,
80-
reason='Test requires at least 2 GPUs.')
8176
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
8277
reason='bitsandbytes is not supported on this GPU type.')
8378
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
84-
@create_new_process_for_each_test()
79+
@multi_gpu_test(num_gpus=2)
8580
def test_load_tp_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
8681
model_name, description) -> None:
8782

@@ -96,12 +91,10 @@ def test_load_tp_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
9691
vllm_tp_size=2)
9792

9893

99-
@pytest.mark.skipif(torch.cuda.device_count() < 2,
100-
reason='Test requires at least 2 GPUs.')
10194
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
10295
reason='bitsandbytes is not supported on this GPU type.')
10396
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
104-
@create_new_process_for_each_test()
97+
@multi_gpu_test(num_gpus=2)
10598
def test_load_pp_4bit_bnb_model(model_name, description) -> None:
10699
common_args = [
107100
"--disable-log-stats",
@@ -127,7 +120,6 @@ def test_load_pp_4bit_bnb_model(model_name, description) -> None:
127120
@pytest.mark.parametrize("model_name, description",
128121
models_4bit_to_embedding_test)
129122
@pytest.mark.parametrize("dtype", ["half"])
130-
@create_new_process_for_each_test()
131123
def test_4bit_bnb_embedding_model(
132124
model_name,
133125
description,
@@ -146,6 +138,13 @@ def test_4bit_bnb_embedding_model(
146138
example_prompts = [str(s).strip() for s in example_prompts]
147139

148140
# 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+
149148
hf_model_kwargs = dict(quantization_config=BitsAndBytesConfig(
150149
load_in_4bit=True))
151150
with hf_runner(
@@ -156,12 +155,6 @@ def test_4bit_bnb_embedding_model(
156155
) as hf_model:
157156
hf_outputs = hf_model.encode(example_prompts)
158157

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)
165158
check_embeddings_close(
166159
embeddings_0_lst=hf_outputs,
167160
embeddings_1_lst=vllm_outputs,

0 commit comments

Comments
 (0)