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, 16:00 UTC
[Misc] Auto detect bitsandbytes pre-quantized models (#16027)
Signed-off-by: Tristan Leclercq <tristanleclercq@gmail.com>
Tristan Leclercq committed
Apr 5, 2025 at 08:30 UTC
4285e423a6be0bcfdf7ab45fac3815bf37cc2272
3 files changed
+16
-7
docs/source/features/quantization/bnb.md
+6
-3
@@ -19,17 +19,20 @@ And usually, these repositories have a config.json file that includes a quantiza
19
20
## Read quantized checkpoint
21
22
+For pre-quantized checkpoints, vLLM will try to infer the quantization method from the config file, so you don't need to explicitly specify the quantization argument.
23
+
24
```python
25
from vllm import LLM
26
import torch
27
# unsloth/tinyllama-bnb-4bit is a pre-quantized checkpoint.
28
model_id = "unsloth/tinyllama-bnb-4bit"
27
-llm = LLM(model=model_id, dtype=torch.bfloat16, trust_remote_code=True, \
28
-quantization="bitsandbytes")
29
+llm = LLM(model=model_id, dtype=torch.bfloat16, trust_remote_code=True)
30
```
31
32
## Inflight quantization: load as 4bit quantization
33
34
+For inflight 4bit quantization with BitsAndBytes, you need to explicitly specify the quantization argument.
35
+
36
```python
37
from vllm import LLM
38
import torch
@@ -40,7 +43,7 @@ quantization="bitsandbytes")
43
44
## OpenAI Compatible Server
45
43
-Append the following to your 4bit model arguments:
46
+Append the following to your model arguments for 4bit inflight quantization:
47
48
```console
49
--quantization bitsandbytes
tests/quantization/test_bitsandbytes.py
+6
-4
@@ -41,7 +41,7 @@ def test_load_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
41
42
hf_model_kwargs = {"load_in_4bit": True}
43
validate_generated_texts(hf_runner, vllm_runner, example_prompts[:1],
44
- model_name, hf_model_kwargs)
44
+ model_name, False, hf_model_kwargs)
45
46
47
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
@@ -53,7 +53,7 @@ def test_load_pre_quant_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
53
model_name, description) -> None:
54
55
validate_generated_texts(hf_runner, vllm_runner, example_prompts[:1],
56
- model_name)
56
+ model_name, True)
57
58
59
@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
@@ -65,7 +65,7 @@ def test_load_8bit_bnb_model(hf_runner, vllm_runner, example_prompts,
65
model_name, description) -> None:
66
67
validate_generated_texts(hf_runner, vllm_runner, example_prompts[:1],
68
- model_name)
68
+ model_name, True)
69
70
71
@pytest.mark.skipif(torch.cuda.device_count() < 2,
@@ -82,6 +82,7 @@ def test_load_tp_4bit_bnb_model(hf_runner, vllm_runner, example_prompts,
82
vllm_runner,
83
example_prompts[:1],
84
model_name,
85
+ False,
86
hf_model_kwargs,
87
vllm_tp_size=2)
88
@@ -128,13 +129,14 @@ def validate_generated_texts(hf_runner,
129
vllm_runner,
130
prompts,
131
model_name,
132
+ pre_quant=False,
133
hf_model_kwargs=None,
134
vllm_tp_size=1):
135
136
# NOTE: run vLLM first, as it requires a clean process
137
# when using distributed inference
138
with vllm_runner(model_name,
137
- quantization='bitsandbytes',
139
+ quantization=None if pre_quant else 'bitsandbytes',
140
tensor_parallel_size=vllm_tp_size,
141
enforce_eager=False) as llm:
142
vllm_outputs = llm.generate_greedy(prompts, 8)
vllm/engine/arg_utils.py
+4
@@ -1275,6 +1275,10 @@ class EngineArgs:
1275
self.model_loader_extra_config[
1276
"qlora_adapter_name_or_path"] = self.qlora_adapter_name_or_path
1277
1278
+ # bitsandbytes pre-quantized model need a specific model loader
1279
+ if model_config.quantization == "bitsandbytes":
1280
+ self.quantization = self.load_format = "bitsandbytes"
1281
+
1282
load_config = self.create_load_config()
1283
1284
prompt_adapter_config = PromptAdapterConfig(