fix

Signed-off-by: Pamela <[email protected]>

Author: pamelap-nvidia

tensorrt_llm/_torch/models/modeling_qwen2vl.py

@@ -7,7 +7,11 @@
 from transformers import (AutoProcessor, AutoTokenizer, PretrainedConfig,
                           PreTrainedModel, Qwen2_5_VLForConditionalGeneration,
                           Qwen2VLForConditionalGeneration)
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import \
+    Qwen2_5_VisionTransformerPretrainedModel
 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
+from transformers.models.qwen2_vl.modeling_qwen2_vl import \
+    Qwen2VisionTransformerPretrainedModel
 
 from tensorrt_llm.inputs.multimodal import MultimodalParams
 
@@ -358,15 +362,21 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig],
         # TODO: Change the model class to TRT-LLM's Qwen2VisionModel
         # Currently, copying vision encoder on all devices.
         # NOTE: Using attn_implementation='flash_attention_2' to avoid the issue of vision model's GPU OOM.
-        model = model_class.from_pretrained(
-            model_path,
-            torch_dtype=pretrained_config.torch_dtype,
-            attn_implementation='flash_attention_2',
-            ignore_mismatched_sizes=True).eval()
+        hf_model_config = AutoConfig.from_pretrained(model_path)
+        vision_model = model_class(config=hf_model_config.vision_config,
+                                   torch_dtype=pretrained_config.torch_dtype,
+                                   attn_implementation='flash_attention_2')
         # TODO: Make vision model compatible with meta init mode and load_weights at the same place
-        self.visual = model.visual.to(self.device)
+        self.visual = vision_model.to(self.device)
         self.post_config()
 
+    def load_weights(self, weights):
+        filtered_weights = {
+            k.replace('visual.', ''): v
+            for k, v in weights.items() if k.startswith('visual.')
+        }
+        self.visual.load_state_dict(filtered_weights)
+
     def post_config(self):
         self.config = self.visual.config
 
@@ -504,6 +514,7 @@ def init_rotary_cos_sin_ori(self):
 
     def load_weights(self, weights):
         self.llm.load_weights(weights)
+        self.mm_encoder.load_weights(weights)
         self.init_rotary_cos_sin_ori()
 
     def infer_max_seq_len(self) -> int:
@@ -676,7 +687,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
         super().__init__(model_config, *args, **kwargs)
         if not DISAGG:
             self.mm_encoder = Qwen2VisionModelBase(
-                model_config, Qwen2VLForConditionalGeneration)
+                model_config, Qwen2VisionTransformerPretrainedModel)
 
 
 @register_vision_encoder(Qwen2VisionModelBase,
@@ -697,4 +708,4 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
         super().__init__(model_config, *args, **kwargs)
         if not DISAGG:
             self.mm_encoder = Qwen2VisionModelBase(
-                model_config, Qwen2_5_VLForConditionalGeneration)
+                model_config, Qwen2_5_VisionTransformerPretrainedModel)