biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Model] Add BNB quantization support for Mllama (#9720)

Browse Source
This commit is contained in:
Isotr0py
2024-10-29 20:20:02 +08:00
committed by GitHub
parent ef7865b4f9
commit 09500f7dde
3 changed files with 84 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
vllm/model_executor/model_loader/loader.py
View File

@@ -899,6 +899,19 @@ class BitsAndBytesModelLoader(BaseModelLoader):
return self._unquantized_generator(hf_weights_files, use_safetensors,
quant_state_dict), quant_state_dict
def _is_8bit_weight_name(self, weight_name: str):
quantized_suffix = {".scb", ".weight_format"}
return any(weight_name.lower().endswith(suffix)
for suffix in quantized_suffix)
def _is_4bit_weight_name(self, weight_name: str):
quantized_suffix = {
"absmax", "quant_map", "nested_absmax", "nested_quant_map",
"bitsandbytes"
}
suffix = weight_name.split(".")[-1]
return any(q_suffix in suffix for q_suffix in quantized_suffix)
def _quantized_8bit_generator(self, hf_weights_files, use_safetensors,
quant_state_dict) -> Generator:
for weight_name, weight_tensor in self._hf_weight_iter(
@@ -912,7 +925,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
for weight_name, weight_tensor in self._hf_weight_iter(
hf_weights_files, use_safetensors):
if not weight_name.endswith((".weight", ".bias")):
if self._is_8bit_weight_name(weight_name):
continue
qweight_name = weight_name.replace(".weight", ".qweight")
@@ -932,7 +945,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
use_safetensors)
temp_state_dict = {}
for weight_name, weight_tensor in weight_iterator:
if weight_name.endswith((".weight", ".bias")):
if not self._is_4bit_weight_name(weight_name):
continue
# bitsandbytes library requires
# weight.quant_state.bitsandbytes__* in CPU
@@ -956,7 +969,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
for weight_name, weight_tensor in self._hf_weight_iter(
hf_weights_files, use_safetensors):
if not weight_name.endswith((".weight", ".bias")):
if self._is_4bit_weight_name(weight_name):
continue
if (f"{weight_name}.quant_state.bitsandbytes__nf4" \