update tests

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

Author: HuiyingLi

tests/unit_tests/_transformers/test_auto_model.py

@@ -28,7 +28,6 @@
     _consume_config_overrides,
     _filter_kwargs_for_init,
 )
-from nemo_automodel._transformers.model_init import _ensure_pad_token_id
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 
 
@@ -700,237 +699,3 @@ def __init__(self):
 
         assert is_custom is False
         mock_wrap.assert_called_once_with(FakeModel)
-
-
-# =============================================================================
-# Tests for pad_token_id fix and config forwarding in _init_model
-# =============================================================================
-
-
-class TestEnsurePadTokenId:
-    """Test _ensure_pad_token_id for transformers v5 compatibility."""
-
-    def test_pad_token_id_set_when_missing(self):
-        """Config without pad_token_id gets it set to None."""
-
-        class BareConfig:
-            pass
-
-        config = BareConfig()
-        assert not hasattr(config, "pad_token_id")
-
-        _ensure_pad_token_id(config)
-
-        assert hasattr(config, "pad_token_id")
-        assert config.pad_token_id is None
-
-    def test_pad_token_id_preserved_when_present(self):
-        """Config that already has pad_token_id keeps its value unchanged."""
-
-        class ConfigWithPad:
-            pad_token_id = 42
-
-        config = ConfigWithPad()
-        _ensure_pad_token_id(config)
-        assert config.pad_token_id == 42
-
-    def test_sub_config_pad_token_id_patched(self):
-        """Nested sub-configs (e.g. text_config) also get pad_token_id set."""
-
-        class SubConfig:
-            """Mimics a HF sub-config (has to_dict)."""
-            def to_dict(self):
-                return {}
-
-        class TopConfig:
-            pad_token_id = 0  # top-level already has it
-
-            def __init__(self):
-                self.text_config = SubConfig()
-                self.vision_config = SubConfig()
-
-        config = TopConfig()
-        assert not hasattr(config.text_config, "pad_token_id")
-        assert not hasattr(config.vision_config, "pad_token_id")
-
-        _ensure_pad_token_id(config)
-
-        # Top-level preserved
-        assert config.pad_token_id == 0
-        # Sub-configs patched
-        assert config.text_config.pad_token_id is None
-        assert config.vision_config.pad_token_id is None
-
-    def test_sub_config_pad_token_id_preserved_when_present(self):
-        """Sub-config that already has pad_token_id keeps its value."""
-
-        class SubConfig:
-            pad_token_id = 99
-            def to_dict(self):
-                return {}
-
-        class TopConfig:
-            def __init__(self):
-                self.text_config = SubConfig()
-
-        config = TopConfig()
-        _ensure_pad_token_id(config)
-
-        assert config.pad_token_id is None  # top had none, got patched
-        assert config.text_config.pad_token_id == 99  # preserved
-
-    def test_non_config_attributes_ignored(self):
-        """Attributes without to_dict (plain ints, strings, etc.) are skipped."""
-
-        class TopConfig:
-            def __init__(self):
-                self.hidden_size = 768
-                self.model_type = "test"
-                self.some_list = [1, 2, 3]
-
-        config = TopConfig()
-        _ensure_pad_token_id(config)
-
-        assert config.pad_token_id is None
-        # Other attributes untouched
-        assert config.hidden_size == 768
-        assert config.model_type == "test"
-
-    def test_integration_with_init_model(self):
-        """_init_model applies _ensure_pad_token_id to the config."""
-
-        class BareConfig:
-            name_or_path = "test-model"
-
-        config = BareConfig()
-        assert not hasattr(config, "pad_token_id")
-
-        cls = MagicMock()
-        cls._model_mapping = {}
-        cls._from_config_parent_class = MagicMock(return_value=MagicMock())
-
-        with (
-            patch("nemo_automodel._transformers.model_init.get_architectures", return_value=[]),
-            patch("nemo_automodel._transformers.model_init._get_mixin_wrapped_class", side_effect=lambda c: c),
-        ):
-            _init_model(
-                cls,
-                config,
-                attn_implementation="eager",
-                torch_dtype="auto",
-                quantization_config=None,
-                force_hf=False,
-            )
-
-        assert config.pad_token_id is None
-
-
-class TestConfigForwardingInPretrainedPaths:
-    """Test that the patched config is forwarded to HF from_pretrained calls
-    so they don't reload a fresh copy missing the pad_token_id fix."""
-
-    def _make_cls(self, model_mapping_dict=None):
-        cls = MagicMock()
-        cls._model_mapping = model_mapping_dict or {}
-        return cls
-
-    def _make_fake_config(self, *, has_pad_token_id=False):
-        config = MagicMock()
-        config.architectures = []
-        config.to_dict.return_value = {}
-        if not has_pad_token_id:
-            del config.pad_token_id  # ensure hasattr returns False
-        return config
-
-    def test_force_hf_pretrained_forwards_config_in_kwargs(self):
-        """force_hf + from_pretrained path passes config in kwargs."""
-        fake_config = self._make_fake_config()
-        fake_model = MagicMock()
-
-        cls = self._make_cls({type(fake_config): type(fake_model)})
-        cls._from_pretrained_parent_class = MagicMock(return_value=fake_model)
-
-        with (
-            patch("nemo_automodel._transformers.model_init.get_hf_config", return_value=fake_config),
-            patch("nemo_automodel._transformers.model_init._get_mixin_wrapped_class", side_effect=lambda c: c),
-        ):
-            _init_model(
-                cls,
-                "some-model-name",  # string triggers is_pretrained_init=True
-                attn_implementation="eager",
-                torch_dtype="auto",
-                quantization_config=None,
-                force_hf=True,
-            )
-
-        # Verify _from_pretrained_parent_class was called with config in kwargs
-        call_kwargs = cls._from_pretrained_parent_class.call_args[1]
-        assert "config" in call_kwargs
-        assert call_kwargs["config"] is fake_config
-
-    def test_fallback_hf_pretrained_forwards_config_in_kwargs(self):
-        """Fallback HF + from_pretrained path passes config in kwargs."""
-        fake_config = self._make_fake_config()
-        fake_model = MagicMock()
-
-        cls = self._make_cls({type(fake_config): type(fake_model)})
-        cls._from_pretrained_parent_class = MagicMock(return_value=fake_model)
-
-        with (
-            patch("nemo_automodel._transformers.model_init.get_hf_config", return_value=fake_config),
-            patch("nemo_automodel._transformers.model_init.get_architectures", return_value=[]),
-            patch("nemo_automodel._transformers.model_init._get_mixin_wrapped_class", side_effect=lambda c: c),
-        ):
-            _init_model(
-                cls,
-                "some-model-name",  # string triggers is_pretrained_init=True
-                attn_implementation="eager",
-                torch_dtype="auto",
-                quantization_config=None,
-                force_hf=False,
-            )
-
-        call_kwargs = cls._from_pretrained_parent_class.call_args[1]
-        assert "config" in call_kwargs
-        assert call_kwargs["config"] is fake_config
-
-    def test_custom_model_pretrained_does_not_receive_config_in_kwargs(self):
-        """Custom model path must NOT get config in kwargs (it's passed positionally)."""
-        fake_config = MagicMock()
-        fake_config.architectures = ["FakeArch"]
-        fake_config.to_dict.return_value = {}
-        fake_config.torch_dtype = "bfloat16"
-        # Ensure pad_token_id is missing so the fix sets it
-        del fake_config.pad_token_id
-
-        class FakeCustomModel(torch.nn.Module):
-            def __init__(self, config, **kwargs):
-                super().__init__()
-                self.config = config
-                # Store kwargs so we can inspect them in the test
-                self._init_kwargs = kwargs
-
-        cls = MagicMock()
-        registry_mock = {"FakeArch": FakeCustomModel}
-
-        with (
-            patch("nemo_automodel._transformers.model_init.get_hf_config", return_value=fake_config),
-            patch("nemo_automodel._transformers.model_init.get_architectures", return_value=["FakeArch"]),
-            patch("nemo_automodel._transformers.model_init.ModelRegistry") as mock_registry,
-            patch("nemo_automodel._transformers.model_init._download_model_weights"),
-        ):
-            mock_registry.model_arch_name_to_cls = registry_mock
-            is_custom, model = _init_model(
-                cls,
-                "some-model-name",  # string triggers is_pretrained_init=True
-                attn_implementation="eager",
-                torch_dtype="bfloat16",
-                quantization_config=None,
-                force_hf=False,
-            )
-
-        assert is_custom is True
-        # config was passed positionally
-        assert model.config is fake_config
-        # config must NOT be in kwargs (would cause TypeError: got multiple values)
-        assert "config" not in model._init_kwargs

tests/unit_tests/datasets/vlm/test_collate_fns.py

@@ -523,6 +523,24 @@ def test_kimi_vl_collate_fn_extracts_images(collate_mod, monkeypatch):
     assert forward_call["images"] == ["test_image.jpg"]
 
 
+def test_kimi_vl_collate_fn_passes_add_special_tokens_false(collate_mod, monkeypatch):
+    """Test that kimi_vl_collate_fn passes add_special_tokens=False to processor."""
+    processor = DummyKimiVLProcessor()
+
+    labels_stub = torch.tensor([[10, 11, 12, 13, 14]], dtype=torch.long)
+    monkeypatch.setattr(
+        collate_mod, "build_labels", lambda *args, **kwargs: labels_stub, raising=True
+    )
+
+    examples = [{"conversation": CONVERSATION}]
+    collate_mod.kimi_vl_collate_fn(examples, processor)
+
+    assert len(processor.forward_calls) == 1
+    forward_call = processor.forward_calls[0]
+    assert "add_special_tokens" in forward_call
+    assert forward_call["add_special_tokens"] is False
+
+
 def test_kimi_vl_collate_fn_multiple_examples(collate_mod, monkeypatch):
     """Test kimi_vl_collate_fn handles multiple examples."""
     processor = DummyKimiVLProcessor()

tests/unit_tests/recipes/test_finetune_vlm_helpers.py

@@ -154,7 +154,7 @@ def get(self, key, default=None):
 
     class FreezeConfig:
         def to_dict(self):
-            return {"freeze_embeddings": True, "freeze_language_model": False}
+            return {"freeze_language_model": False, "freeze_vision_tower": True}
 
     with patch('nemo_automodel.recipes.vlm.finetune._supports_logits_to_keep', return_value=True):
         model = build_model(
@@ -166,7 +166,117 @@ def to_dict(self):
 
     # Verify freeze_config was passed to model instantiation
     assert "freeze_config" in captured_kwargs
-    assert captured_kwargs["freeze_config"] == {"freeze_embeddings": True, "freeze_language_model": False}
+    assert captured_kwargs["freeze_config"] == {"freeze_language_model": False, "freeze_vision_tower": True}
+
+
+def test_build_model_passes_moe_config_from_parallelizer_config():
+    """Test that cfg_moe as MoEParallelizerConfig is forwarded directly."""
+    from nemo_automodel._transformers import NeMoAutoModelForImageTextToText
+    from nemo_automodel.components.moe.config import MoEParallelizerConfig
+
+    captured_kwargs = {}
+
+    class CapturingModelConfig:
+        def __init__(self):
+            self._target_ = NeMoAutoModelForImageTextToText.from_pretrained
+
+        def instantiate(self, **kwargs):
+            captured_kwargs.update(kwargs)
+            return DummyModel()
+
+        def get(self, key, default=None):
+            return getattr(self, key, default)
+
+    cfg_model = CapturingModelConfig()
+    moe_cfg = MoEParallelizerConfig()
+
+    with patch('nemo_automodel.recipes.vlm.finetune._supports_logits_to_keep', return_value=True):
+        build_model(
+            cfg_model=cfg_model,
+            cfg_freeze=None,
+            cfg_peft=None,
+            seed=123,
+            cfg_moe=moe_cfg,
+            activation_checkpointing=True,
+        )
+
+    assert "moe_config" in captured_kwargs
+    assert captured_kwargs["moe_config"] is moe_cfg
+    assert captured_kwargs["activation_checkpointing"] is True
+
+
+def test_build_model_passes_moe_config_from_dict_like():
+    """Test that cfg_moe with to_dict() is converted to MoEParallelizerConfig."""
+    from nemo_automodel._transformers import NeMoAutoModelForImageTextToText
+    from nemo_automodel.components.moe.config import MoEParallelizerConfig
+
+    captured_kwargs = {}
+
+    class CapturingModelConfig:
+        def __init__(self):
+            self._target_ = NeMoAutoModelForImageTextToText.from_pretrained
+
+        def instantiate(self, **kwargs):
+            captured_kwargs.update(kwargs)
+            return DummyModel()
+
+        def get(self, key, default=None):
+            return getattr(self, key, default)
+
+    class DictLikeMoeConfig:
+        def to_dict(self):
+            return {
+                "activation_checkpointing": True,  # should be stripped
+                "_target_": "some.target",  # should be stripped
+            }
+
+    cfg_model = CapturingModelConfig()
+
+    with patch('nemo_automodel.recipes.vlm.finetune._supports_logits_to_keep', return_value=True):
+        build_model(
+            cfg_model=cfg_model,
+            cfg_freeze=None,
+            cfg_peft=None,
+            seed=123,
+            cfg_moe=DictLikeMoeConfig(),
+            activation_checkpointing=False,
+        )
+
+    assert "moe_config" in captured_kwargs
+    assert isinstance(captured_kwargs["moe_config"], MoEParallelizerConfig)
+    assert captured_kwargs["activation_checkpointing"] is False
+
+
+def test_build_model_no_moe_config_when_cfg_moe_is_none():
+    """Test that moe_config and activation_checkpointing are not in kwargs when cfg_moe is None."""
+    from nemo_automodel._transformers import NeMoAutoModelForImageTextToText
+
+    captured_kwargs = {}
+
+    class CapturingModelConfig:
+        def __init__(self):
+            self._target_ = NeMoAutoModelForImageTextToText.from_pretrained
+
+        def instantiate(self, **kwargs):
+            captured_kwargs.update(kwargs)
+            return DummyModel()
+
+        def get(self, key, default=None):
+            return getattr(self, key, default)
+
+    cfg_model = CapturingModelConfig()
+
+    with patch('nemo_automodel.recipes.vlm.finetune._supports_logits_to_keep', return_value=True):
+        build_model(
+            cfg_model=cfg_model,
+            cfg_freeze=None,
+            cfg_peft=None,
+            seed=123,
+            cfg_moe=None,
+        )
+
+    assert "moe_config" not in captured_kwargs
+    assert "activation_checkpointing" not in captured_kwargs
 
 
 # -----------------------------------------------------------------------------