[MM] Keep raw pixels through dp-sharded ViT path (#56872)

Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>

vllm/model_executor/models/qwen2_5_vl.py

@@ -1441,7 +1441,7 @@ class Qwen2_5_VLForConditionalGeneration(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values.type(self.visual.dtype),
+                    pixel_values,
                     grid_thw_list,
                     rope_type="rope_3d",
                 )
@@ -1501,7 +1501,7 @@ class Qwen2_5_VLForConditionalGeneration(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values_videos.type(self.visual.dtype),
+                    pixel_values_videos,
                     grid_thw_list,
                     rope_type="rope_3d",
                 )

vllm/model_executor/models/qwen2_vl.py

@@ -1379,7 +1379,7 @@ class Qwen2VLForConditionalGeneration(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values.type(self.visual.dtype),
+                    pixel_values,
                     grid_thw.tolist(),
                     rope_type="rope_3d",
                 )
@@ -1404,7 +1404,7 @@ class Qwen2VLForConditionalGeneration(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values_videos.type(self.visual.dtype),
+                    pixel_values_videos,
                     grid_thw.tolist(),
                     rope_type="rope_3d",
                 )

vllm/model_executor/models/vision.py

@@ -481,6 +481,7 @@ def run_dp_sharded_mrope_vision_model(
     # to work
     max_len_per_rank = max(grouped_pixel_values_len) // embed_dim_reduction_factor
     local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]
+    embed_dtype = next(vision_model.parameters()).dtype
 
     # Run the vision model on the local pixel_values_local
     if rope_type == "rope_2d":
@@ -495,7 +496,7 @@ def run_dp_sharded_mrope_vision_model(
             image_embeds_local = torch.empty(
                 (0, embed_dim_reduction_factor, out_dim),
                 device=pixel_values.device,
-                dtype=pixel_values.dtype,
+                dtype=embed_dtype,
             )
     else:
         if pixel_values_local.shape[0] > 0:
@@ -505,7 +506,7 @@ def run_dp_sharded_mrope_vision_model(
             image_embeds_local = torch.empty(
                 (0, vision_model.out_hidden_size),
                 device=pixel_values.device,
-                dtype=pixel_values.dtype,
+                dtype=embed_dtype,
             )
 
     # Pad the output based on max_len_per_rank