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 15, 02:15 UTC
[Quantization] add BNB for MixtralForCausalLM (#20893)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Jul 14, 2025 at 15:34 UTC
a99b9f7dee0ad261284cbcd823f5b37381d15ac1
7 files changed
+128
-20
vllm/model_executor/model_loader/utils.py
+6
-1
@@ -227,7 +227,12 @@ def get_model_architecture(
227
# Special handling for quantized Mixtral.
228
# FIXME(woosuk): This is a temporary hack.
229
mixtral_supported = [
230
- "fp8", "compressed-tensors", "gptq_marlin", "awq_marlin", "quark"
230
+ "fp8",
231
+ "compressed-tensors",
232
+ "gptq_marlin",
233
+ "awq_marlin",
234
+ "quark",
235
+ "bitsandbytes",
236
]
237
238
vllm_supported_archs = ModelRegistry.get_supported_archs()
vllm/model_executor/models/granitemoe.py
+102
-3
@@ -45,12 +45,14 @@ from vllm.model_executor.layers.quantization.base_config import (
45
from vllm.model_executor.layers.rotary_embedding import get_rope
46
from vllm.model_executor.layers.vocab_parallel_embedding import (
47
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
48
+from vllm.model_executor.model_loader.weight_utils import (
49
+ default_weight_loader, maybe_remap_kv_scale_name)
50
from vllm.model_executor.sampling_metadata import SamplingMetadata
51
from vllm.sequence import IntermediateTensors
52
51
-from . import mixtral
53
from .interfaces import SupportsLoRA, SupportsPP
53
-from .utils import AutoWeightsLoader, make_layers, maybe_prefix
54
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter, make_layers,
55
+ maybe_prefix)
56
57
58
class GraniteMoeMoE(nn.Module):
@@ -307,6 +309,103 @@ class GraniteMoeModel(nn.Module):
309
hidden_states = self.norm(hidden_states)
310
return hidden_states
311
312
+ def _load_weights(self,
313
+ weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
314
+ """
315
+ This function is copied from `MixtralModel.load_weights`, mainly to
316
+ decouple from mixtral, avoiding impact on support like BNB
317
+ quantization.
318
+ """
319
+ stacked_params_mapping = [
320
+ # (param_name, shard_name, shard_id)
321
+ ("qkv_proj", "q_proj", "q"),
322
+ ("qkv_proj", "k_proj", "k"),
323
+ ("qkv_proj", "v_proj", "v"),
324
+ ]
325
+
326
+ # Params for weights, fp8 weight scales, fp8 activation scales
327
+ # (param_name, weight_name, expert_id, shard_id)
328
+ expert_params_mapping = FusedMoE.make_expert_params_mapping(
329
+ ckpt_gate_proj_name="w1",
330
+ ckpt_down_proj_name="w2",
331
+ ckpt_up_proj_name="w3",
332
+ num_experts=self.config.num_local_experts)
333
+
334
+ params_dict = dict(self.named_parameters())
335
+ loaded_params: set[str] = set()
336
+ for name, loaded_weight in weights:
337
+ if (self.quant_config is not None and
338
+ (scale_name := self.quant_config.get_cache_scale(name))):
339
+ # Loading kv cache quantization scales
340
+ param = params_dict[scale_name]
341
+ weight_loader = getattr(param, "weight_loader",
342
+ default_weight_loader)
343
+ loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
344
+ loaded_weight[0])
345
+ weight_loader(param, loaded_weight)
346
+ loaded_params.add(scale_name)
347
+ continue
348
+
349
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
350
+ if weight_name not in name:
351
+ continue
352
+ name = name.replace(weight_name, param_name)
353
+ # Skip loading extra bias for GPTQ models.
354
+ if ((name.endswith(".bias") or name.endswith("_bias"))
355
+ and name not in params_dict):
356
+ continue
357
+ # Skip layers on other devices.
358
+ if is_pp_missing_parameter(name, self):
359
+ continue
360
+ if name.endswith("scale"):
361
+ # Remapping the name of FP8 kv-scale.
362
+ name = maybe_remap_kv_scale_name(name, params_dict)
363
+ if name is None:
364
+ continue
365
+ param = params_dict[name]
366
+ weight_loader = param.weight_loader
367
+ weight_loader(param, loaded_weight, shard_id)
368
+ break
369
+ else:
370
+ for mapping in expert_params_mapping:
371
+ param_name, weight_name, expert_id, shard_id = mapping
372
+ if weight_name not in name:
373
+ continue
374
+ name = name.replace(weight_name, param_name)
375
+ # Skip layers on other devices.
376
+ if is_pp_missing_parameter(name, self):
377
+ continue
378
+ if ((name.endswith(".bias") or name.endswith("_bias"))
379
+ and name not in params_dict):
380
+ continue
381
+ param = params_dict[name]
382
+ weight_loader = param.weight_loader
383
+ weight_loader(param,
384
+ loaded_weight,
385
+ name,
386
+ shard_id=shard_id,
387
+ expert_id=expert_id)
388
+ break
389
+ else:
390
+ # Skip loading extra bias for GPTQ models.
391
+ if ((name.endswith(".bias") or name.endswith("_bias"))
392
+ and name not in params_dict):
393
+ continue
394
+ # Skip layers on other devices.
395
+ if is_pp_missing_parameter(name, self):
396
+ continue
397
+ # Remapping the name of FP8 kv-scale.
398
+ name = maybe_remap_kv_scale_name(name, params_dict)
399
+ if name is None:
400
+ continue
401
+
402
+ param = params_dict[name]
403
+ weight_loader = getattr(param, "weight_loader",
404
+ default_weight_loader)
405
+ weight_loader(param, loaded_weight)
406
+ loaded_params.add(name)
407
+ return loaded_params
408
+
409
def load_weights(self, weights: Iterable[tuple[str,
410
torch.Tensor]]) -> set[str]:
411
new_weights = {}
@@ -339,7 +438,7 @@ class GraniteMoeModel(nn.Module):
438
new_weights[gate_name] = p
439
else:
440
new_weights[n] = p
342
- return mixtral.MixtralModel.load_weights(self, new_weights.items())
441
+ return self._load_weights(new_weights.items())
442
443
444
class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
vllm/model_executor/models/granitemoeshared.py
+2
-3
@@ -27,8 +27,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
27
from vllm.model_executor.sampling_metadata import SamplingMetadata
28
from vllm.sequence import IntermediateTensors
29
30
-from . import mixtral
31
-from .granitemoe import GraniteMoeAttention, GraniteMoeMoE
30
+from .granitemoe import GraniteMoeAttention, GraniteMoeModel, GraniteMoeMoE
31
from .interfaces import SupportsLoRA, SupportsPP
32
from .utils import AutoWeightsLoader, make_layers, maybe_prefix
33
@@ -242,7 +241,7 @@ class GraniteMoeSharedModel(nn.Module):
241
new_weights[gate_name] = p
242
else:
243
new_weights[n] = p
245
- return mixtral.MixtralModel.load_weights(self, new_weights.items())
244
+ return GraniteMoeModel._load_weights(self, new_weights.items())
245
246
247
class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
vllm/model_executor/models/mixtral.py
+13
-8
@@ -317,6 +317,15 @@ class MixtralModel(nn.Module):
317
hidden_states, _ = self.norm(hidden_states, residual)
318
return hidden_states
319
320
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
321
+ # Params for weights, fp8 weight scales, fp8 activation scales
322
+ # (param_name, weight_name, expert_id, shard_id)
323
+ return FusedMoE.make_expert_params_mapping(
324
+ ckpt_gate_proj_name="w1",
325
+ ckpt_down_proj_name="w2",
326
+ ckpt_up_proj_name="w3",
327
+ num_experts=self.config.num_local_experts)
328
+
329
def load_weights(self, weights: Iterable[tuple[str,
330
torch.Tensor]]) -> set[str]:
331
stacked_params_mapping = [
@@ -326,16 +335,9 @@ class MixtralModel(nn.Module):
335
("qkv_proj", "v_proj", "v"),
336
]
337
329
- # Params for weights, fp8 weight scales, fp8 activation scales
330
- # (param_name, weight_name, expert_id, shard_id)
331
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
332
- ckpt_gate_proj_name="w1",
333
- ckpt_down_proj_name="w2",
334
- ckpt_up_proj_name="w3",
335
- num_experts=self.config.num_local_experts)
336
-
338
params_dict = dict(self.named_parameters())
339
loaded_params: set[str] = set()
340
+ expert_params_mapping = self.get_expert_mapping()
341
for name, loaded_weight in weights:
342
if (self.quant_config is not None and
343
(scale_name := self.quant_config.get_cache_scale(name))):
@@ -486,3 +488,6 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
488
torch.Tensor]]) -> set[str]:
489
loader = AutoWeightsLoader(self)
490
return loader.load_weights(weights)
491
+
492
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
493
+ return self.model.get_expert_mapping()
vllm/model_executor/models/olmoe.py
+2
-1
@@ -352,6 +352,7 @@ class OlmoeModel(nn.Module):
352
353
params_dict = dict(self.named_parameters())
354
loaded_params: set[str] = set()
355
+ expert_params_mapping = self.get_expert_mapping()
356
for name, loaded_weight in weights:
357
for (param_name, weight_name, shard_id) in stacked_params_mapping:
358
# Skip non-stacked layers and experts (experts handled below).
@@ -380,7 +381,7 @@ class OlmoeModel(nn.Module):
381
weight_loader(param, loaded_weight, shard_id)
382
break
383
else:
383
- for mapping in self.get_expert_mapping():
384
+ for mapping in expert_params_mapping:
385
param_name, weight_name, expert_id, shard_id = mapping
386
if weight_name not in name:
387
continue
vllm/model_executor/models/qwen2_moe.py
+2
-1
@@ -413,6 +413,7 @@ class Qwen2MoeModel(nn.Module):
413
414
params_dict = dict(self.named_parameters())
415
loaded_params: set[str] = set()
416
+ expert_params_mapping = self.get_expert_mapping()
417
for name, loaded_weight in weights:
418
for (param_name, weight_name, shard_id) in stacked_params_mapping:
419
# Skip non-stacked layers and experts (experts handled below).
@@ -442,7 +443,7 @@ class Qwen2MoeModel(nn.Module):
443
weight_loader(param, loaded_weight, shard_id)
444
break
445
else:
445
- for mapping in self.get_expert_mapping():
446
+ for mapping in expert_params_mapping:
447
param_name, weight_name, expert_id, shard_id = mapping
448
if weight_name not in name:
449
continue
vllm/model_executor/models/qwen3_moe.py
+1
-3
@@ -400,11 +400,9 @@ class Qwen3MoeModel(nn.Module):
400
".v_scale", "_v_scale", ".weight_scale",
401
"_weight_scale", ".input_scale", "_input_scale")
402
403
- # Params for weights, fp8 weight scales, fp8 activation scales
404
- # (param_name, weight_name, expert_id, shard_id)
405
- expert_params_mapping = self.get_expert_mapping()
403
params_dict = dict(self.named_parameters())
404
loaded_params: set[str] = set()
405
+ expert_params_mapping = self.get_expert_mapping()
406
for name, loaded_weight in weights:
407
for (param_name, weight_name, shard_id) in stacked_params_mapping:
408
# Skip non-stacked layers and experts (experts handled below).