[Bugfix] Make mm_device_do_normalize encoder-cudagraph safe (#55370)

Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>

tests/models/multimodal/generation/test_qwen2_vl.py

@@ -126,8 +126,10 @@ def batch_make_image_embeddings(
     # image to pixel values
     image_processor = processor.image_processor
 
+    mm_config = llm.get_llm().model_config.multimodal_config
+    mm_kwargs = mm_config.merge_mm_processor_kwargs({})
     preprocess_result = image_processor.preprocess(
-        images=images, return_tensors="pt"
+        images=images, return_tensors="pt", **mm_kwargs
     ).data
     pixel_values = preprocess_result["pixel_values"]
     image_grid_thw = preprocess_result["image_grid_thw"]
@@ -210,8 +212,10 @@ def batch_make_video_embeddings(
     # video to pixel values
     video_processor = processor.video_processor
 
+    mm_config = llm.get_llm().model_config.multimodal_config
+    mm_kwargs = mm_config.merge_mm_processor_kwargs({})
     preprocess_result = video_processor.preprocess(
-        videos=videos, return_tensors="pt"
+        videos=videos, return_tensors="pt", **mm_kwargs
     ).data
     pixel_values = preprocess_result["pixel_values_videos"]
     video_grid_thw = preprocess_result["video_grid_thw"]

vllm/model_executor/models/qwen2_5_vl.py

@@ -630,6 +630,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
         vision_config: Qwen2_5_VLVisionConfig,
         norm_eps: float = 1e-6,
         quant_config: QuantizationConfig | None = None,
+        input_norm: nn.Module | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -661,6 +662,9 @@ class Qwen2_5_VisionTransformer(nn.Module):
             in_channels=in_channels,
             hidden_size=self.hidden_size,
         )
+        self.input_norm = (
+            input_norm if input_norm is not None else FusedInputNorm.identity()
+        )
 
         norm_layer = partial(RMSNorm, eps=norm_eps)
         head_dim = self.hidden_size // self.num_heads
@@ -1076,7 +1080,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
         *,
         encoder_metadata: dict[str, torch.Tensor] | None = None,
     ) -> torch.Tensor:
-        hidden_states = x.to(device=self.device, dtype=self.dtype)
+        hidden_states = self.input_norm(x.to(device=self.device), self.dtype)
         hidden_states = self.patch_embed(hidden_states)
 
         seq_len = hidden_states.shape[0]
@@ -1359,9 +1363,9 @@ class Qwen2_5_VLForConditionalGeneration(
                 vision_config=config.vision_config,
                 norm_eps=getattr(config, "rms_norm_eps", 1e-6),
                 quant_config=self.quant_config,
+                input_norm=FusedInputNorm.from_model_config(self.model_config),
                 prefix=maybe_prefix(prefix, "visual"),
             )
-            self.input_norm = FusedInputNorm.from_model_config(self.model_config)
 
         with self._mark_language_model(vllm_config):
             self.language_model = init_vllm_registered_model(
@@ -1436,11 +1440,13 @@ class Qwen2_5_VLForConditionalGeneration(
             image_embeds = image_input["image_embeds"].type(self.visual.dtype)
         else:
             pixel_values = image_input["pixel_values"]
-            pixel_values = self.input_norm(pixel_values, self.visual.dtype)
 
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
-                    self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"
+                    self.visual,
+                    pixel_values.type(self.visual.dtype),
+                    grid_thw_list,
+                    rope_type="rope_3d",
                 )
             else:
                 image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
@@ -1494,14 +1500,11 @@ class Qwen2_5_VLForConditionalGeneration(
             video_embeds = video_input["video_embeds"].type(self.visual.dtype)
         else:
             pixel_values_videos = video_input["pixel_values_videos"]
-            pixel_values_videos = self.input_norm(
-                pixel_values_videos, self.visual.dtype
-            )
 
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values_videos,
+                    pixel_values_videos.type(self.visual.dtype),
                     grid_thw_list,
                     rope_type="rope_3d",
                 )

vllm/model_executor/models/qwen2_vl.py

@@ -537,6 +537,7 @@ class Qwen2VisionTransformer(nn.Module):
         vision_config: Qwen2VLVisionConfig,
         norm_eps: float = 1e-6,
         quant_config: QuantizationConfig | None = None,
+        input_norm: nn.Module | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -564,6 +565,9 @@ class Qwen2VisionTransformer(nn.Module):
             in_channels=in_channels,
             embed_dim=embed_dim,
         )
+        self.input_norm = (
+            input_norm if input_norm is not None else FusedInputNorm.identity()
+        )
 
         norm_layer = partial(nn.LayerNorm, eps=norm_eps)
         head_dim = embed_dim // num_heads
@@ -720,7 +724,7 @@ class Qwen2VisionTransformer(nn.Module):
         encoder_metadata: dict[str, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         # patchify
-        x = x.to(device=self.device, dtype=self.dtype)
+        x = self.input_norm(x.to(device=self.device), self.dtype)
         x = self.patch_embed(x)
 
         if encoder_metadata is None:
@@ -1301,9 +1305,9 @@ class Qwen2VLForConditionalGeneration(
                 config.vision_config,
                 norm_eps=getattr(config, "rms_norm_eps", 1e-6),
                 quant_config=quant_config,
+                input_norm=FusedInputNorm.from_model_config(self.model_config),
                 prefix=maybe_prefix(prefix, "visual"),
             )
-            self.input_norm = FusedInputNorm.from_model_config(self.model_config)
 
         with self._mark_language_model(vllm_config):
             self.language_model = init_vllm_registered_model(
@@ -1374,11 +1378,13 @@ class Qwen2VLForConditionalGeneration(
             image_embeds = image_input["image_embeds"].type(self.visual.dtype)
         else:
             pixel_values = image_input["pixel_values"]
-            pixel_values = self.input_norm(pixel_values, self.visual.dtype)
 
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
-                    self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"
+                    self.visual,
+                    pixel_values.type(self.visual.dtype),
+                    grid_thw.tolist(),
+                    rope_type="rope_3d",
                 )
             else:
                 image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
@@ -1398,13 +1404,10 @@ class Qwen2VLForConditionalGeneration(
             video_embeds = video_input["video_embeds"].type(self.visual.dtype)
         else:
             pixel_values_videos = video_input["pixel_values_videos"]
-            pixel_values_videos = self.input_norm(
-                pixel_values_videos, self.visual.dtype
-            )
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values_videos,
+                    pixel_values_videos.type(self.visual.dtype),
                     grid_thw.tolist(),
                     rope_type="rope_3d",
                 )