[Model] Optimize MiniCPM-o multimodal encoder memory (#5188)

Signed-off-by: xsmccc <247721718+xsmccc@users.noreply.github.com>
Co-authored-by: xsmccc <247721718+xsmccc@users.noreply.github.com>

tests/model_executor/models/minicpmo_4_5/test_vision_batching.py

@@ -0,0 +1,175 @@
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    MiniCPMO45OmniLLMForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class _VisionEncoderOutput:
+    last_hidden_state: torch.Tensor
+
+
+@dataclass
+class _VisionConfig:
+    vision_batch_size: int
+
+
+@dataclass
+class _AudioConfig:
+    audio_chunk_length: float = 0
+    audio_pool_step: int = 1
+
+
+@dataclass
+class _FakeConv:
+    weight: torch.Tensor
+
+
+@dataclass
+class _AudioEncoderOutput:
+    last_hidden_state: torch.Tensor
+    hidden_states: tuple[torch.Tensor, ...] | None
+
+
+class _FakeVisionEncoder:
+    def __init__(self) -> None:
+        self.batch_sizes: list[int] = []
+
+    def __call__(
+        self,
+        pixel_values: torch.Tensor,
+        *,
+        patch_attention_mask: torch.Tensor,
+        tgt_sizes: torch.Tensor,
+    ) -> _VisionEncoderOutput:
+        batch_size = pixel_values.shape[0]
+        self.batch_sizes.append(batch_size)
+        assert patch_attention_mask.shape[0] == batch_size
+        assert tgt_sizes.shape[0] == batch_size
+
+        per_item = pixel_values.sum(dim=(1, 2, 3))
+        hidden_states = per_item[:, None, None].expand(-1, 3, 4).clone()
+        return _VisionEncoderOutput(last_hidden_state=hidden_states)
+
+
+class _FakeResampler:
+    def __init__(self) -> None:
+        self.batch_sizes: list[int] = []
+
+    def __call__(self, hidden_states: torch.Tensor, tgt_sizes: torch.Tensor) -> torch.Tensor:
+        self.batch_sizes.append(hidden_states.shape[0])
+        return hidden_states + tgt_sizes.sum(dim=1)[:, None, None]
+
+
+class _FakeAudioEncoder:
+    def __init__(self) -> None:
+        self.conv1 = _FakeConv(weight=torch.empty(1, dtype=torch.float32))
+        self.output_hidden_states: list[bool] = []
+
+    def __call__(
+        self,
+        input_features: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor,
+        output_hidden_states: bool,
+    ) -> _AudioEncoderOutput:
+        self.output_hidden_states.append(output_hidden_states)
+        assert attention_mask.shape == (1, 1, 2, 2)
+        final = torch.full((1, 2, 4), 2.0)
+        earlier = torch.full((1, 2, 4), 1.0)
+        return _AudioEncoderOutput(
+            last_hidden_state=final,
+            hidden_states=(earlier, final) if output_hidden_states else None,
+        )
+
+
+class _FakeVisionModel:
+    def __init__(self, vision_batch_size: int) -> None:
+        self.config = _VisionConfig(vision_batch_size=vision_batch_size)
+        self.vpm = _FakeVisionEncoder()
+        self.resampler = _FakeResampler()
+
+
+class _FakeAudioModel:
+    def __init__(self, audio_encoder: _FakeAudioEncoder, audio_encoder_layer: int) -> None:
+        self.config = _AudioConfig()
+        self.apm = audio_encoder
+        self.audio_projection_layer = torch.nn.Identity()
+        self.audio_avg_pooler = torch.nn.Identity()
+        self.audio_encoder_layer = audio_encoder_layer
+
+    def _get_feat_extract_output_lengths(
+        self,
+        input_lengths: torch.LongTensor,
+    ) -> tuple[torch.LongTensor, torch.LongTensor]:
+        input_lengths_after_cnn = (input_lengths - 1) // 2 + 1
+        input_lengths_after_pooling = (
+            input_lengths_after_cnn - self.config.audio_pool_step
+        ) // self.config.audio_pool_step + 1
+        return input_lengths_after_cnn, input_lengths_after_pooling.to(dtype=torch.int32)
+
+
+def _run_vision_encoder(vision_batch_size: int) -> tuple[torch.Tensor, list[int], list[int]]:
+    model = _FakeVisionModel(vision_batch_size)
+    pixel_values = [
+        torch.full((3, 2, length), fill_value=index + 1, dtype=torch.float32)
+        for index, length in enumerate((4, 3, 2, 4, 1))
+    ]
+    tgt_sizes = torch.tensor([[1, 2], [1, 2], [1, 1], [2, 2], [1, 1]])
+
+    result = MiniCPMO45OmniLLMForConditionalGeneration.get_vision_hidden_states(
+        model,
+        {"pixel_values": pixel_values, "tgt_sizes": tgt_sizes},
+    )
+    return result, model.vpm.batch_sizes, model.resampler.batch_sizes
+
+
+def test_vision_encoder_batches_vpm_before_resampling() -> None:
+    chunked, vision_batches, resampler_batches = _run_vision_encoder(2)
+    unchunked, unchunked_vision_batches, unchunked_resampler_batches = _run_vision_encoder(16)
+
+    assert vision_batches == [2, 2, 1]
+    assert resampler_batches == [5]
+    assert unchunked_vision_batches == [5]
+    assert unchunked_resampler_batches == [5]
+    torch.testing.assert_close(chunked, unchunked, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("vision_batch_size", [0, -2])
+def test_vision_encoder_clamps_nonpositive_batch_size(vision_batch_size: int) -> None:
+    chunked, vision_batches, resampler_batches = _run_vision_encoder(vision_batch_size)
+    unchunked, _, _ = _run_vision_encoder(16)
+
+    assert vision_batches == [1, 1, 1, 1, 1]
+    assert resampler_batches == [5]
+    torch.testing.assert_close(chunked, unchunked, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    ("audio_encoder_layer", "expected_value", "expected_hidden_states"),
+    [(-1, 2.0, False), (0, 1.0, True)],
+)
+def test_audio_encoder_retains_layers_only_for_nonfinal_selection(
+    audio_encoder_layer: int,
+    expected_value: float,
+    expected_hidden_states: bool,
+) -> None:
+    audio_encoder = _FakeAudioEncoder()
+    model = _FakeAudioModel(audio_encoder, audio_encoder_layer)
+
+    result = MiniCPMO45OmniLLMForConditionalGeneration.get_audio_hidden_states(
+        model,
+        {
+            "audio_features": torch.zeros((1, 80, 4)),
+            "audio_feature_lens": [torch.tensor([4])],
+        },
+    )
+
+    assert audio_encoder.output_hidden_states == [expected_hidden_states]
+    torch.testing.assert_close(result[0], torch.full((2, 4), expected_value))

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -4170,6 +4170,8 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         device = pixel_values[0].device
         dtype = pixel_values[0].dtype
 
+        # VPM batching bounds encoder activations below, while this padding
+        # buffer intentionally remains sized for the full input batch.
         all_pixel_values = torch.zeros((B, 3, P, L), dtype=dtype, device=device)
         for i, pixel_values_item in enumerate(pixel_values):
             L_item = pixel_values_item.shape[-1]
@@ -4183,16 +4185,29 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         for i, num_patches_item in enumerate(num_patches):
             patch_attn_mask[i, :num_patches_item] = True
 
-        vision_embedding = self.vpm(
-            all_pixel_values,
-            patch_attention_mask=patch_attn_mask.unsqueeze(1),
-            tgt_sizes=tgt_sizes,
-        )
+        patch_attn_mask = patch_attn_mask.unsqueeze(1)
 
-        if not isinstance(vision_embedding, torch.Tensor):
-            vision_embedding = vision_embedding.last_hidden_state
+        def encode_vision(start: int, end: int) -> torch.Tensor:
+            chunk_tgt_sizes = tgt_sizes[start:end]
+            vision_embedding = self.vpm(
+                all_pixel_values[start:end],
+                patch_attention_mask=patch_attn_mask[start:end],
+                tgt_sizes=chunk_tgt_sizes,
+            )
 
-        return self.resampler(vision_embedding, tgt_sizes)
+            if not isinstance(vision_embedding, torch.Tensor):
+                vision_embedding = vision_embedding.last_hidden_state
+
+            return vision_embedding
+
+        vision_batch_size = max(1, int(self.config.vision_batch_size))
+        if B <= vision_batch_size:
+            return self.resampler(encode_vision(0, B), tgt_sizes)
+
+        vision_embeddings = [
+            encode_vision(start, min(start + vision_batch_size, B)) for start in range(0, B, vision_batch_size)
+        ]
+        return self.resampler(torch.cat(vision_embeddings, dim=0), tgt_sizes)
 
     def _process_vision_input(
         self,
@@ -4309,9 +4324,19 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             audio_attention_mask_ = torch.logical_or(audio_attention_mask_, torch.logical_not(chunk_mask))
 
         audio_attention_mask[audio_attention_mask_] = float("-inf")
-        audio_states = self.apm(wavforms, attention_mask=audio_attention_mask, output_hidden_states=True).hidden_states[
-            self.audio_encoder_layer
-        ]
+        selects_final_layer = self.audio_encoder_layer == -1
+        audio_outputs = self.apm(
+            wavforms,
+            attention_mask=audio_attention_mask,
+            output_hidden_states=not selects_final_layer,
+        )
+        if selects_final_layer:
+            audio_states = audio_outputs.last_hidden_state
+        else:
+            # Whisper follows the Hugging Face hidden-state ordering: the
+            # embedding output, followed by each encoder layer output.
+            assert audio_outputs.hidden_states is not None
+            audio_states = audio_outputs.hidden_states[self.audio_encoder_layer]
         audio_embeds = self.audio_projection_layer(audio_states)
 
         audio_embeds = audio_embeds.transpose(1, 2)