[Model] Use merge_by_field_config for MM models (R-T)

Signed-off-by: Ayush Satyam <[email protected]>

Author: ayushsatyam146

vllm/model_executor/models/step3_vl.py

@@ -49,7 +49,6 @@
 from .utils import (
     AutoWeightsLoader,
     WeightsMapper,
-    flatten_bn,
     init_vllm_registered_model,
     maybe_prefix,
 )
@@ -895,6 +894,8 @@ def forward(
     dummy_inputs=Step3VLDummyInputsBuilder,
 )
 class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+    merge_by_field_config = True
+
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "model.": "language_model.model.",
@@ -982,18 +983,25 @@ def _parse_and_validate_image_input(
             return None
 
         if pixel_values is not None:
-            pixel_values = flatten_bn(pixel_values, concat=True)
+            if isinstance(pixel_values, list):
+                pixel_values = torch.cat(pixel_values)
             if pixel_values.dim() >= 3:
                 pixel_values = pixel_values.view(-1, *pixel_values.shape[-3:])
             if patch_pixel_values is not None:
-                patch_pixel_values = flatten_bn(patch_pixel_values, concat=True)
+                if isinstance(patch_pixel_values, list):
+                    patch_pixel_values = torch.cat(patch_pixel_values)
                 patch_pixel_values = patch_pixel_values.view(
                     -1, *patch_pixel_values.shape[-3:]
                 )
                 # Handle empty patch_pixel_values by setting to None
                 if patch_pixel_values.shape[0] == 0:
                     patch_pixel_values = None
-            num_patches = flatten_bn(num_patches, concat=True).tolist()
+            if isinstance(num_patches, torch.Tensor):
+                num_patches = num_patches.tolist()
+            elif isinstance(num_patches, list):
+                num_patches = [
+                    n.item() if isinstance(n, torch.Tensor) else n for n in num_patches
+                ]
 
             return Step3VLImagePixelInputs(
                 type="pixel_values",

vllm/model_executor/models/tarsier.py

@@ -47,12 +47,7 @@
 from .clip import CLIPVisionModel
 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
 from .siglip import SiglipVisionModel
-from .utils import (
-    AutoWeightsLoader,
-    flatten_bn,
-    init_vllm_registered_model,
-    maybe_prefix,
-)
+from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
 from .vision import (
     VisionEncoderInfo,
     get_num_selected_vision_tokens,
@@ -404,6 +399,8 @@ def _get_layer_index(feature_layer_index: int, num_hidden_layers_total: int) ->
     dummy_inputs=TarsierDummyInputsBuilder,
 )
 class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+    merge_by_field_config = True
+
     packed_modules_mapping = {
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
         "gate_up_proj": ["gate_proj", "up_proj"],
@@ -472,9 +469,11 @@ def _parse_and_validate_image_input(
                     f"Incorrect type of pixel values. Got type: {type(pixel_values)}"
                 )
 
+            if isinstance(pixel_values, list):
+                pixel_values = torch.cat(pixel_values)
             return TarsierImagePixelInputs(
                 type="pixel_values",
-                pixel_values=flatten_bn(pixel_values, concat=True),
+                pixel_values=pixel_values,
             )
 
         if image_embeds is not None:
@@ -483,9 +482,11 @@ def _parse_and_validate_image_input(
                     "Incorrect type of image embeddings. "
                     f"Got type: {type(image_embeds)}"
                 )
+            if isinstance(image_embeds, list):
+                image_embeds = torch.cat(image_embeds)
             return TarsierImageEmbeddingInputs(
                 type="image_embeds",
-                data=flatten_bn(image_embeds, concat=True),
+                data=image_embeds,
             )
 
         raise AssertionError("This line should be unreachable.")

vllm/model_executor/models/terratorch.py

@@ -226,6 +226,7 @@ def apply(
     dummy_inputs=TerratorchInputBuilder,
 )
 class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):
+    merge_by_field_config = True
     supports_multimodal_raw_input_only = True
     is_pooling_model = True
 

vllm/model_executor/models/transformers.py

@@ -79,7 +79,6 @@
     AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
-    flatten_bn,
     make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
@@ -842,17 +841,6 @@ def compute_logits(
         return logits
 
 
-def flatten_and_concat(x: list[torch.Tensor]) -> torch.Tensor:
-    """Flatten until a list of tensors can be concatenated then do concat"""
-
-    def _can_concat(x: list[torch.Tensor]):
-        return len(set(map(lambda _x: _x.shape[1:], x))) == 1
-
-    if _can_concat(x):
-        return torch.concat(x)
-    return flatten_and_concat(flatten_bn(x))
-
-
 @MULTIMODAL_REGISTRY.register_processor(
     MultiModalProcessor,
     info=MultiModalProcessingInfo,