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 14, 22:35 UTC
[Quantization] Enable BNB support for InternS1 (#21953)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Aug 1, 2025 at 19:09 UTC
28b18cc741e596ea6f9981b8365c4819523fc24b
2 files changed
+43
-16
vllm/model_executor/model_loader/bitsandbytes_loader.py
+25
-14
@@ -34,7 +34,8 @@ from vllm.model_executor.model_loader.weight_utils import (
34
filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
35
pt_weights_iterator, safetensors_weights_iterator)
36
from vllm.model_executor.models import is_pooling_model
37
-from vllm.model_executor.utils import (get_packed_modules_mapping,
37
+from vllm.model_executor.utils import (get_moe_expert_mapping,
38
+ get_packed_modules_mapping,
39
set_weight_attrs)
40
from vllm.platforms import current_platform
41
@@ -43,6 +44,12 @@ from vllm.platforms import current_platform
44
logger = init_logger(__name__)
45
46
47
+def is_moe_model(model: torch.nn.Module) -> bool:
48
+ """Checks if the model contains FusedMoE layers."""
49
+ return bool(any(
50
+ isinstance(module, FusedMoE) for module in model.modules()))
51
+
52
+
53
class BitsAndBytesModelLoader(BaseModelLoader):
54
"""Model loader to load model weights with BitAndBytes quantization."""
55
@@ -61,6 +68,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):
68
# Store all module names (from transformers) that support
69
# BNB quantization.
70
self.target_modules: list[str] = []
71
+ # Store the mapping of expert parameters for MoE models.
72
+ self.expert_params_mapping: list[tuple[str, str, int, str]] = []
73
# mapping weight names from transformers to vllm.
74
self.weight_mapper: Callable = lambda name: name
75
self.pre_quant: bool = False
@@ -413,13 +422,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):
422
# in case model has a mixture of disk-merged and disk-split
423
# weights with same last name.
424
self.target_modules.append(name)
416
- elif (isinstance(module, FusedMoE)
417
- and hasattr(module.quant_method, "quant_config")):
418
- if not hasattr(model, "get_expert_mapping"):
419
- raise AttributeError(
420
- f"MoE Model {type(model).__name__} does not support "
421
- "BitsAndBytes quantization yet. Ensure this model has "
422
- "'get_expert_mapping' method.")
425
+ elif isinstance(module, FusedMoE) and hasattr(
426
+ module.quant_method, "quant_config"):
427
# TODO: support FusedMoE with prequant and 8bit.
428
if self.pre_quant:
429
raise ValueError(
@@ -430,9 +434,9 @@ class BitsAndBytesModelLoader(BaseModelLoader):
434
"BitsAndBytes 8bit quantization with FusedMoE is not "
435
"supported yet.")
436
# Get the corresponding weight name using module name and
433
- # get_expert_mapping.
434
- expert_mapping = model.get_expert_mapping()
435
- for exp in expert_mapping:
437
+ # expert_params_mapping.
438
+
439
+ for exp in self.expert_params_mapping:
440
weight_name = exp[1]
441
rep_name = name.replace("experts",
442
"") + weight_name.removesuffix(".")
@@ -464,7 +468,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
468
elif isinstance(module, (RowParallelLinear, )):
469
self.column_sharded_weights_modules.append(name)
470
elif isinstance(module, FusedMoE):
467
- expert_mapping = model.get_expert_mapping()
471
+ expert_mapping = self.expert_params_mapping
472
for exp in expert_mapping:
473
if exp[-1] == "w2":
474
weight_name = exp[1]
@@ -516,6 +520,13 @@ class BitsAndBytesModelLoader(BaseModelLoader):
520
self.is_pool_model = is_pooling_model(model)
521
self.modules_mapping = ParamMapping(get_packed_modules_mapping(model))
522
523
+ if is_moe_model(model):
524
+ self.expert_params_mapping = get_moe_expert_mapping(model)
525
+ if not self.expert_params_mapping:
526
+ raise AttributeError(
527
+ f"MoE Model {type(model).__name__} does not support "
528
+ "BitsAndBytes quantization yet. Ensure this model has "
529
+ "'get_expert_mapping' method.")
530
# For some models like Molmo, we need to use hf_to_vllm_mapper
531
# to ensure correct loading of weights.
532
if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):
@@ -569,10 +580,10 @@ class BitsAndBytesModelLoader(BaseModelLoader):
580
"""
581
from bitsandbytes.functional import QuantState
582
572
- if not hasattr(model, "get_expert_mapping"):
583
+ if not self.expert_params_mapping:
584
return dict()
585
575
- expert_mapping = model.get_expert_mapping()
586
+ expert_mapping = self.expert_params_mapping
587
expert_qs_dict = {}
588
for name, module in model.named_modules():
589
if not isinstance(module, FusedMoE):
vllm/model_executor/utils.py
+18
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Utils for model executor."""
4
+
5
import copy
6
from typing import Any, Optional
7
@@ -9,6 +10,7 @@ import torch
10
11
def set_random_seed(seed: int) -> None:
12
from vllm.platforms import current_platform
13
+
14
current_platform.seed_everything(seed)
15
16
@@ -29,7 +31,7 @@ def set_weight_attrs(
31
return
32
for key, value in weight_attrs.items():
33
assert not hasattr(
32
- weight, key), (f"Overwriting existing tensor attribute: {key}")
34
+ weight, key), f"Overwriting existing tensor attribute: {key}"
35
36
# NOTE(woosuk): During weight loading, we often do something like:
37
# narrowed_tensor = param.data.narrow(0, offset, len)
@@ -41,6 +43,7 @@ def set_weight_attrs(
43
# we sync the param tensor after its weight loader is called.
44
# TODO(woosuk): Remove this hack once we have a better solution.
45
from vllm.platforms import current_platform
46
+
47
if current_platform.is_tpu() and key == "weight_loader":
48
value = _make_synced_weight_loader(value)
49
setattr(weight, key, value)
@@ -77,4 +80,17 @@ def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]:
80
f"safely because of conflicts from {type(child).__name__}.")
81
else:
82
parent_map.update(child_map)
80
- return parent_map
\ No newline at end of file
83
+ return parent_map
84
+
85
+
86
+def get_moe_expert_mapping(
87
+ model: torch.nn.Module, ) -> list[tuple[str, str, int, str]]:
88
+ if parent_map := getattr(model, "get_expert_mapping", None):
89
+ return parent_map()
90
+ else:
91
+ # We only check main components instead of whole model submodules
92
+ for child in model.children():
93
+ child_map = getattr(child, "get_expert_mapping", None)
94
+ if child_map is not None:
95
+ return child_map()
96
+ return []