refactor(multimodal R-T): Migrate MM models to merge_by_field_config

Migrate step3_vl, tarsier, terratorch, ultravox, voxtral, and whisper to
use merge_by_field_config = True, enabling HF-compatible input shapes.
Remove flatten_bn calls and dead flatten_and_concat function.

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

Author: ayushsatyam146

vllm/model_executor/models/step3_vl.py

@@ -37,7 +37,7 @@
 from vllm.transformers_utils.tokenizer import AnyTokenizer
 
 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
-from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
+from .utils import (AutoWeightsLoader, WeightsMapper,
                     init_vllm_registered_model, maybe_prefix)
 from .vision import run_dp_sharded_vision_model
 
@@ -836,6 +836,7 @@ 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.",
@@ -917,18 +918,21 @@ def _parse_and_validate_image_input(
             return None
 
         if pixel_values is not None:
-            pixel_values = flatten_bn(pixel_values, concat=True)
             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)
                 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

@@ -38,8 +38,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,
                      get_vision_encoder_info)
 
@@ -386,6 +385,8 @@ def _get_layer_index(feature_layer_index: 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"]
@@ -450,7 +451,7 @@ def _parse_and_validate_image_input(
 
             return TarsierImagePixelInputs(
                 type="pixel_values",
-                pixel_values=flatten_bn(pixel_values, concat=True),
+                pixel_values=pixel_values,
             )
 
         if image_embeds is not None:
@@ -459,7 +460,7 @@ def _parse_and_validate_image_input(
                                  f"Got type: {type(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

@@ -215,6 +215,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

@@ -59,8 +59,7 @@
 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
                          SupportsMultiModal, SupportsPP, SupportsQuant)
 from .utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
-                    flatten_bn, make_empty_intermediate_tensors_factory,
-                    maybe_prefix)
+                    make_empty_intermediate_tensors_factory, maybe_prefix)
 
 logger = init_logger(__name__)
 
@@ -812,17 +811,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,