wassname/vllm
1
0
Fork 0
mirror of https://github.com/wassname/vllm.git synced 2026-09-12 13:01:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Misc]Add BNB quantization for MolmoForCausalLM (#11551)

Browse Source 
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
This commit is contained in:
Jee Jee Li
2024-12-27 18:48:24 +00:00
committed by GitHub
parent 55509c2114
commit 0240402c46
2 changed files with 83 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/model_loader/loader.py
+18 -8
View File
@@ -11,7 +11,8 @@ import os
import warnings
from abc import ABC, abstractmethod
from contextlib import contextmanager
from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, cast
from typing import (Any, Callable, Dict, Generator, Iterable, List, Optional,
Tuple, cast)
import gguf
import huggingface_hub
@@ -706,6 +707,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):
# Store all module names (from transformers) that support
# BNB quantization.
self.target_modules: List[str] = []
# mapping weight names from transformers to vllm.
self.weight_mapper: Callable = lambda name: name
def _get_weight_files(
self,
@@ -763,9 +766,12 @@ class BitsAndBytesModelLoader(BaseModelLoader):
def _hf_weight_iter(self, hf_weights_files, use_safetensors: bool):
if use_safetensors:
return safetensors_weights_iterator(hf_weights_files)
iterator = safetensors_weights_iterator(hf_weights_files)
else:
return pt_weights_iterator(hf_weights_files)
iterator = pt_weights_iterator(hf_weights_files)
for name, param in iterator:
# mapping weight names from transformers to vllm.
yield self.weight_mapper(name), param
def _get_quantized_weights_iterator(
self,
@@ -782,12 +788,12 @@ class BitsAndBytesModelLoader(BaseModelLoader):
try:
import bitsandbytes
if bitsandbytes.__version__ < "0.44.0":
if bitsandbytes.__version__ < "0.45.0":
raise ImportError("bitsandbytes version is wrong. Please "
"install bitsandbytes>=0.44.0.")
"install bitsandbytes>=0.45.0.")
except ImportError as err:
raise ImportError("Please install bitsandbytes>=0.44.0 via "
"`pip install bitsandbytes>=0.44.0` to use "
raise ImportError("Please install bitsandbytes>=0.45.0 via "
"`pip install bitsandbytes>=0.45.0` to use "
"bitsandbytes quantizer.") from err
hf_weights_files, use_safetensors = self._prepare_weights(
@@ -991,7 +997,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
if isinstance(module, (LinearBase, )):
last_name = name.split(".")[-1]
if sub_modules := inverse_stacked_mapping.get(last_name, []):
# Map vllm's names to transformers' names.
# Map vllm's names to transformers's names.
for sub_name in sub_modules:
self.target_modules.append(
name.replace(last_name, sub_name))
@@ -1013,6 +1019,10 @@ class BitsAndBytesModelLoader(BaseModelLoader):
f"Model {type(model).__name__} does not support BitsAndBytes "
"quantization yet.")
# For some models like Molmo, we need to use hf_to_vllm_mapper
# to ensure correct loading of weights.
if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):
self.weight_mapper = lambda name: hf_to_vllm_mapper._map_name(name)
# Modules whose weights might have fused on disk
# we need their output_sizes to make shard in flight correctly with TP
self.maybe_fused_weights_modules: Dict[str, List[int]] = {}