@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 08abfa78ec

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 14, 05:30 UTC

[Bugfix] fix modelopt exclude_modules name mapping (#24178)

Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

tomeras91 committed Sep 10, 2025 at 20:20 UTC 08abfa78ec2fca59889045bb5aec41e383dac9bd

3 files changed +59 -38

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -291,6 +291,7 @@ class MambaMixer2(MambaBase, CustomOp):

291 output_size=self.conv_dim,

292 bias=use_conv_bias,

293 quant_config=None,

294 + prefix=f"{prefix}.conv1d",

295 )

296 # unsqueeze to fit conv1d weights shape into the linear weights shape.

297 # Can't do this in `weight_loader` since it already exists in

@@ -303,6 +304,7 @@ class MambaMixer2(MambaBase, CustomOp):

304 output_size=intermediate_size + self.conv_dim + self.num_heads,

305 bias=use_bias,

306 quant_config=quant_config,

307 + prefix=f"{prefix}.in_proj",

308 )

309

310 # - because in_proj is a concatenation of 3 weights, we

@@ -402,6 +404,7 @@ class MambaMixer2(MambaBase, CustomOp):

404 bias=use_bias,

405 input_is_parallel=True,

406 quant_config=quant_config,

407 + prefix=f"{prefix}.out_proj",

408 )

409

410 self.norm = Mixer2RMSNormGated(intermediate_size,

vllm/model_executor/layers/quantization/modelopt.py

+20 -4

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from typing import Any, Callable, Optional, Union

4 +from typing import TYPE_CHECKING, Any, Callable, Optional, Union

6 import torch

7 from torch.nn import Module

@@ -45,6 +45,9 @@ from vllm.utils import next_power_of_2

45 from vllm.utils.flashinfer import (flashinfer_scaled_fp4_mm, has_flashinfer,

46 has_flashinfer_moe)

48 +if TYPE_CHECKING:

49 + from vllm.model_executor.models.utils import WeightsMapper

50 +

51 logger = init_logger(__name__)

53 QUANT_ALGOS = ["FP8", "NVFP4"]

@@ -63,7 +66,7 @@ class ModelOptFp8Config(QuantizationConfig):

66 super().__init__()

67 self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized

68 self.kv_cache_quant_method = kv_cache_quant_method

66 - self.exclude_modules = exclude_modules

69 + self.exclude_modules = exclude_modules or []

70 if is_checkpoint_fp8_serialized:

71 logger.warning("Detected ModelOpt fp8 checkpoint. Please note that"

72 " the format is experimental and could change.")

@@ -84,6 +87,11 @@ class ModelOptFp8Config(QuantizationConfig):

87 def get_config_filenames(cls) -> list[str]:

88 return ["hf_quant_config.json"]

90 + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):

91 + if self.exclude_modules is not None:

92 + self.exclude_modules = hf_to_vllm_mapper.apply_list(

93 + self.exclude_modules)

94 +

95 @classmethod

96 def override_quantization_method(

97 cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:

@@ -170,7 +178,9 @@ class ModelOptFp8Config(QuantizationConfig):

178 prefix: str) -> Optional["QuantizeMethodBase"]:

179 from vllm.attention.layer import Attention # Avoid circular import

180 if isinstance(layer, LinearBase):

173 - if self.is_layer_excluded(prefix):

181 + if (is_layer_skipped(prefix, self.exclude_modules,

182 + self.packed_modules_mapping)

183 + or self.is_layer_excluded(prefix)):

184 return UnquantizedLinearMethod()

185 return ModelOptFp8LinearMethod(self)

186 elif isinstance(layer, Attention):

@@ -615,6 +625,11 @@ class ModelOptNvFp4Config(QuantizationConfig):

625 def get_config_filenames(cls) -> list[str]:

626 return ["hf_quant_config.json"]

627

628 + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):

629 + if self.exclude_modules is not None:

630 + self.exclude_modules = hf_to_vllm_mapper.apply_list(

631 + self.exclude_modules)

632 +

633 @classmethod

634 def override_quantization_method(

635 cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:

@@ -763,7 +778,8 @@ class ModelOptNvFp4Config(QuantizationConfig):

778 prefix: str) -> Optional["QuantizeMethodBase"]:

779 from vllm.attention.layer import Attention # Avoid circular import

780 if isinstance(layer, LinearBase):

766 - if (is_layer_skipped(prefix, self.exclude_modules)

781 + if (is_layer_skipped(prefix, self.exclude_modules,

782 + self.packed_modules_mapping)

783 or self.is_layer_excluded(prefix, self.exclude_modules)):

784 return UnquantizedLinearMethod()

785 return ModelOptNvFp4LinearMethod(self)

vllm/model_executor/models/nemotron_h.py

+36 -34

@@ -44,15 +44,16 @@ from vllm.model_executor.layers.mamba.mamba_utils import (

44 from vllm.model_executor.layers.quantization import QuantizationConfig

45 from vllm.model_executor.layers.vocab_parallel_embedding import (

46 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

47 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

47 +from vllm.model_executor.model_loader.weight_utils import (

48 + default_weight_loader, maybe_remap_kv_scale_name)

49 from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,

50 SupportsLoRA, SupportsPP,

51 SupportsQuant)

52 from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

53 MambaCacheParams)

54 from vllm.model_executor.models.utils import (

54 - AutoWeightsLoader, make_empty_intermediate_tensors_factory, make_layers,

55 - maybe_prefix)

55 + AutoWeightsLoader, WeightsMapper, make_empty_intermediate_tensors_factory,

56 + make_layers, maybe_prefix)

57 from vllm.model_executor.sampling_metadata import SamplingMetadata

58 from vllm.sequence import IntermediateTensors

59 from vllm.transformers_utils.configs import NemotronHConfig

@@ -426,38 +427,36 @@ class NemotronHModel(nn.Module):

427

428 def load_weights(self, weights: Iterable[tuple[str,

429 torch.Tensor]]) -> set[str]:

429 - attb_params_mapping = {

430 - "q_proj": "q",

431 - "k_proj": "k",

432 - "v_proj": "v",

433 - }

430 + stacked_params_mapping = [

431 + # (param_name, shard_name, shard_id)

432 + ("qkv_proj", "q_proj", "q"),

433 + ("qkv_proj", "k_proj", "k"),

434 + ("qkv_proj", "v_proj", "v"),

435 + ]

436

437 params_dict = dict(self.named_parameters())

438 loaded_params: set[str] = set()

439 for name, loaded_weight in weights:

438 - if "embeddings" in name:

439 - name = name.replace("embeddings", "embed_tokens")

440 + if "scale" in name:

441 + # Remapping the name of FP8 kv-scale.

442 + name = maybe_remap_kv_scale_name(name, params_dict)

443 + if name is None:

444 + continue

445 +

446 + # load stacked params

447 + for param_name, weight_name, shard_id in stacked_params_mapping:

448 + if weight_name not in name:

449 + continue

450 + name = name.replace(weight_name, param_name)

451 + # Skip loading extra bias for GPTQ models.

452 + if name.endswith(".bias") and name not in params_dict:

453 + continue

454

441 - if "A_log" in name:

442 - name = name.replace("A_log", "A")

443 - loaded_weight = loaded_weight.to(torch.float32)

444 -

445 - if "D" in name:

446 - loaded_weight = loaded_weight.to(torch.float32)

447 -

448 - if "dt_bias" in name:

449 - loaded_weight = loaded_weight.to(torch.float32)

450 -

451 - # load attn params

452 - if any(proj in name for proj in ["q_proj", "k_proj", "v_proj"]):

453 - weight_name = next(proj

454 - for proj in ["q_proj", "k_proj", "v_proj"]

455 - if proj in name)

456 - name = name.replace(weight_name, "qkv_proj")

455 param = params_dict[name]

456 weight_loader = param.weight_loader

459 - weight_loader(param, loaded_weight,

460 - attb_params_mapping[weight_name])

457 + weight_loader(param, loaded_weight, shard_id)

458 + break

459 +

460 # load other params

461 else:

462 param = params_dict[name]

@@ -471,6 +470,14 @@ class NemotronHModel(nn.Module):

470

471 class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

472 IsHybrid, SupportsQuant):

473 + hf_to_vllm_mapper = WeightsMapper(

474 + orig_to_new_prefix={"backbone": "model"},

475 + orig_to_new_substr={

476 + "A_log": "A",

477 + "embeddings": "embed_tokens"

478 + },

479 + )

480 +

481 packed_modules_mapping = {

482 "qkv_proj": [

483 "q_proj",

@@ -622,10 +629,5 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

629

630 def load_weights(self, weights: Iterable[tuple[str,

631 torch.Tensor]]) -> set[str]:

625 - # update name in weights before passing to loader

626 - updated_weights = []

627 - for name, loaded_weight in weights:

628 - name = name.replace("backbone", "model")

629 - updated_weights.append((name, loaded_weight))

632 loader = AutoWeightsLoader(self)

631 - return loader.load_weights(updated_weights)

633 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)