diff --git a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py new file mode 100644 index 000000000..b1cc48177 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_thinker_forward_returns_bare_hidden_states(): + """Verify that MiniCPMO45OmniLLMForConditionalGeneration.forward returns + a bare hidden_states Tensor (not a tuple), which stock vLLM GPUModelRunner + expects. + """ + model = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + torch.nn.Module.__init__(model) + model.llm = MagicMock() + mock_hidden_states = torch.randn(4, 4096) + model.llm.model = MagicMock(return_value=mock_hidden_states) + + out = model.forward( + input_ids=torch.tensor([1, 2, 3, 4]), + positions=torch.tensor([0, 1, 2, 3]), + inputs_embeds=torch.randn(4, 4096), + ) + + assert isinstance(out, torch.Tensor), f"Expected bare torch.Tensor, got {type(out)}" + assert out.shape == (4, 4096) + assert torch.equal(out, mock_hidden_states) + + +def test_thinker_forward_builds_embeddings_once_when_inputs_embeds_none(): + """Verify that when inputs_embeds is None, forward builds input embeddings + exactly once and forwards them to self.llm.model with input_ids=None.""" + model = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + torch.nn.Module.__init__(model) + model.llm = MagicMock() + mock_hidden_states = torch.randn(4, 4096) + model.llm.model = MagicMock(return_value=mock_hidden_states) + + mock_mm_embeds = [torch.randn(2, 4096)] + mock_input_embeds = torch.randn(4, 4096) + model.get_multimodal_embeddings = MagicMock(return_value=mock_mm_embeds) + model.get_input_embeddings = MagicMock(return_value=mock_input_embeds) + + input_ids = torch.tensor([1, 2, 3, 4]) + positions = torch.tensor([0, 1, 2, 3]) + + out = model.forward( + input_ids=input_ids, + positions=positions, + inputs_embeds=None, + ) + + # Assert input embeddings are built exactly once + model.get_multimodal_embeddings.assert_called_once() + model.get_input_embeddings.assert_called_once_with(input_ids, mock_mm_embeds) + + # Verify forwarded to self.llm.model with input_ids=None and inputs_embeds + model.llm.model.assert_called_once_with( + None, + positions, + None, + inputs_embeds=mock_input_embeds, + ) + + assert isinstance(out, torch.Tensor), f"Expected bare torch.Tensor, got {type(out)}" + assert torch.equal(out, mock_hidden_states) + + +@pytest.mark.parametrize( + ("thinker_shape", "expected_shape"), + [ + ((1, 4096), (1, 4096)), + ((5, 4096), (5, 4096)), + ((1, 1, 4096), (1, 4096)), + ((1, 5, 4096), (5, 4096)), + ], +) +def test_wrapper_shape_normalization(thinker_shape, expected_shape): + """Verify that MiniCPMO45OmniForConditionalGeneration normalizes 3D + hidden states of shape (1, seq_len, H) to (seq_len, H), while leaving + 2D hidden states intact (especially (1, H) for single-token decode).""" + wrapper = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + torch.nn.Module.__init__(wrapper) + wrapper.model_stage = "llm" + wrapper.thinker = torch.nn.Module() + mock_hidden = torch.randn(*thinker_shape) + wrapper.thinker.forward = MagicMock(return_value=mock_hidden) + + out = wrapper.forward( + input_ids=torch.tensor([1]), + positions=torch.tensor([0]), + ) + + assert out.text_hidden_states.shape == expected_shape, ( + f"Expected text_hidden_states shape {expected_shape}, got {out.text_hidden_states.shape}" + ) + assert out.multimodal_outputs["latent"].shape == expected_shape, ( + f"Expected latent shape {expected_shape}, got {out.multimodal_outputs['latent'].shape}" + ) + + +if __name__ == "__main__": + test_thinker_forward_returns_bare_hidden_states() + test_thinker_forward_builds_embeddings_once_when_inputs_embeds_none() + for shape, exp in [ + ((1, 4096), (1, 4096)), + ((5, 4096), (5, 4096)), + ((1, 1, 4096), (1, 4096)), + ((1, 5, 4096), (5, 4096)), + ]: + test_wrapper_shape_normalization(shape, exp) + print("ALL TESTS PASSED SUCCESSFULLY!") diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index e9ed05e75..60252e78c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -547,13 +547,13 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp ) if isinstance(thinker_output, tuple): - embeds, text_hidden_states = thinker_output + _, text_hidden_states = thinker_output else: text_hidden_states = thinker_output # Prepare hidden states for downstream stages - # Ensure correct shape: (batch_size, seq_len, hidden_dim) - if added_batch_dim: + # Ensure correct shape: (seq_len, hidden_dim) + if text_hidden_states.ndim == 3 and text_hidden_states.shape[0] == 1: text_hidden_states = text_hidden_states.squeeze(0) # Return hidden states with latent in multimodal_outputs for stage_input_processors diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 6d1c5de32..2781fcc30 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -4611,9 +4611,6 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S inputs_embeds: torch.Tensor | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors: - """Forward pass through thinker model.""" - text_inputs_embeds = None - if intermediate_tensors is not None: inputs_embeds = None @@ -4621,22 +4618,12 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S elif inputs_embeds is None: multimodal_embeddings = self.get_multimodal_embeddings(**kwargs) inputs_embeds = self.get_input_embeddings(input_ids, multimodal_embeddings) - text_inputs_embeds = self.get_input_embeddings( - input_ids, - ( - [(torch.zeros_like(embeddings), "image") for embeddings in multimodal_embeddings] - if multimodal_embeddings is not None - else None - ), - ) input_ids = None - else: - text_inputs_embeds = inputs_embeds # Forward through language model hidden_states = self.llm.model(input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds) - return text_inputs_embeds, hidden_states.unsqueeze(0) if hidden_states.ndim == 2 else hidden_states + return hidden_states def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: """Compute logits from hidden states."""