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

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 2, 06:00 UTC

[Bugfix] Convert Gemma4-MM ViT linear layers to vllm native impl (#43798)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: ZiTian Zhao <zitian.zhao@tencentmusic.com> Co-authored-by: B-201 <Joy25810@foxmail.com>

Isotr0py committed Jun 2, 2026 at 12:41 UTC f91fb2fcf3f100d46baf4db2e5d86fa5596b276b

4 files changed +111 -11

vllm/model_executor/layers/quantization/bitsandbytes.py

+8 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

4 +from functools import cached_property

5 from typing import Any, Union

7 import torch

@@ -168,6 +169,12 @@ class BitsAndBytesConfig(QuantizationConfig):

169 return None

170

171

172 +class BitsAndBytesWeightParameter(torch.nn.Parameter):

173 + @cached_property

174 + def dtype(self) -> torch.dtype:

175 + return torch.get_default_dtype()

176 +

177 +

178 def is_layer_skipped_bnb(prefix: str, llm_int8_skip_modules: list[str]):

179 # Split the prefix into its dot-separated components

180 components = prefix.split(".")

@@ -246,7 +253,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):

253 "The input size is not aligned with the quantized weight shape."

254 )

255

249 - qweight = torch.nn.Parameter(

256 + qweight = BitsAndBytesWeightParameter(

257 torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),

258 requires_grad=False,

259 )

vllm/model_executor/model_loader/bitsandbytes_loader.py

+23

@@ -744,6 +744,29 @@ class BitsAndBytesModelLoader(BaseModelLoader):

744 stacked_quant_state_dict[quant_param_name][shard_index] = quant_state_dict[

745 non_stacked_param_name

746 ]

747 +

748 + # repeat k_proj for v_proj for k_eq_v models (e.g. Gemma4)

749 + config = getattr(model, "config", None)

750 + if config is not None:

751 + text_config = config.get_text_config()

752 + if getattr(text_config, "attention_k_eq_v", False):

753 + shard_packed = {

754 + name

755 + for name, subs in self.modules_mapping.packed_mapping.items()

756 + if len(subs) == 3

757 + }

758 + for param_name, shards in stacked_quant_state_dict.items():

759 + is_target = (

760 + isinstance(shards, dict)

761 + and len(shards) == 2

762 + and any(

763 + param_name.endswith(f"{p}.weight") for p in shard_packed

764 + )

765 + )

766 + if is_target:

767 + assert 1 in shards and 2 not in shards

768 + shards[2] = shards[1]

769 +

770 return stacked_quant_state_dict

771

772 def _bind_quant_states_to_params(

vllm/model_executor/models/gemma4_mm.py

+51 -10

@@ -16,7 +16,7 @@ reason about temporal order.

17 import math

18 from collections.abc import Iterable, Mapping, Sequence

19 -from typing import Annotated, Any, Literal

19 +from typing import TYPE_CHECKING, Annotated, Any, Literal

21 import numpy as np

22 import torch

@@ -41,6 +41,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm

41 from vllm.model_executor.layers.linear import ReplicatedLinear

42 from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM

43 from vllm.model_executor.models.module_mapping import MultiModelKeys

44 +from vllm.model_executor.models.transformers.utils import recursive_replace_linear

45 from vllm.multimodal import MULTIMODAL_REGISTRY

46 from vllm.multimodal.inputs import (

47 MultiModalFieldConfig,

@@ -71,6 +72,7 @@ from .interfaces import (

72 SupportsLoRA,

73 SupportsMultiModal,

74 SupportsPP,

75 + SupportsQuant,

76 )

77 from .utils import (

78 AutoWeightsLoader,

@@ -79,6 +81,9 @@ from .utils import (

81 maybe_prefix,

82 )

84 +if TYPE_CHECKING:

85 + from vllm.model_executor.layers.quantization import QuantizationConfig

86 +

87 logger = init_logger(__name__)

89 # Video constants — match transformers Gemma4VideoProcessor defaults.

@@ -872,6 +877,9 @@ class Gemma4MultimodalEmbedder(nn.Module):

877 self,

878 multimodal_config: Gemma4VisionConfig | Gemma4AudioConfig,

879 text_config: Gemma4TextConfig,

880 + *,

881 + quant_config: "QuantizationConfig | None" = None,

882 + prefix: str = "",

883 ):

884 super().__init__()

885

@@ -895,6 +903,8 @@ class Gemma4MultimodalEmbedder(nn.Module):

903 embedding_dim,

904 self.text_hidden_size,

905 bias=False,

906 + quant_config=quant_config,

907 + prefix=maybe_prefix(prefix, "embedding_projection"),

908 )

909

910 def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:

@@ -917,6 +927,7 @@ class Gemma4MultimodalEmbedder(nn.Module):

927 class Gemma4ForConditionalGeneration(

928 nn.Module,

929 SupportsMultiModal,

930 + SupportsQuant,

931 SupportsPP,

932 SupportsLoRA,

933 SupportsEagle3,

@@ -936,11 +947,14 @@ class Gemma4ForConditionalGeneration(

947 # Maps checkpoint prefixes to vLLM module paths.

948 hf_to_vllm_mapper = WeightsMapper(

949 orig_to_new_prefix={

950 + # vision tower

951 + "model.vision_tower": "vision_tower",

952 + "model.embed_vision": "embed_vision",

953 + # audio tower

954 + "model.audio_tower.": "audio_tower.",

955 "model.embed_audio.": "embed_audio.",

940 - "model.embed_vision.": "embed_vision.",

956 + # backbone

957 "model.language_model.": "language_model.model.",

942 - "model.vision_tower.": "vision_tower.",

943 - "model.audio_tower.": "audio_tower.",

958 "lm_head.": "language_model.lm_head.",

959 "model": "language_model.model",

960 }

@@ -959,7 +973,15 @@ class Gemma4ForConditionalGeneration(

973 with self._mark_tower_model(vllm_config, {"image", "video"}):

974 self.vision_tower = AutoModel.from_config(config=config.vision_config)

975 self.embed_vision = Gemma4MultimodalEmbedder(

962 - config.vision_config, config.text_config

976 + config.vision_config,

977 + config.text_config,

978 + quant_config=quant_config,

979 + prefix=maybe_prefix(prefix, "embed_vision"),

980 + )

981 + recursive_replace_linear(

982 + self.vision_tower,

983 + quant_config,

984 + prefix=maybe_prefix(prefix, "vision_tower"),

985 )

986

987 # ---- Audio tower (variants with audio_config) ----

@@ -972,7 +994,15 @@ class Gemma4ForConditionalGeneration(

994 # position embeddings, softcap, gradient_clipping).

995 self.audio_tower.post_init()

996 self.embed_audio = Gemma4MultimodalEmbedder(

975 - config.audio_config, config.text_config

997 + config.audio_config,

998 + config.text_config,

999 + quant_config=quant_config,

1000 + prefix=maybe_prefix(prefix, "embed_audio"),

1001 + )

1002 + recursive_replace_linear(

1003 + self.audio_tower,

1004 + quant_config,

1005 + prefix=maybe_prefix(prefix, "audio_tower"),

1006 )

1007 else:

1008 self.audio_tower = None

@@ -1153,6 +1183,7 @@ class Gemma4ForConditionalGeneration(

1183 vt = self.vision_tower

1184 vision_cfg = self.config.vision_config

1185 pooling_k2 = vision_cfg.pooling_kernel_size**2

1186 + target_dtype = self.language_model.model.embed_tokens.weight.dtype

1187

1188 # Concurrent requests with different image resolutions may

1189 # arrive as a list of per-image tensors, while same-resolution

@@ -1193,7 +1224,11 @@ class Gemma4ForConditionalGeneration(

1224 )

1225 pad_tensor = (pp_tensor == -1).all(dim=-1)

1226

1196 - inputs_embeds = vt.patch_embedder(pv_tensor, pp_tensor, pad_tensor)

1227 + inputs_embeds = vt.patch_embedder(

1228 + pv_tensor,

1229 + pp_tensor,

1230 + pad_tensor,

1231 + ).to(target_dtype)

1232 encoder_outputs = vt.encoder(

1233 inputs_embeds=inputs_embeds,

1234 attention_mask=~pad_tensor,

@@ -1230,7 +1265,9 @@ class Gemma4ForConditionalGeneration(

1265 all_valid_states[orig_idx] = valid_states

1266 valid_lens[orig_idx] = valid_states.shape[0]

1267

1233 - target_dtype = self.embed_vision.embedding_projection.weight.dtype

1268 + # Use embed_tokens dtype as compute dtype; embedding_projection.weight

1269 + # may be uint8 under BnB 4-bit, which would corrupt the cast.

1270 + target_dtype = self.language_model.model.embed_tokens.weight.dtype

1271

1272 # Project all images in a single batched call.

1273 flat_valid_states = torch.cat(all_valid_states, dim=0).to(target_dtype)

@@ -1273,7 +1310,7 @@ class Gemma4ForConditionalGeneration(

1310 vt = self.vision_tower

1311 vision_cfg = self.config.vision_config

1312 pooling_k2 = vision_cfg.pooling_kernel_size**2

1276 - target_dtype = self.embed_vision.embedding_projection.weight.dtype

1313 + target_dtype = self.language_model.model.embed_tokens.weight.dtype

1314

1315 if isinstance(frame_counts, torch.Tensor):

1316 fc_list = frame_counts.tolist()

@@ -1301,7 +1338,11 @@ class Gemma4ForConditionalGeneration(

1338 pp_chunk = pixel_position_ids[i : i + max_batch_size]

1339 pad_chunk = padding_positions[i : i + max_batch_size]

1340

1304 - inputs_embeds = vt.patch_embedder(pv_chunk, pp_chunk, pad_chunk)

1341 + inputs_embeds = vt.patch_embedder(

1342 + pv_chunk,

1343 + pp_chunk,

1344 + pad_chunk,

1345 + ).to(target_dtype)

1346 encoder_outputs = vt.encoder(

1347 inputs_embeds=inputs_embeds,

1348 attention_mask=~pad_chunk,

vllm/model_executor/models/transformers/utils.py

+29

@@ -32,6 +32,7 @@ from vllm.model_executor.layers.linear import (

32 ReplicatedLinear,

33 RowParallelLinear,

34 )

35 +from vllm.model_executor.models.utils import maybe_prefix

36 from vllm.transformers_utils.config import is_rope_parameters_nested

38 if TYPE_CHECKING:

@@ -227,6 +228,34 @@ def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm:

228 return RMSNorm(**kwargs)

229

230

231 +def recursive_replace_linear(

232 + model: nn.Module,

233 + quant_config: "QuantizationConfig | None",

234 + prefix: str = "",

235 +):

236 + """Recursively replace linear modules in the model as needed."""

237 +

238 + def _recursive_replace(module: nn.Module, prefix: str):

239 + for child_name, child_module in module.named_children():

240 + new_module = child_module

241 + qual_name = maybe_prefix(prefix, child_name)

242 + # Replace modules as needed

243 + if isinstance(child_module, nn.Linear):

244 + style = "replicate"

245 + new_module = replace_linear_class(

246 + child_module,

247 + style,

248 + quant_config,

249 + prefix=qual_name,

250 + )

251 + else:

252 + _recursive_replace(child_module, prefix=qual_name)

253 + if new_module is not child_module:

254 + setattr(module, child_name, new_module)

255 +

256 + _recursive_replace(model, prefix=prefix)

257 +

258 +

259 def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module):

260 logger.debug("%s: %s -> %s", name, old_module, new_module)

261