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

[Misc]Reduce BNB static variable (#9987)

Browse Source 
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
This commit is contained in:
Jee Jee Li
2024-11-05 01:04:40 +08:00
committed by GitHub
parent 8d72bb20fa
commit fb2716d641
10 changed files with 20 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -28,7 +28,8 @@ 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.linear import (ReplicatedLinear,
RowParallelLinear)
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig)
from vllm.model_executor.model_loader.tensorizer import (
@@ -727,6 +728,10 @@ class BitsAndBytesModelLoader(BaseModelLoader):
def __init__(self, load_config: LoadConfig):
super().__init__(load_config)
# Save the module names without sharding.
self.unsharded_weights_modules: List[str] = []
# Save the module names that are sharded by column.
self.column_sharded_weights_modules: List[str] = []
# we don't need to quantize the whole model, only the target modules
# that are specified in the adapter config file. If the adapter config
# file is not provided, we will quantize the default modules.
@@ -744,8 +749,6 @@ 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)
@@ -971,9 +974,9 @@ class BitsAndBytesModelLoader(BaseModelLoader):
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):
elif any(
weight_name.startswith(module)
for module in self.column_sharded_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)
@@ -1028,20 +1031,17 @@ class BitsAndBytesModelLoader(BaseModelLoader):
else:
self.target_modules = self.default_target_modules
if hasattr(model, 'column_parallel_weights_modules'):
self.column_parallel_weights_modules = \
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, ))
]
for name, module in model.named_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.
if isinstance(module, (ReplicatedLinear, )):
self.unsharded_weights_modules.append(name)
# In TP, these weights are partitioned along the column
# dimension (dim=-1)
elif isinstance(module, (RowParallelLinear, )):
self.column_sharded_weights_modules.append(name)
self.model_type = type(model).__name__
logger.info("Loading weights with BitsAndBytes quantization. "