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, 17:15 UTC
[Quantization][V1] BitsAndBytes support V1 (#15611)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Mar 28, 2025 at 10:12 UTC
726efc6a320ad9a4ef0b0378b40abbd0561ea394
7 files changed
+52
-24
tests/models/encoder_decoder/vision_language/test_mllama.py
-1
@@ -425,7 +425,6 @@ def test_bnb_regression(
425
max_model_len=4096,
426
max_num_seqs=2,
427
quantization="bitsandbytes",
428
- load_format="bitsandbytes",
428
)
429
sampling_params = SamplingParams(
430
temperature=0,
tests/models/test_transformers.py
-1
@@ -72,7 +72,6 @@ def test_distributed(
72
"meta-llama/Llama-3.2-1B-Instruct",
73
{
74
"quantization": "bitsandbytes",
75
- "load_format": "bitsandbytes",
75
},
76
),
77
])
tests/quantization/test_bitsandbytes.py
-3
@@ -101,8 +101,6 @@ def test_load_pp_4bit_bnb_model(model_name, description) -> None:
101
"--enable-prefix-caching",
102
"--quantization",
103
"bitsandbytes",
104
- "--load-format",
105
- "bitsandbytes",
104
"--gpu-memory-utilization",
105
"0.7",
106
]
@@ -137,7 +135,6 @@ def validate_generated_texts(hf_runner,
135
# when using distributed inference
136
with vllm_runner(model_name,
137
quantization='bitsandbytes',
140
- load_format='bitsandbytes',
138
tensor_parallel_size=vllm_tp_size,
139
enforce_eager=False) as llm:
140
vllm_outputs = llm.generate_greedy(prompts, 8)
vllm/config.py
+4
-2
@@ -682,8 +682,9 @@ class ModelConfig:
682
683
def _verify_bnb_config(self) -> None:
684
"""
685
- The current version of bitsandbytes (0.44.0) with 8-bit models does not
685
+ The current version of bitsandbytes (0.45.3) with 8-bit models does not
686
yet support CUDA graph.
687
+ # TODO Remove this when bitsandbytes supports.
688
"""
689
is_bitsandbytes = self.quantization == "bitsandbytes"
690
has_quantization_config = (getattr(self.hf_config,
@@ -698,8 +699,9 @@ class ModelConfig:
699
not self.enforce_eager,
700
]):
701
logger.warning(
701
- "CUDA graph is not supported on BitAndBytes 8bit yet, "
702
+ "CUDA graph is not supported on BitsAndBytes 8bit yet, "
703
"fallback to the eager mode.")
704
+
705
self.enforce_eager = True
706
707
def _verify_with_expert_parallelism(self) -> None:
vllm/engine/arg_utils.py
+1
-1
@@ -1616,7 +1616,7 @@ class EngineArgs:
1616
return False
1617
1618
# Some quantization is not compatible with torch.compile.
1619
- V1_UNSUPPORTED_QUANT = ["bitsandbytes", "gguf"]
1619
+ V1_UNSUPPORTED_QUANT = ["gguf"]
1620
if model_config.quantization in V1_UNSUPPORTED_QUANT:
1621
_raise_or_fallback(
1622
feature_name=f"--quantization {model_config.quantization}",
vllm/model_executor/layers/quantization/bitsandbytes.py
+45
-16
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
9
set_weight_attrs)
10
from vllm.model_executor.layers.quantization.base_config import (
11
QuantizationConfig)
12
+from vllm.utils import direct_register_custom_op
13
14
15
class BitsAndBytesConfig(QuantizationConfig):
@@ -321,9 +322,6 @@ class BitsAndBytesLinearMethod(LinearMethodBase):
322
x: torch.Tensor,
323
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
324
324
- # only load the bitsandbytes module when needed
325
- from bitsandbytes import matmul_4bit
326
-
325
original_type = x.dtype
326
original_shape = x.shape
327
reshape_after_matmul = False
@@ -343,19 +341,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):
341
out_dim_1,
342
dtype=torch.bfloat16,
343
device=x.device)
346
-
347
- current_index = 0
348
- for i in range(len(quant_states)):
349
- output_size = quant_states[i].shape[0]
350
- # It is more efficient to use out kwarg like
351
- # matmul_4bit(..., out = ...). Infeasible now due to the bug
352
- # https://github.com/TimDettmers/bitsandbytes/issues/1235.
353
- # Need to change after the bug is fixed.
354
- out[:, current_index:current_index + output_size] = matmul_4bit(
355
- bf_x, qweight[offsets[i]:offsets[i + 1]].t(), quant_states[i])
356
-
357
- current_index += output_size
358
-
344
+ apply_bnb_4bit(bf_x, qweight, offsets, out)
345
out = out.to(original_type)
346
347
if reshape_after_matmul:
@@ -365,3 +351,46 @@ class BitsAndBytesLinearMethod(LinearMethodBase):
351
out += bias
352
353
return out
354
+
355
+
356
+def _apply_bnb_4bit(
357
+ x: torch.Tensor,
358
+ weight: torch.Tensor,
359
+ offsets: torch.Tensor,
360
+ out: torch.Tensor,
361
+) -> None:
362
+ # only load the bitsandbytes module when needed
363
+ from bitsandbytes import matmul_4bit
364
+ quant_states = weight.bnb_quant_state
365
+ current_index = 0
366
+ for i in range(len(quant_states)):
367
+ output_size = quant_states[i].shape[0]
368
+ # It is more efficient to use out kwarg like
369
+ # matmul_4bit(..., out = ...). Infeasible now due to the bug
370
+ # https://github.com/TimDettmers/bitsandbytes/issues/1235.
371
+ # Need to change after the bug is fixed.
372
+ out[:, current_index:current_index + output_size] = matmul_4bit(
373
+ x, weight[offsets[i]:offsets[i + 1]].t(), quant_states[i])
374
+ current_index += output_size
375
+
376
+
377
+def _apply_bnb_4bit_fake(
378
+ x: torch.Tensor,
379
+ weight: torch.Tensor,
380
+ offsets: torch.Tensor,
381
+ out: torch.Tensor,
382
+) -> None:
383
+ return
384
+
385
+
386
+try:
387
+ direct_register_custom_op(
388
+ op_name="apply_bnb_4bit",
389
+ op_func=_apply_bnb_4bit,
390
+ mutates_args=["out"],
391
+ fake_impl=_apply_bnb_4bit_fake,
392
+ )
393
+ apply_bnb_4bit = torch.ops.vllm.apply_bnb_4bit
394
+
395
+except AttributeError as error:
396
+ raise error
vllm/model_executor/model_loader/loader.py
+2
@@ -1259,6 +1259,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):
1259
pack_ratio)
1260
1261
offsets = np.concatenate(([0], np.cumsum(num_elements)))
1262
+ # Make torch infer_schema happy
1263
+ offsets = torch.tensor(offsets).cpu()
1264
set_weight_attrs(param, {"bnb_shard_offsets": offsets})
1265
1266
if load_8bit: