From 38648e99a973f2d7563b7619cf1161a731f576f6 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Fri, 1 Aug 2025 23:48:29 +0000 Subject: [PATCH] [None][feat] Switch to internal version of MMProjector in Gemma3 Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../_torch/models/modeling_gemma3vl.py | 79 +++++++++++++++---- 1 file changed, 62 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py index b2f30ba63e6f..699e464151d7 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py @@ -5,8 +5,6 @@ import torch from transformers import AutoProcessor, Gemma3Config, PreTrainedModel -from transformers.modeling_utils import no_init_weights -from transformers.models.gemma3.modeling_gemma3 import Gemma3MultiModalProjector from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ BaseWeightMapper @@ -18,6 +16,8 @@ from ...sampling_params import SamplingParams from ..attention_backend import AttentionMetadata from ..model_config import ModelConfig +from ..modules.linear import Linear +from ..modules.rms_norm import RMSNorm from .modeling_gemma3 import Gemma3ForCausalLM from .modeling_multimodal_utils import fuse_input_embeds from .modeling_siglip import SiglipVisionModel @@ -81,6 +81,61 @@ def __call__( return input_ids[0].to(torch.int32).tolist(), multimodal_data +# Original HF implementation: +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma3/modeling_gemma3.py#L684 +class Gemma3MultiModalProjector(torch.nn.Module): + """Gemma3MultiModalProjector using TRTLLM's Linear and RMSNorm.""" + + def __init__(self, model_config: ModelConfig[Gemma3Config]): + super().__init__() + config = model_config.pretrained_config + self.dtype = config.torch_dtype + self.mm_input_projection = Linear( + in_features=config.vision_config.hidden_size, + out_features=config.text_config.hidden_size, + bias=False, + dtype=self.dtype, + mapping=model_config.mapping) + self.mm_soft_emb_norm = RMSNorm( + hidden_size=config.vision_config.hidden_size, + eps=config.vision_config.layer_norm_eps, + dtype=self.dtype) + self.patches_per_image = int(config.vision_config.image_size // + config.vision_config.patch_size) + self.tokens_per_side = int(config.mm_tokens_per_image**0.5) + self.kernel_size = self.patches_per_image // self.tokens_per_side + self.avg_pool = torch.nn.AvgPool2d(kernel_size=self.kernel_size, + stride=self.kernel_size) + + def load_weights(self, weights): + # Original `mm_input_projection_weight` is a matmul while we use a linear op with no bias. + self.mm_input_projection.weight.data.copy_( + weights["mm_input_projection_weight"].transpose(0, 1)) + # Gemma3RmsNorm is a layernorm-1P and needs a +1.0. + self.mm_soft_emb_norm.weight.data.copy_( + weights["mm_soft_emb_norm.weight"] + 1.0) + + @torch.inference_mode() + def forward(self, vision_outputs: torch.Tensor): + batch_size, _, seq_length = vision_outputs.shape + reshaped_vision_outputs = vision_outputs.transpose(1, 2) + reshaped_vision_outputs = reshaped_vision_outputs.reshape( + batch_size, seq_length, self.patches_per_image, + self.patches_per_image) + reshaped_vision_outputs = reshaped_vision_outputs.contiguous() + pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs).to( + self.dtype) + pooled_vision_outputs = pooled_vision_outputs.flatten(2) # [B, T, P*P]. + pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2).reshape( + -1, seq_length) + # FlashInfer's rmsnorm needs input to be contiguous. + normed_vision_outputs = self.mm_soft_emb_norm( + pooled_vision_outputs.contiguous()) + projected_vision_outputs = self.mm_input_projection( + normed_vision_outputs) + return projected_vision_outputs.type_as(vision_outputs) + + @register_auto_model("Gemma3ForConditionalGeneration") @register_input_processor(Gemma3InputProcessor, model_type="gemma3") class Gemma3VLM(PreTrainedModel): @@ -114,11 +169,8 @@ def __init__(self, model_config: ModelConfig[Gemma3Config]): self.siglip_tower = SiglipVisionModel(vision_model_config, use_post_layernorm=True) - # NOTE: Use HF implementation. We init the weights after transferring to the `device` since it can take a much - # longer time to initialize them on the CPU. - with no_init_weights(): - self.mm_projector = Gemma3MultiModalProjector(config).eval().to( - self._device) + self.mm_projector = Gemma3MultiModalProjector(model_config).eval().to( + self._device) self.post_config() self.is_loaded = True @@ -153,12 +205,8 @@ def load_weights(self, weights, weight_mapper: BaseWeightMapper): vit_weights = filter_weights("vision_tower", weights) self.siglip_tower.load_weights(vit_weights) - _load_weights_into_hf_module( - model=self.mm_projector, - weights=weights, - prefix="multi_modal_projector", - model_name="multi modal projector", - ) + mm_projector_weights = filter_weights("multi_modal_projector", weights) + self.mm_projector.load_weights(mm_projector_weights) def post_config(self): self.config = self.llm.config @@ -191,12 +239,9 @@ def forward( mm_embeds = [] mm_token_mask = None if len(pixel_values) > 0: - # The shape of `image_features` is `[B, T, embed_dim]`. image_features = self._get_image_features( pixel_values=torch.cat(pixel_values)) - # We need to reshape it to `[B * T, embed_dim]` before passing to `fuse_input_embeds`. - B, T, embed_dim = image_features.shape - mm_embeds = [image_features.reshape(B * T, embed_dim).contiguous()] + mm_embeds = [image_features.contiguous()] # Get token type ids. 0 corresponds to text tokens, 1 corresponds to image tokens. mm_token_mask = torch.isin(input_ids, self.image_token_ids)