add tests

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

Author: HuiyingLi

tests/unit_tests/_transformers/test_auto_model.py

@@ -14,20 +14,19 @@
 
 import logging
 import types
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
 from nemo_automodel._transformers.auto_model import (
+    _consume_config_overrides,
     _get_next_fallback_attn,
     _init_model,
     _patch_attention,
-    _consume_config_overrides,
 )
 from nemo_automodel._transformers.infrastructure import _apply_peft_and_lower_precision
-from nemo_automodel._transformers.model_init import _filter_kwargs_for_init
-from nemo_automodel._transformers.model_init import _get_mixin_wrapped_class
+from nemo_automodel._transformers.model_init import _filter_kwargs_for_init, _get_mixin_wrapped_class
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 
 
@@ -514,6 +513,7 @@ class TestNeedSetupCacheClassesMapping:
     def test_shim_does_not_overwrite_existing_attribute(self):
         """If NEED_SETUP_CACHE_CLASSES_MAPPING already exists, shim doesn't overwrite."""
         import importlib
+
         import transformers.generation.utils as gen_utils
 
         sentinel = {"test": "sentinel_value"}
@@ -532,6 +532,7 @@ def test_shim_does_not_overwrite_existing_attribute(self):
     def test_shim_creates_attribute_when_missing(self):
         """If NEED_SETUP_CACHE_CLASSES_MAPPING is missing, shim creates it."""
         import importlib
+
         import transformers.generation.utils as gen_utils
 
         # Remove the attribute if it exists
@@ -699,3 +700,41 @@ def __init__(self):
 
         assert is_custom is False
         mock_wrap.assert_called_once_with(FakeModel)
+
+
+class TestNeMoAutoModelForMultimodalLM:
+    """Tests for the NeMoAutoModelForMultimodalLM class and its exports."""
+
+    def test_class_exists_and_inherits_correctly(self):
+        from transformers import AutoModelForMultimodalLM
+
+        from nemo_automodel._transformers.auto_model import NeMoAutoModelForMultimodalLM, _BaseNeMoAutoModelClass
+
+        assert issubclass(NeMoAutoModelForMultimodalLM, _BaseNeMoAutoModelClass)
+        assert issubclass(NeMoAutoModelForMultimodalLM, AutoModelForMultimodalLM)
+
+    def test_has_from_pretrained_and_from_config(self):
+        from nemo_automodel._transformers.auto_model import NeMoAutoModelForMultimodalLM
+
+        assert callable(NeMoAutoModelForMultimodalLM.from_pretrained)
+        assert callable(NeMoAutoModelForMultimodalLM.from_config)
+
+    def test_lazy_export_from_transformers_subpackage(self):
+        from nemo_automodel._transformers import NeMoAutoModelForMultimodalLM
+
+        assert NeMoAutoModelForMultimodalLM is not None
+
+    def test_lazy_export_from_top_level_package(self):
+        from nemo_automodel import NeMoAutoModelForMultimodalLM
+
+        assert NeMoAutoModelForMultimodalLM is not None
+
+    def test_top_level_dir_includes_multimodal(self):
+        import nemo_automodel
+
+        assert "NeMoAutoModelForMultimodalLM" in dir(nemo_automodel)
+
+    def test_transformers_subpackage_all_includes_multimodal(self):
+        import nemo_automodel._transformers as pkg
+
+        assert "NeMoAutoModelForMultimodalLM" in pkg.__all__

tests/unit_tests/_transformers/test_transformers_utils.py

@@ -16,7 +16,7 @@
 
 import pytest
 
-from nemo_automodel._transformers.utils import sliding_window_overwrite
+from nemo_automodel._transformers.utils import apply_qwen3_omni_config_patch, sliding_window_overwrite
 
 
 class TestSlidingWindowOverwrite:
@@ -165,3 +165,43 @@ def test_sliding_window_overwrite_hasattr_behavior(self, mock_from_pretrained):
 
         # Should return empty dict when use_sliding_window is not exactly False
         assert result == {}
+
+
+class TestApplyQwen3OmniConfigPatch:
+    """Test cases for apply_qwen3_omni_config_patch function."""
+
+    def test_patch_sets_use_sliding_window_default(self):
+        """Verify the patch adds use_sliding_window=False to the config class."""
+        from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+            Qwen3OmniMoeTalkerCodePredictorConfig,
+        )
+
+        apply_qwen3_omni_config_patch()
+        assert hasattr(Qwen3OmniMoeTalkerCodePredictorConfig, "use_sliding_window")
+
+    def test_patch_is_idempotent(self):
+        """Calling the patch twice does not raise or change the value."""
+        from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+            Qwen3OmniMoeTalkerCodePredictorConfig,
+        )
+
+        apply_qwen3_omni_config_patch()
+        apply_qwen3_omni_config_patch()
+        assert Qwen3OmniMoeTalkerCodePredictorConfig.use_sliding_window is False
+
+    def test_patch_does_not_overwrite_existing_attribute(self):
+        """If the attribute already exists (e.g. fixed upstream), patch is a no-op."""
+        from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+            Qwen3OmniMoeTalkerCodePredictorConfig,
+        )
+
+        original = getattr(Qwen3OmniMoeTalkerCodePredictorConfig, "use_sliding_window", None)
+        Qwen3OmniMoeTalkerCodePredictorConfig.use_sliding_window = True
+        try:
+            apply_qwen3_omni_config_patch()
+            assert Qwen3OmniMoeTalkerCodePredictorConfig.use_sliding_window is True
+        finally:
+            if original is None:
+                del Qwen3OmniMoeTalkerCodePredictorConfig.use_sliding_window
+            else:
+                Qwen3OmniMoeTalkerCodePredictorConfig.use_sliding_window = original

tests/unit_tests/models/qwen3_omni_moe/test_qwen3_omni_moe_model.py

@@ -25,13 +25,12 @@
     Qwen3OmniMoeThinkerForConditionalGeneration as HFQwen3OmniMoeThinkerForConditionalGeneration,
 )
 
+from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.qwen3_omni_moe.model import (
     Qwen3OmniMoeThinkerForConditionalGeneration,
     Qwen3OmniMoeThinkerTextModel,
 )
 from nemo_automodel.components.moe.config import MoEConfig
-from nemo_automodel.components.models.common import BackendConfig
-
 
 pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 
@@ -259,3 +258,47 @@ def test_modelclass_export_exists():
 
     assert hasattr(omni_module, "ModelClass")
     assert omni_module.ModelClass is Qwen3OmniMoeThinkerForConditionalGeneration
+
+
+@patch.object(HFQwen3OmniMoeThinkerForConditionalGeneration, "__init__", new=_stub_hf_init)
+@patch("nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerTextRotaryEmbedding")
+def test_forward_unpacks_vision_features_from_named_output(rotary_cls, thinker_config, backend_config, moe_config, device):
+    """get_image_features / get_video_features return a named output object;
+    forward() must use .pooler_output and .deepstack_features, not tuple unpacking."""
+    rotary_cls.return_value = MagicMock(side_effect=lambda x, y: (torch.zeros_like(x), torch.zeros_like(x)))
+    model = Qwen3OmniMoeThinkerForConditionalGeneration(thinker_config, moe_config=moe_config, backend=backend_config).to(device)
+    model.config = thinker_config
+
+    hidden_size = thinker_config.text_config.hidden_size
+    vocab_size = thinker_config.text_config.vocab_size
+    batch, seq_len = 1, 6
+
+    fake_image_embed = torch.randn(2, hidden_size, device=device)
+    fake_deepstack = [torch.randn(2, hidden_size, device=device)]
+
+    fake_vision_output = SimpleNamespace(
+        last_hidden_state=torch.randn(2, hidden_size, device=device),
+        pooler_output=fake_image_embed,
+        hidden_states=None,
+        attentions=None,
+        deepstack_features=fake_deepstack,
+    )
+
+    hidden = torch.randn(batch, seq_len, hidden_size, device=device, dtype=model.lm_head.weight.dtype)
+    input_ids = torch.randint(0, vocab_size, (batch, seq_len), device=device)
+
+    with (
+        patch.object(model.model, "forward", return_value=hidden),
+        patch.object(model, "get_image_features", return_value=fake_vision_output) as mock_gif,
+        patch.object(model, "get_placeholder_mask", return_value=(
+            torch.zeros(batch, seq_len, 1, dtype=torch.bool, device=device),
+            torch.zeros(batch, seq_len, 1, dtype=torch.bool, device=device),
+            torch.zeros(batch, seq_len, 1, dtype=torch.bool, device=device),
+        )),
+    ):
+        pixel_values = torch.randn(1, 3, 224, 224, device=device)
+        image_grid_thw = torch.tensor([[1, 14, 14]], device=device)
+        logits = model(input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw)
+
+    mock_gif.assert_called_once()
+    assert logits.shape == (batch, seq_len, vocab_size)

tests/unit_tests/recipes/test_finetune_vlm_helpers.py

@@ -11,13 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
 import pytest
 import torch
 import torch.nn as nn
-from unittest.mock import patch, MagicMock
-from types import SimpleNamespace
-
-from contextlib import nullcontext
 
 from nemo_automodel.components.loggers.metric_logger import MetricsSample
 from nemo_automodel.recipes.vlm.finetune import (
@@ -41,7 +41,6 @@ def test_get_model_name_prefers_pretrained_path():
     assert _get_model_name(_Cfg()) is None
 
 
-from nemo_automodel.components.checkpoint.checkpointing import Checkpointer, CheckpointingConfig
 
 
 
@@ -401,6 +400,7 @@ def test_autoprocessor_success():
 def test_autoprocessor_exception_handling(caplog):
     """Test AutoProcessor exception handling and logging in build_dataloader."""
     import logging
+
     from nemo_automodel.recipes.vlm.finetune import build_dataloader
 
     with patch('transformers.AutoProcessor.from_pretrained') as mock_from_pretrained, \
@@ -434,6 +434,7 @@ def test_autoprocessor_exception_handling(caplog):
 def test_autoprocessor_with_processor_kwargs(caplog):
     """Test AutoProcessor exception handling when cfg_processor has no instantiate method."""
     import logging
+
     from nemo_automodel.recipes.vlm.finetune import build_dataloader
 
     # Simple processor config class without instantiate method
@@ -707,13 +708,12 @@ def get(self, key, default=None):
 
 
 from nemo_automodel.recipes.vlm.finetune import (
-    build_step_scheduler,
-    build_lr_scheduler,
     build_checkpoint_config,
+    build_lr_scheduler,
+    build_step_scheduler,
     calculate_loss,
 )
 
-
 # -----------------------------------------------------------------------------
 # build_step_scheduler tests
 # -----------------------------------------------------------------------------
@@ -1847,3 +1847,33 @@ def get(self, key, default=None):
         )
 
     assert model is not None
+
+
+@pytest.mark.parametrize("entry_point", ["from_config", "from_pretrained"])
+def test_vlm_build_model_accepts_multimodal_lm_entry_points(entry_point):
+    """Test that VLM build_model accepts NeMoAutoModelForMultimodalLM entry points."""
+    from nemo_automodel._transformers import NeMoAutoModelForMultimodalLM
+
+    target = getattr(NeMoAutoModelForMultimodalLM, entry_point)
+
+    class NeMoVLMModelConfig:
+        def __init__(self):
+            self._target_ = target
+
+        def instantiate(self, **kwargs):
+            return DummyModel()
+
+        def get(self, key, default=None):
+            return getattr(self, key, default)
+
+    cfg_model = NeMoVLMModelConfig()
+
+    with patch('nemo_automodel.recipes.vlm.finetune._supports_logits_to_keep', return_value=True):
+        model = build_model(
+            cfg_model=cfg_model,
+            cfg_freeze=None,
+            cfg_peft=None,
+            seed=42,
+        )
+
+    assert model is not None