@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8020e98c9f
Fork 0

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][1/N] MoE support BNB-Inflight Quantization (#20061)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Jul 11, 2025 at 16:01 UTC 8020e98c9f033e76c97eb8261f772d59eba49c9a
8 files changed +561 -88
tests/models/quantization/test_bitsandbytes.py
+39 -6
@@ -14,7 +14,7 @@ from transformers import BitsAndBytesConfig
14 from tests.quantization.utils import is_quant_method_supported
15
16 from ...utils import compare_two_settings, multi_gpu_test
17 -from ..utils import check_embeddings_close
17 +from ..utils import check_embeddings_close, check_logprobs_close
18
19 models_4bit_to_test = [
20 ("facebook/opt-125m", "quantize opt model inflight"),
@@ -26,6 +26,10 @@ models_4bit_to_embedding_test = [
26 ("intfloat/e5-mistral-7b-instruct", "quantize embedding model inflight"),
27 ]
28
29 +models_4bit_to_moe_test = [
30 + ("allenai/OLMoE-1B-7B-0125-Instruct", "quantize moe model inflight"),
31 +]
32 +
33 models_pre_qaunt_4bit_to_test = [
34 ('PrunaAI/Einstein-v6.1-Llama3-8B-bnb-4bit-smashed',
35 'read pre-quantized 4-bit FP4 model'),
@@ -115,6 +119,35 @@ def test_load_pp_4bit_bnb_model(model_name, description) -> None:
119 compare_two_settings(model_name, common_args, pp_args)
120
121
122 +@pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
123 + reason='bitsandbytes is not supported on this GPU type.')
124 +@pytest.mark.parametrize("model_name, description", models_4bit_to_moe_test)
125 +def test_4bit_bnb_moe_model(hf_runner, vllm_runner, example_prompts,
126 + model_name, description) -> None:
127 +
128 + hf_model_kwargs = dict(quantization_config=BitsAndBytesConfig(
129 + load_in_4bit=True,
130 + bnb_4bit_quant_type="nf4",
131 + bnb_4bit_use_double_quant=True,
132 + ))
133 + with vllm_runner(model_name,
134 + quantization='bitsandbytes',
135 + enforce_eager=False) as llm:
136 + vllm_outputs = llm.generate_greedy_logprobs(example_prompts,
137 + max_tokens=32,
138 + num_logprobs=5)
139 +
140 + with hf_runner(model_name, model_kwargs=hf_model_kwargs) as llm:
141 + transformers_outputs = llm.generate_greedy_logprobs_limit(
142 + example_prompts, max_tokens=32, num_logprobs=5)
143 + check_logprobs_close(
144 + outputs_0_lst=transformers_outputs,
145 + outputs_1_lst=vllm_outputs,
146 + name_0="transformers",
147 + name_1="vllm",
148 + )
149 +
150 +
151 @pytest.mark.skipif(not is_quant_method_supported("bitsandbytes"),
152 reason='bitsandbytes is not supported on this GPU type.')
153 @pytest.mark.parametrize("model_name, description",
@@ -182,7 +215,8 @@ def validate_generated_texts(hf_runner,
215 model_name,
216 pre_quant=False,
217 hf_model_kwargs=None,
185 - vllm_tp_size=1):
218 + vllm_tp_size=1,
219 + max_tokens=8):
220
221 # NOTE: run vLLM first, as it requires a clean process
222 # when using distributed inference
@@ -190,7 +224,8 @@ def validate_generated_texts(hf_runner,
224 quantization=None if pre_quant else 'bitsandbytes',
225 tensor_parallel_size=vllm_tp_size,
226 enforce_eager=False) as llm:
193 - vllm_outputs = llm.generate_greedy(prompts, 8)
227 +
228 + vllm_outputs = llm.generate_greedy(prompts, max_tokens)
229 vllm_logs = log_generated_texts(prompts, vllm_outputs, "VllmRunner")
230
231 # Clean up the GPU memory for the next test
@@ -202,19 +237,17 @@ def validate_generated_texts(hf_runner,
237
238 # Run with HF runner
239 with hf_runner(model_name, model_kwargs=hf_model_kwargs) as llm:
205 - hf_outputs = llm.generate_greedy(prompts, 8)
240 + hf_outputs = llm.generate_greedy(prompts, max_tokens)
241 hf_logs = log_generated_texts(prompts, hf_outputs, "HfRunner")
242
243 # Clean up the GPU memory for the next test
244 gc.collect()
245 torch.cuda.empty_cache()
211 -
246 # Compare the generated strings
247 for hf_log, vllm_log in zip(hf_logs, vllm_logs):
248 hf_str = hf_log["generated_text"]
249 vllm_str = vllm_log["generated_text"]
250 prompt = hf_log["prompt"]
217 -
251 assert hf_str == vllm_str, (f"Model: {model_name}"
252 f"Mismatch between HF and vLLM outputs:\n"
253 f"Prompt: {prompt}\n"
vllm/model_executor/layers/fused_moe/layer.py
+32 -4
@@ -883,14 +883,21 @@ class FusedMoE(torch.nn.Module):
883 expert_data=expert_data,
884 tp_rank=tp_rank)
885
886 - def _load_w13(self, expert_data: torch.Tensor, shard_dim: int,
887 - shard_id: str, loaded_weight: torch.Tensor, tp_rank: int):
886 + def _load_w13(self,
887 + expert_data: torch.Tensor,
888 + shard_dim: int,
889 + shard_id: str,
890 + loaded_weight: torch.Tensor,
891 + tp_rank: int,
892 + load_full: bool = False):
893
894 # Index the loaded weight for tp sharding.
895 # gate_up_proj: "MergedColumnParallel", so tp sharding on output_dim
896 shard_size = expert_data.shape[shard_dim] // 2
892 - loaded_weight = loaded_weight.narrow(shard_dim, shard_size * tp_rank,
893 - shard_size)
897 + if not load_full:
898 + loaded_weight = loaded_weight.narrow(shard_dim,
899 + shard_size * tp_rank,
900 + shard_size)
901 # Narrow parameter and load.
902 # w1, gate_proj: Load into first logical weight of w13.
903 if shard_id == "w1":
@@ -998,6 +1005,27 @@ class FusedMoE(torch.nn.Module):
1005 param.data.copy_(loaded_weight)
1006 return True if return_success else None
1007
1008 + # Case for BitsAndBytes
1009 + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
1010 + if use_bitsandbytes_4bit:
1011 + shard_dim = 0
1012 +
1013 + expert_data = param.data[expert_id]
1014 + if shard_id == "w2":
1015 + expert_data.copy_(loaded_weight)
1016 + elif shard_id in ("w1", "w3"):
1017 + # BNB inflight quantization has already sharded the weights
1018 + full_load = True
1019 + self._load_w13(
1020 + shard_id=shard_id,
1021 + shard_dim=shard_dim,
1022 + loaded_weight=loaded_weight,
1023 + expert_data=expert_data,
1024 + tp_rank=self.tp_rank,
1025 + load_full=full_load,
1026 + )
1027 + return True if return_success else None
1028 +
1029 # is_transposed: if the dim to shard the weight
1030 # should be flipped. Required by GPTQ, compressed-tensors
1031 # should be whatever dimension intermediate_size_per_partition is
vllm/model_executor/layers/quantization/bitsandbytes.py
+223 -9
@@ -1,10 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any, Optional
4 +from typing import Any, Callable, Optional, Union
5
6 import torch
7
8 +from vllm.model_executor.layers.fused_moe import fused_experts
9 +from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
10 + FusedMoEMethodBase)
11 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
12 UnquantizedLinearMethod,
13 set_weight_attrs)
@@ -120,12 +123,15 @@ class BitsAndBytesConfig(QuantizationConfig):
123 llm_int8_skip_modules=llm_int8_skip_modules,
124 llm_int8_threshold=llm_int8_threshold)
125
123 - def get_quant_method(self, layer: torch.nn.Module,
124 - prefix: str) -> Optional["LinearMethodBase"]:
126 + def get_quant_method(
127 + self, layer: torch.nn.Module, prefix: str
128 + ) -> Optional[Union["LinearMethodBase", "BitsAndBytesMoEMethod"]]:
129 if isinstance(layer, LinearBase):
130 if is_layer_skipped_bnb(prefix, self.llm_int8_skip_modules):
131 return UnquantizedLinearMethod()
132 return BitsAndBytesLinearMethod(self)
133 + elif isinstance(layer, FusedMoE):
134 + return BitsAndBytesMoEMethod(self)
135 return None
136
137
@@ -146,6 +152,13 @@ def is_layer_skipped_bnb(prefix: str, llm_int8_skip_modules: list[str]):
152 return substr_check or prefix_check
153
154
155 +def calculate_quant_ratio(dtype):
156 + if dtype.is_floating_point:
157 + return torch.finfo(dtype).bits // torch.iinfo(torch.uint8).bits
158 + else:
159 + return torch.iinfo(dtype).bits // torch.iinfo(torch.uint8).bits
160 +
161 +
162 class BitsAndBytesLinearMethod(LinearMethodBase):
163 """Linear method for BitsAndBytes.
164
@@ -173,12 +186,6 @@ class BitsAndBytesLinearMethod(LinearMethodBase):
186 **extra_weight_attrs):
187 from bitsandbytes.nn import Int8Params
188
176 - def calculate_quant_ratio(dtype):
177 - if dtype.is_floating_point:
178 - return torch.finfo(dtype).bits // torch.iinfo(torch.uint8).bits
179 - else:
180 - return torch.iinfo(dtype).bits // torch.iinfo(torch.uint8).bits
181 -
189 def create_qweight_for_8bit():
190 qweight = Int8Params(
191 data=torch.empty(sum(output_partition_sizes),
@@ -394,3 +401,210 @@ try:
401
402 except AttributeError as error:
403 raise error
404 +
405 +
406 +class BitsAndBytesMoEMethod(FusedMoEMethodBase):
407 + """MoE method for BitsAndBytes.
408 +
409 + Args:
410 + quant_config: The BitsAndBytes quantization config.
411 + """
412 +
413 + def __init__(self, quant_config: BitsAndBytesConfig):
414 + try:
415 + import bitsandbytes
416 + if bitsandbytes.__version__ < "0.45.3":
417 + raise ImportError("bitsandbytes version is wrong. Please "
418 + "install bitsandbytes>=0.45.3.")
419 + except ImportError as err:
420 + raise ImportError("Please install bitsandbytes>=0.45.3 via "
421 + "`pip install bitsandbytes>=0.45.3` to use "
422 + "bitsandbytes quantizer.") from err
423 + self.topk_indices_dtype = None
424 + self.quant_config = quant_config
425 +
426 + def create_weights(
427 + self,
428 + layer: torch.nn.Module,
429 + num_experts: int,
430 + hidden_size: int,
431 + intermediate_size_per_partition: int,
432 + params_dtype: torch.dtype,
433 + **extra_weight_attrs,
434 + ):
435 + if self.quant_config.load_in_8bit:
436 + call_fun = self._create_weights_8bit
437 + else:
438 + call_fun = self._create_weights_4bit
439 + call_fun(
440 + layer,
441 + num_experts,
442 + hidden_size,
443 + intermediate_size_per_partition,
444 + params_dtype,
445 + **extra_weight_attrs,
446 + )
447 +
448 + def apply(
449 + self,
450 + layer: torch.nn.Module,
451 + x: torch.Tensor,
452 + router_logits: torch.Tensor,
453 + top_k: int,
454 + renormalize: bool,
455 + use_grouped_topk: bool = False,
456 + topk_group: Optional[int] = None,
457 + num_expert_group: Optional[int] = None,
458 + global_num_experts: int = -1,
459 + expert_map: Optional[torch.Tensor] = None,
460 + custom_routing_function: Optional[Callable] = None,
461 + scoring_func: str = "softmax",
462 + e_score_correction_bias: Optional[torch.Tensor] = None,
463 + apply_router_weight_on_input: bool = False,
464 + activation: str = "silu",
465 + enable_eplb: bool = False,
466 + expert_load_view: Optional[torch.Tensor] = None,
467 + logical_to_physical_map: Optional[torch.Tensor] = None,
468 + logical_replica_count: Optional[torch.Tensor] = None,
469 + ) -> torch.Tensor:
470 +
471 + if enable_eplb:
472 + raise NotImplementedError(
473 + "EPLB not supported for `BitsAndBytesMoEMethod` yet.")
474 + topk_weights, topk_ids = FusedMoE.select_experts(
475 + hidden_states=x,
476 + router_logits=router_logits,
477 + use_grouped_topk=use_grouped_topk,
478 + top_k=top_k,
479 + renormalize=renormalize,
480 + topk_group=topk_group,
481 + num_expert_group=num_expert_group,
482 + custom_routing_function=custom_routing_function,
483 + scoring_func=scoring_func,
484 + e_score_correction_bias=e_score_correction_bias,
485 + indices_type=self.topk_indices_dtype)
486 + if self.quant_config.load_in_8bit:
487 + w13, w2 = self._apply_8bit_dequant(layer)
488 + else:
489 + w13, w2 = self._apply_4bit_dequnt(layer)
490 + return fused_experts(
491 + hidden_states=x,
492 + w1=w13,
493 + w2=w2,
494 + topk_weights=topk_weights,
495 + topk_ids=topk_ids,
496 + inplace=True,
497 + activation=activation,
498 + apply_router_weight_on_input=apply_router_weight_on_input,
499 + global_num_experts=global_num_experts,
500 + expert_map=expert_map,
501 + )
502 +
503 + def _create_weights_4bit(
504 + self,
505 + layer: torch.nn.Module,
506 + num_experts: int,
507 + hidden_size: int,
508 + intermediate_size_per_partition: int,
509 + params_dtype: torch.dtype,
510 + **extra_weight_attrs,
511 + ):
512 + quant_ratio = calculate_quant_ratio(params_dtype)
513 + # Fused gate_up_proj (column parallel)
514 + w13_total_size = (hidden_size * 2 *
515 + intermediate_size_per_partition) // quant_ratio
516 + w13_qweight = torch.nn.Parameter(
517 + torch.empty(
518 + num_experts,
519 + w13_total_size,
520 + 1,
521 + dtype=torch.uint8,
522 + ),
523 + requires_grad=False,
524 + )
525 + layer.register_parameter("w13_weight", w13_qweight)
526 + set_weight_attrs(w13_qweight, extra_weight_attrs)
527 + set_weight_attrs(
528 + w13_qweight,
529 + {
530 + "num_experts":
531 + num_experts,
532 + "input_dim":
533 + hidden_size,
534 + "output_dim":
535 + 2 * intermediate_size_per_partition,
536 + "experts_shape": (
537 + num_experts,
538 + intermediate_size_per_partition * 2,
539 + hidden_size,
540 + ),
541 + "pack_factor":
542 + quant_ratio,
543 + "use_bitsandbytes_4bit":
544 + True,
545 + },
546 + )
547 + # down_proj (row parallel)
548 + w2_total_size = (hidden_size *
549 + intermediate_size_per_partition) // quant_ratio
550 + w2_qweight = torch.nn.Parameter(
551 + torch.empty(
552 + num_experts,
553 + w2_total_size,
554 + 1,
555 + dtype=torch.uint8,
556 + ),
557 + requires_grad=False,
558 + )
559 + set_weight_attrs(
560 + w2_qweight,
561 + {
562 + "num_experts":
563 + num_experts,
564 + "input_dim":
565 + intermediate_size_per_partition,
566 + "output_dim":
567 + hidden_size,
568 + "experts_shape": (
569 + num_experts,
570 + hidden_size,
571 + intermediate_size_per_partition,
572 + ),
573 + "pack_factor":
574 + quant_ratio,
575 + "use_bitsandbytes_4bit":
576 + True,
577 + },
578 + )
579 + layer.register_parameter("w2_weight", w2_qweight)
580 + set_weight_attrs(w2_qweight, extra_weight_attrs)
581 +
582 + def _create_weights_8bit(
583 + self,
584 + layer: torch.nn.Module,
585 + num_experts: int,
586 + hidden_size: int,
587 + intermediate_size_per_partition: int,
588 + params_dtype: torch.dtype,
589 + **extra_weight_attrs,
590 + ):
591 + raise NotImplementedError
592 +
593 + def _apply_4bit_dequnt(
594 + self, layer: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
595 + from bitsandbytes.functional import dequantize_4bit
596 + w13 = dequantize_4bit(
597 + layer.w13_weight.reshape(-1, 1),
598 + layer.w13_weight.bnb_quant_state,
599 + )
600 + w2 = dequantize_4bit(
601 + layer.w2_weight.reshape(-1, 1),
602 + layer.w2_weight.bnb_quant_state,
603 + )
604 + w13 = w13.reshape(layer.w13_weight.experts_shape)
605 + w2 = w2.reshape(layer.w2_weight.experts_shape)
606 + return w13, w2
607 +
608 + def _apply_8bit_dequant(
609 + self, layer: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
610 + raise NotImplementedError
vllm/model_executor/model_loader/bitsandbytes_loader.py
+193 -45
@@ -20,6 +20,7 @@ from vllm.distributed import (get_tensor_model_parallel_rank,
20 get_tensor_model_parallel_world_size)
21 # yapf: enable
22 from vllm.logger import init_logger
23 +from vllm.model_executor.layers.fused_moe import FusedMoE
24 from vllm.model_executor.layers.linear import (LinearBase,
25 MergedColumnParallelLinear,
26 QKVParallelLinear,
@@ -411,9 +412,33 @@ class BitsAndBytesModelLoader(BaseModelLoader):
412 # in case model has a mixture of disk-merged and disk-split
413 # weights with same last name.
414 self.target_modules.append(name)
415 + elif (isinstance(module, FusedMoE)
416 + and hasattr(module.quant_method, "quant_config")):
417 + if not hasattr(model, "get_expert_mapping"):
418 + raise AttributeError(
419 + f"MoE Model {type(model).__name__} does not support "
420 + "BitsAndBytes quantization yet. Ensure this model has "
421 + "'get_expert_mapping' method.")
422 + # TODO: support FusedMoE with prequant and 8bit.
423 + if self.pre_quant:
424 + raise ValueError(
425 + "Prequant BitsAndBytes models with FusedMoE is not "
426 + "supported yet.")
427 + if self.load_8bit:
428 + raise ValueError(
429 + "BitsAndBytes 8bit quantization with FusedMoE is not "
430 + "supported yet.")
431 + # Get the corresponding weight name using module name and
432 + # get_expert_mapping.
433 + expert_mapping = model.get_expert_mapping()
434 + for exp in expert_mapping:
435 + weight_name = exp[1]
436 + rep_name = name.replace("experts",
437 + "") + weight_name.removesuffix(".")
438 + self.target_modules.append(rep_name)
439
440 assert (self.target_modules
416 - ), "vllm currently does not support BNB quantization for"
441 + ), "vLLM currently does not support BNB quantization for"
442 f" {type(model).__name__}"
443
444 def _classify_module_sharding(self, model: nn.Module):
@@ -437,6 +462,14 @@ class BitsAndBytesModelLoader(BaseModelLoader):
462 # dimension (dim=-1)
463 elif isinstance(module, (RowParallelLinear, )):
464 self.column_sharded_weights_modules.append(name)
465 + elif isinstance(module, FusedMoE):
466 + expert_mapping = model.get_expert_mapping()
467 + for exp in expert_mapping:
468 + if exp[-1] == "w2":
469 + weight_name = exp[1]
470 + rep_name = name.replace(
471 + "experts", "") + weight_name.removesuffix(".")
472 + self.column_sharded_weights_modules.append(rep_name)
473
474 def _verify_model_compatibility(self, model: nn.Module,
475 model_config: ModelConfig) -> None:
@@ -490,34 +523,132 @@ class BitsAndBytesModelLoader(BaseModelLoader):
523 self._get_bnb_target_modules(model)
524 self._classify_module_sharding(model)
525
493 - def load_weights(self, model: nn.Module,
494 - model_config: ModelConfig) -> None:
526 + def _dequantize_dq(self, quant_states: Any):
527 + """
528 + When BNB employs Double Quantization, we perform the dequantization of
529 + these constants during weight loading rather than at inference time,
530 + thereby avoiding this computational overhead during inference. This
531 + comes at the cost of increased memory usage.
532 + """
533 + from bitsandbytes.functional import QuantState, dequantize_blockwise
534
496 - self._verify_model_compatibility(model, model_config)
497 - self._initialize_loader_state(model, model_config)
535 + def _dequantize_single_state(quant_state):
536 + """Helper function to dequantize a single QuantState object."""
537 + if not (isinstance(quant_state, QuantState)
538 + and quant_state.nested):
539 + return
540
499 - logger.info("Loading weights with BitsAndBytes quantization. "
500 - "May take a while ...")
501 - qweight_iterator, quant_state_dict = (
502 - self._get_quantized_weights_iterator(
503 - model_config.model,
504 - model_config.revision,
505 - ))
506 - weights_to_load = {name for name, _ in model.named_parameters()}
507 - loaded_weights = model.load_weights(qweight_iterator)
508 - # Some models may have weights loading tracker unimplemented.
509 - if loaded_weights is not None:
510 - weights_not_loaded = weights_to_load - loaded_weights
511 - if weights_not_loaded:
512 - raise ValueError("Following weights were not initialized from "
513 - f"checkpoint: {weights_not_loaded}")
541 + # Copied from: https://github.com/bitsandbytes-foundation/bitsandbytes/blob/0.45.3/bitsandbytes/functional.py#L1352-#L1356
542 + absmax = dequantize_blockwise(quant_state.absmax,
543 + quant_state.state2)
544 + absmax += quant_state.offset
545
515 - param_dict = dict(model.named_parameters())
546 + # Ensure float32 dtype
547 + if absmax.dtype != torch.float32:
548 + absmax = absmax.float()
549 +
550 + quant_state.absmax = absmax
551 + quant_state.nested = False
552 + quant_state.offset = None
553 + quant_state.state2 = None
554 +
555 + if isinstance(quant_states, dict):
556 + for quant_state in quant_states.values():
557 + _dequantize_single_state(quant_state)
558 + else:
559 + _dequantize_single_state(quant_states)
560 + return quant_states
561 +
562 + def _fuse_moe_quant_states(self, model: nn.Module,
563 + quant_states_dict: dict) -> dict:
564 + """
565 +
566 + This function consolidates individual expert quantization states into
567 + fused representations for w13 and w2.
568 + """
569 + from bitsandbytes.functional import QuantState
570 +
571 + if not hasattr(model, "get_expert_mapping"):
572 + return dict()
573 +
574 + expert_mapping = model.get_expert_mapping()
575 + expert_qs_dict = {}
576 + for name, module in model.named_modules():
577 + if not isinstance(module, FusedMoE):
578 + continue
579 + w1_states_lst = []
580 + w2_states_lst = []
581 + w3_states_lst = []
582 + for exp in expert_mapping:
583 + shard_id = exp[-1]
584 + if shard_id not in ("w1", "w2", "w3"):
585 + raise ValueError(f"shard_id must be ['w1','w2','w3'] but "
586 + f"got {shard_id}.")
587 + layer_prefix = name.split("experts")[0]
588 + weight_qual_name = layer_prefix + exp[1] + "weight"
589 + quant_state = self._dequantize_dq(
590 + quant_states_dict[weight_qual_name])
591 + if shard_id == "w1":
592 + w1_states_lst.append(quant_state)
593 + elif shard_id == "w2":
594 + w2_states_lst.append(quant_state)
595 + else:
596 + w3_states_lst.append(quant_state)
597 + del quant_states_dict[weight_qual_name]
598 + assert (len(w1_states_lst) == len(w2_states_lst) ==
599 + len(w3_states_lst))
600 + w13_absmax_lst = []
601 + w2_absmax_lst = []
602 + w13_total_dim0 = 0
603 + w2_total_dim0 = 0
604 + for w1_qs, w2_qs, w3_qs in zip(w1_states_lst, w2_states_lst,
605 + w3_states_lst):
606 + assert w1_qs.shape == w3_qs.shape
607 + assert w1_qs.blocksize == w2_qs.blocksize == w3_qs.blocksize
608 + assert w1_qs.dtype == w2_qs.dtype == w3_qs.dtype
609 + # w1 and w3 are interleaved in storage
610 + w13_absmax_lst.append(w1_qs.absmax)
611 + w13_absmax_lst.append(w3_qs.absmax)
612 + w2_absmax_lst.append(w2_qs.absmax)
613 + w13_total_dim0 += w1_qs.shape[0] + w3_qs.shape[0]
614 + w2_total_dim0 += w2_qs.shape[0]
615 +
616 + w13_absmax = torch.cat(w13_absmax_lst)
617 + w2_absmax = torch.cat(w2_absmax_lst)
618 + # Create fused quantization state for w13.
619 + w13_qs = QuantState(
620 + absmax=w13_absmax,
621 + shape=(w13_total_dim0, w1_states_lst[0].shape[1]),
622 + code=w1_states_lst[0].code,
623 + blocksize=w1_states_lst[0].blocksize,
624 + quant_type="nf4",
625 + dtype=w1_states_lst[0].dtype,
626 + )
627 + # Create fused quantization state for w2.
628 + w2_qs = QuantState(
629 + absmax=w2_absmax,
630 + shape=(w2_total_dim0, w2_states_lst[0].shape[1]),
631 + code=w2_states_lst[0].code,
632 + blocksize=w2_states_lst[0].blocksize,
633 + quant_type="nf4",
634 + dtype=w2_states_lst[0].dtype,
635 + )
636 + # The weight suffixes .w13_weight and .w2_weight are consistent
637 + # with the param in BitsAndBytesMoEMethod.
638 + w13_weight_name = name + ".w13_weight"
639 + w2_weight_name = name + ".w2_weight"
640 + expert_qs_dict[w13_weight_name] = w13_qs
641 + expert_qs_dict[w2_weight_name] = w2_qs
642 + return expert_qs_dict
643 +
644 + def _stack_quantization_states(
645 + self, model: nn.Module,
646 + quant_state_dict: dict) -> dict[str, dict[int, Any]]:
647 stacked_quant_state_dict: dict[str, dict[int, Any]] = {}
648 # TODO: Change this lazy import to normal import
649 # after the checks are updated to run on a new version
650 from vllm.model_executor.models.utils import is_pp_missing_parameter
520 -
651 + param_dict = dict(model.named_parameters())
652 for quant_param_name in quant_state_dict:
653 if is_pp_missing_parameter(quant_param_name, model):
654 continue
@@ -558,14 +689,20 @@ class BitsAndBytesModelLoader(BaseModelLoader):
689
690 stacked_quant_state_dict[quant_param_name][shard_index] = (
691 quant_state_dict[non_stacked_param_name])
692 + return stacked_quant_state_dict
693
694 + def _bind_quant_states_to_params(self, model: nn.Module,
695 + stacked_quant_state_dict: dict) -> None:
696 # save quant_states and offsets as the attributes of the parameters
697 + param_dict = dict(model.named_parameters())
698 for param_name, param in param_dict.items():
699 if param_name in stacked_quant_state_dict:
700 quant_states = stacked_quant_state_dict[param_name]
701 # Dequantize double quantized values during weight loading.
567 - dequantize_dq(quant_states)
702 + self._dequantize_dq(quant_states)
703 set_weight_attrs(param, {"bnb_quant_state": quant_states})
704 + if not isinstance(quant_states, dict):
705 + continue
706
707 pack_ratio = getattr(param, "pack_factor", -1)
708 if pack_ratio == -1:
@@ -585,29 +722,40 @@ class BitsAndBytesModelLoader(BaseModelLoader):
722 if self.load_8bit:
723 set_weight_attrs(
724 param, {"matmul_state": [None] * len(quant_states)})
725 +
726 + def load_weights(self, model: nn.Module,
727 + model_config: ModelConfig) -> None:
728 +
729 + self._verify_model_compatibility(model, model_config)
730 + self._initialize_loader_state(model, model_config)
731 +
732 + logger.info("Loading weights with BitsAndBytes quantization. "
733 + "May take a while ...")
734 + qweight_iterator, quant_state_dict = (
735 + self._get_quantized_weights_iterator(
736 + model_config.model,
737 + model_config.revision,
738 + ))
739 + weights_to_load = {name for name, _ in model.named_parameters()}
740 + loaded_weights = model.load_weights(qweight_iterator)
741 + # Some models may have weights loading tracker unimplemented.
742 + if loaded_weights is not None:
743 + weights_not_loaded = weights_to_load - loaded_weights
744 + if weights_not_loaded:
745 + raise ValueError("Following weights were not initialized from "
746 + f"checkpoint: {weights_not_loaded}")
747 + expert_quant_state_dict = self._fuse_moe_quant_states(
748 + model, quant_state_dict)
749 +
750 + stacked_quant_state_dict = self._stack_quantization_states(
751 + model, quant_state_dict)
752 +
753 + stacked_quant_state_dict = {
754 + **expert_quant_state_dict,
755 + **stacked_quant_state_dict
756 + }
757 + self._bind_quant_states_to_params(model, stacked_quant_state_dict)
758 torch.cuda.empty_cache()
759
760 def download_model(self, model_config: ModelConfig) -> None:
761 self._prepare_weights(model_config.model, model_config.revision)
592 -
593 -
594 -def dequantize_dq(quant_states: dict) -> None:
595 - """
596 - When BNB employs Double Quantization, we perform the dequantization of
597 - these constants during weight loading rather than at inference time,
598 - thereby avoiding this computational overhead during inference. This comes
599 - at the cost of increased memory usage.
600 - """
601 - from bitsandbytes.functional import QuantState, dequantize_blockwise
602 - for _, quant_state in quant_states.items():
603 - # Copied from: https://github.com/bitsandbytes-foundation/bitsandbytes/blob/0.45.3/bitsandbytes/functional.py#L1352-#L1356
604 - if isinstance(quant_state, QuantState) and quant_state.nested:
605 - absmax = dequantize_blockwise(quant_state.absmax,
606 - quant_state.state2)
607 - absmax += quant_state.offset
608 - if absmax.dtype != torch.float32:
609 - absmax = absmax.float()
610 - quant_state.absmax = absmax
611 - quant_state.nested = False
612 - quant_state.offset = None
613 - quant_state.state2 = None
vllm/model_executor/models/olmoe.py
+24 -9
@@ -330,6 +330,15 @@ class OlmoeModel(nn.Module):
330 hidden_states, _ = self.norm(hidden_states, residual)
331 return hidden_states
332
333 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
334 + # Params for weights, fp8 weight scales, fp8 activation scales
335 + # (param_name, weight_name, expert_id, shard_id)
336 + return FusedMoE.make_expert_params_mapping(
337 + ckpt_gate_proj_name="gate_proj",
338 + ckpt_down_proj_name="down_proj",
339 + ckpt_up_proj_name="up_proj",
340 + num_experts=self.config.num_experts)
341 +
342 def load_weights(self, weights: Iterable[tuple[str,
343 torch.Tensor]]) -> set[str]:
344 stacked_params_mapping = [
@@ -341,14 +350,6 @@ class OlmoeModel(nn.Module):
350 ("gate_up_proj", "up_proj", 1),
351 ]
352
344 - # Params for weights, fp8 weight scales, fp8 activation scales
345 - # (param_name, weight_name, expert_id, shard_id)
346 - expert_params_mapping = FusedMoE.make_expert_params_mapping(
347 - ckpt_gate_proj_name="gate_proj",
348 - ckpt_down_proj_name="down_proj",
349 - ckpt_up_proj_name="up_proj",
350 - num_experts=self.config.num_experts)
351 -
353 params_dict = dict(self.named_parameters())
354 loaded_params: set[str] = set()
355 for name, loaded_weight in weights:
@@ -379,7 +380,7 @@ class OlmoeModel(nn.Module):
380 weight_loader(param, loaded_weight, shard_id)
381 break
382 else:
382 - for mapping in expert_params_mapping:
383 + for mapping in self.get_expert_mapping():
384 param_name, weight_name, expert_id, shard_id = mapping
385 if weight_name not in name:
386 continue
@@ -425,6 +426,17 @@ class OlmoeModel(nn.Module):
426
427
428 class OlmoeForCausalLM(nn.Module, SupportsPP):
429 + packed_modules_mapping = {
430 + "qkv_proj": [
431 + "q_proj",
432 + "k_proj",
433 + "v_proj",
434 + ],
435 + "gate_up_proj": [
436 + "gate_proj",
437 + "up_proj",
438 + ],
439 + }
440
441 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
442 super().__init__()
@@ -466,3 +478,6 @@ class OlmoeForCausalLM(nn.Module, SupportsPP):
478 torch.Tensor]]) -> set[str]:
479 loader = AutoWeightsLoader(self)
480 return loader.load_weights(weights)
481 +
482 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
483 + return self.model.get_expert_mapping()
vllm/model_executor/models/phimoe.py
+11
@@ -516,6 +516,14 @@ class PhiMoEModel(nn.Module):
516 hidden_states = self.norm(hidden_states)
517 return hidden_states
518
519 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
520 + return FusedMoE.make_expert_params_mapping(
521 + ckpt_gate_proj_name="w1",
522 + ckpt_down_proj_name="w2",
523 + ckpt_up_proj_name="w3",
524 + num_experts=self.config.num_local_experts,
525 + )
526 +
527 def load_weights(self, weights: Iterable[tuple[str,
528 torch.Tensor]]) -> set[str]:
529 stacked_params_mapping = [
@@ -672,3 +680,6 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
680 torch.Tensor]]) -> set[str]:
681 loader = AutoWeightsLoader(self)
682 return loader.load_weights(weights)
683 +
684 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
685 + return self.model.get_expert_mapping()
vllm/model_executor/models/qwen2_moe.py
+26 -9
@@ -391,6 +391,15 @@ class Qwen2MoeModel(nn.Module):
391 hidden_states, _ = self.norm(hidden_states, residual)
392 return hidden_states
393
394 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
395 + # Params for weights, fp8 weight scales, fp8 activation scales
396 + # (param_name, weight_name, expert_id, shard_id)
397 + return FusedMoE.make_expert_params_mapping(
398 + ckpt_gate_proj_name="gate_proj",
399 + ckpt_down_proj_name="down_proj",
400 + ckpt_up_proj_name="up_proj",
401 + num_experts=self.config.num_experts)
402 +
403 def load_weights(self, weights: Iterable[tuple[str,
404 torch.Tensor]]) -> set[str]:
405 stacked_params_mapping = [
@@ -402,14 +411,6 @@ class Qwen2MoeModel(nn.Module):
411 ("gate_up_proj", "up_proj", 1),
412 ]
413
405 - # Params for weights, fp8 weight scales, fp8 activation scales
406 - # (param_name, weight_name, expert_id, shard_id)
407 - expert_params_mapping = FusedMoE.make_expert_params_mapping(
408 - ckpt_gate_proj_name="gate_proj",
409 - ckpt_down_proj_name="down_proj",
410 - ckpt_up_proj_name="up_proj",
411 - num_experts=self.config.num_experts)
412 -
414 params_dict = dict(self.named_parameters())
415 loaded_params: set[str] = set()
416 for name, loaded_weight in weights:
@@ -441,11 +442,13 @@ class Qwen2MoeModel(nn.Module):
442 weight_loader(param, loaded_weight, shard_id)
443 break
444 else:
444 - for mapping in expert_params_mapping:
445 + for mapping in self.get_expert_mapping():
446 param_name, weight_name, expert_id, shard_id = mapping
447 if weight_name not in name:
448 continue
449 name = name.replace(weight_name, param_name)
450 + if "layers.13.mlp.experts.w2_weight" in name:
451 + pass
452 # Skip layers on other devices.
453 if is_pp_missing_parameter(name, self):
454 continue
@@ -493,6 +496,17 @@ class Qwen2MoeModel(nn.Module):
496 class Qwen2MoeForCausalLM(nn.Module, SupportsPP):
497
498 fall_back_to_pt_during_load = False
499 + packed_modules_mapping = {
500 + "qkv_proj": [
501 + "q_proj",
502 + "k_proj",
503 + "v_proj",
504 + ],
505 + "gate_up_proj": [
506 + "gate_proj",
507 + "up_proj",
508 + ],
509 + }
510
511 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
512 super().__init__()
@@ -538,3 +552,6 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP):
552 torch.Tensor]]) -> set[str]:
553 loader = AutoWeightsLoader(self)
554 return loader.load_weights(weights)
555 +
556 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
557 + return self.model.get_expert_mapping()
vllm/model_executor/models/qwen3_moe.py
+13 -6
@@ -375,6 +375,15 @@ class Qwen3MoeModel(nn.Module):
375 hidden_states, _ = self.norm(hidden_states, residual)
376 return hidden_states
377
378 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
379 + # Params for weights, fp8 weight scales, fp8 activation scales
380 + # (param_name, weight_name, expert_id, shard_id)
381 + return FusedMoE.make_expert_params_mapping(
382 + ckpt_gate_proj_name="gate_proj",
383 + ckpt_down_proj_name="down_proj",
384 + ckpt_up_proj_name="up_proj",
385 + num_experts=self.config.num_experts)
386 +
387 def load_weights(self, weights: Iterable[tuple[str,
388 torch.Tensor]]) -> set[str]:
389 stacked_params_mapping = [
@@ -393,12 +402,7 @@ class Qwen3MoeModel(nn.Module):
402
403 # Params for weights, fp8 weight scales, fp8 activation scales
404 # (param_name, weight_name, expert_id, shard_id)
396 - expert_params_mapping = FusedMoE.make_expert_params_mapping(
397 - ckpt_gate_proj_name="gate_proj",
398 - ckpt_down_proj_name="down_proj",
399 - ckpt_up_proj_name="up_proj",
400 - num_experts=self.config.num_experts)
401 -
405 + expert_params_mapping = self.get_expert_mapping()
406 params_dict = dict(self.named_parameters())
407 loaded_params: set[str] = set()
408 for name, loaded_weight in weights:
@@ -539,3 +543,6 @@ class Qwen3MoeForCausalLM(nn.Module, SupportsPP):
543 torch.Tensor]]) -> set[str]:
544 loader = AutoWeightsLoader(self)
545 return loader.load_weights(weights)
546 +
547 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
548 + return self.model.get_expert_mapping()