[Bugfix] Fix MiniCPMV and Mllama BNB bug (#9917)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
This commit is contained in:
Jee Jee Li
2024-11-04 11:36:41 +08:00
committed by GitHub
parent 91c9ebbb1b
commit c49f0407ba
4 changed files with 145 additions and 65 deletions

View File

@@ -28,6 +28,7 @@ from vllm.distributed import (get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size)
from vllm.envs import VLLM_USE_MODELSCOPE
from vllm.logger import init_logger
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig)
from vllm.model_executor.model_loader.tensorizer import (
@@ -771,6 +772,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):
with open(config_file_path, "r") as f:
config = json.load(f)
self.target_modules = config["target_modules"]
# Save the module names without sharding.
self.unsharded_weights_modules: List[str] = []
def _get_config_file(self, qlora_adapter: str) -> str:
is_local = os.path.isdir(qlora_adapter)
@@ -990,16 +993,21 @@ class BitsAndBytesModelLoader(BaseModelLoader):
if any(target_module in weight_name for target_module in
self.target_modules) and weight_name.endswith(".weight"):
weight_name = weight_name.replace(".weight", ".qweight")
if any(module in weight_name
for module in self.column_parallel_weights_modules):
# Without sharding
if any(
weight_name.startswith(module)
for module in self.unsharded_weights_modules):
weight_sub_tensor = weight_tensor
# Shard by column
elif any(module in weight_name
for module in self.column_parallel_weights_modules):
total_size = weight_tensor.size(-1)
start_index = total_size // tp_size * tp_rank
end_index = total_size // tp_size * (tp_rank + 1)
weight_sub_tensor = weight_tensor[...,
start_index:end_index]
# Shard by row
else:
total_size = weight_tensor.size(0)
start_index = total_size // tp_size * tp_rank
@@ -1053,7 +1061,15 @@ class BitsAndBytesModelLoader(BaseModelLoader):
model.column_parallel_weights_modules
else:
self.column_parallel_weights_modules = []
# Some modules like `ReplicatedLinear` should not have their weights
# sharded. The reason for implementing it this way is to avoid new
# static variable in the model implementation.
# TODO: Can we reduce the static variables needed for BNB based on
# model information?
self.unsharded_weights_modules = [
name for name, module in model.named_modules()
if isinstance(module, (ReplicatedLinear, ))
]
self.model_type = type(model).__name__
logger.info("Loading weights with BitsAndBytes quantization. "
@@ -1100,7 +1116,13 @@ class BitsAndBytesModelLoader(BaseModelLoader):
for shard_name, (
weight_name, index
) in model.bitsandbytes_stacked_params_mapping.items():
if shard_name in quant_param_name:
shard_pos = quant_param_name.find(shard_name)
# Some models, such as MiniCPM V2.5/2.6, contain both
# module names 'kv_proj' and 'qkv_proj'. To prevent 'kv_proj'
# from being incorrectly identified as being present in
# 'vpm.encoder.layers.0.self_attn.qkv_proj.qweight
if shard_pos > 0 and quant_param_name[shard_pos - 1] == ".":
shard_index = index
quant_param_name = quant_param_name.replace(
shard_name, weight_name)