test_bitsandbytes.py 5.34 KB
Newer Older
1
2
3
4
'''Tests whether bitsandbytes computation is enabled correctly.

Run `pytest tests/quantization/test_bitsandbytes.py`.
'''
5
6

import gc
7
import os
8

9
10
11
import pytest
import torch

12
from tests.quantization.utils import is_quant_method_supported
13

14
15
from ..utils import fork_new_process_for_each_test, models_path_prefix

youkaichao's avatar
youkaichao committed
16

17
models_4bit_to_test = [
18
    (os.path.join(models_path_prefix, 'huggyllama/llama-7b'), 'quantize model inflight'),
19
20
]

21
models_pre_qaunt_4bit_to_test = [
22
    (os.path.join(models_path_prefix, 'lllyasviel/omost-llama-3-8b-4bits'),
23
     'read pre-quantized 4-bit NF4 model'),
24
    (os.path.join(models_path_prefix, 'PrunaAI/Einstein-v6.1-Llama3-8B-bnb-4bit-smashed'),
25
26
27
28
     'read pre-quantized 4-bit FP4 model'),
]

models_pre_quant_8bit_to_test = [
29
    (os.path.join(models_path_prefix, 'meta-llama/Llama-Guard-3-8B-INT8'), 'read pre-quantized 8-bit model'),
30
31
32
33
34
35
]


@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
                    reason='bitsandbytes is not supported on this GPU type.')
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
youkaichao's avatar
youkaichao committed
36
@fork_new_process_for_each_test
37
38
39
40
41
42
43
44
45
46
47
48
def test_load_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
                             model_name, description) -> None:

    hf_model_kwargs = {"load_in_4bit": True}
    validate_generated_texts(hf_runner, vllm_runner, example_prompts[:1],
                             model_name, hf_model_kwargs)


@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
                    reason='bitsandbytes is not supported on this GPU type.')
@pytest.mark.parametrize("model_name, description",
                         models_pre_qaunt_4bit_to_test)
youkaichao's avatar
youkaichao committed
49
@fork_new_process_for_each_test
50
51
52
53
54
55
def test_load_pre_quant_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
                                       model_name, description) -> None:

    validate_generated_texts(hf_runner, vllm_runner, example_prompts[:1],
                             model_name)

56

57
58
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
                    reason='bitsandbytes is not supported on this GPU type.')
59
60
@pytest.mark.parametrize("model_name, description",
                         models_pre_quant_8bit_to_test)
youkaichao's avatar
youkaichao committed
61
@fork_new_process_for_each_test
62
63
64
65
66
67
68
def test_load_8bit_bnb_model(hf_runner, vllm_runner, example_prompts,
                             model_name, description) -> None:

    validate_generated_texts(hf_runner, vllm_runner, example_prompts[:1],
                             model_name)


69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@pytest.mark.skipif(torch.cuda.device_count() < 2,
                    reason='Test requires at least 2 GPUs.')
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
                    reason='bitsandbytes is not supported on this GPU type.')
@pytest.mark.parametrize("model_name, description", models_4bit_to_test)
@fork_new_process_for_each_test
def test_load_tp_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
                                model_name, description) -> None:

    hf_model_kwargs = {"load_in_4bit": True}
    validate_generated_texts(hf_runner,
                             vllm_runner,
                             example_prompts[:1],
                             model_name,
                             hf_model_kwargs,
                             vllm_tp_size=2)


87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
def log_generated_texts(prompts, outputs, runner_name):
    logged_texts = []
    for i, (_, generated_text) in enumerate(outputs):
        log_entry = {
            "prompt": prompts[i],
            "runner_name": runner_name,
            "generated_text": generated_text,
        }
        logged_texts.append(log_entry)
    return logged_texts


def validate_generated_texts(hf_runner,
                             vllm_runner,
                             prompts,
                             model_name,
103
104
                             hf_model_kwargs=None,
                             vllm_tp_size=1):
105

youkaichao's avatar
youkaichao committed
106
107
    # NOTE: run vLLM first, as it requires a clean process
    # when using distributed inference
108
    with vllm_runner(model_name,
109
110
                     quantization='bitsandbytes',
                     load_format='bitsandbytes',
111
                     tensor_parallel_size=vllm_tp_size,
112
                     enforce_eager=False,
113
114
115
116
117
118
119
120
                     gpu_memory_utilization=0.8) as llm:
        vllm_outputs = llm.generate_greedy(prompts, 8)
        vllm_logs = log_generated_texts(prompts, vllm_outputs, "VllmRunner")

    # Clean up the GPU memory for the next test
    gc.collect()
    torch.cuda.empty_cache()

youkaichao's avatar
youkaichao committed
121
122
123
124
125
126
127
128
129
130
131
132
    if hf_model_kwargs is None:
        hf_model_kwargs = {}

    # Run with HF runner
    with hf_runner(model_name, model_kwargs=hf_model_kwargs) as llm:
        hf_outputs = llm.generate_greedy(prompts, 8)
        hf_logs = log_generated_texts(prompts, hf_outputs, "HfRunner")

    # Clean up the GPU memory for the next test
    gc.collect()
    torch.cuda.empty_cache()

133
134
135
136
137
138
139
140
141
142
    # Compare the generated strings
    for hf_log, vllm_log in zip(hf_logs, vllm_logs):
        hf_str = hf_log["generated_text"]
        vllm_str = vllm_log["generated_text"]
        prompt = hf_log["prompt"]
        assert hf_str == vllm_str, (f"Model: {model_name}"
                                    f"Mismatch between HF and vLLM outputs:\n"
                                    f"Prompt: {prompt}\n"
                                    f"HF Output: '{hf_str}'\n"
                                    f"vLLM Output: '{vllm_str}'")