fix: vlm refac fixes (#1268)

* multiple vlm fixes

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

* revert model init

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

* update tests

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

---------

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

Author: HuiyingLi

examples/vlm_finetune/qwen3/qwen3_omni_moe_30b_te_deepep.yaml

@@ -32,16 +32,17 @@ rng:
 model:
   _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained
   pretrained_model_name_or_path: Qwen/Qwen3-Omni-30B-A3B-Instruct
-  # Customize this backend for fine grained control
-  # backend:
-  #   _target_: nemo_automodel.components.models.common.BackendConfig
-  #   attn: sdpa
-  #   linear: te
-  #   rms_norm: te
-  #   experts: te
-  #   dispatcher: deepep
-  #   fake_balanced_gate: false
-  #   enable_hf_state_dict_adapter: true
+  #Customize this backend for fine grained control
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: sdpa
+    linear: te
+    rms_norm: te
+    rope_fusion: false
+    experts: te
+    enable_deepep: true
+    fake_balanced_gate: false
+    enable_hf_state_dict_adapter: true
 
 
 checkpoint:

examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml

@@ -40,8 +40,9 @@ model:
     attn: sdpa
     linear: te
     rms_norm: te
+    rope_fusion: false
     experts: te
-    dispatcher: deepep
+    enable_deepep: true
     fake_balanced_gate: false
     enable_hf_state_dict_adapter: true
 

nemo_automodel/components/datasets/vlm/collate_fns.py

@@ -308,6 +308,7 @@ def kimi_vl_collate_fn(
         "return_tensors": "pt",
         "padding": True,
         "truncation": True,
+        "add_special_tokens": False,
     }
     if max_length is not None:
         processor_kwargs["max_length"] = max_length

nemo_automodel/recipes/vlm/finetune.py

@@ -88,13 +88,14 @@ def build_model(
     cfg_freeze,
     cfg_peft,
     seed,
-    freeze_embeddings=True,
     cfg_fp8=None,
     cfg_compile=None,
     device_mesh=None,
     moe_mesh=None,
     distributed_config=None,
     pipeline_config=None,
+    cfg_moe=None,
+    activation_checkpointing=False,
 ) -> tuple[nn.Module | AutoPipeline, list["Optimizer"]]:  # noqa: F821
     """Build and initialize a model for VLM.
 
@@ -111,6 +112,20 @@ def build_model(
             "pipeline_config": pipeline_config,
             "freeze_config": cfg_freeze.to_dict() if cfg_freeze is not None else None,
         }
+
+        if cfg_moe is not None:
+            from nemo_automodel.components.moe.config import MoEParallelizerConfig
+
+            if isinstance(cfg_moe, MoEParallelizerConfig):
+                kwargs["moe_config"] = cfg_moe
+            else:
+                moe_dict = cfg_moe.to_dict() if hasattr(cfg_moe, "to_dict") else dict(cfg_moe)
+                # activation_checkpointing is handled separately; strip config keys
+                moe_dict.pop("activation_checkpointing", None)
+                moe_dict.pop("_target_", None)
+                kwargs["moe_config"] = MoEParallelizerConfig(**moe_dict)
+            kwargs["activation_checkpointing"] = activation_checkpointing
+
         if cfg_fp8 is not None:
             fp8_config = build_fp8_config(cfg_fp8)
             kwargs["fp8_config"] = fp8_config
@@ -556,6 +571,8 @@ def setup(self):
             moe_mesh=self.moe_mesh,
             distributed_config=self.distributed_config,
             pipeline_config=self.pipeline_config,
+            cfg_moe=self.dist_setup.moe_config,
+            activation_checkpointing=self.dist_setup.activation_checkpointing,
         )
         self.optimizer = build_optimizer(model, self.cfg.optimizer, self.distributed_config, self.device_mesh)
 

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
 
 
 # -----------------------------------------------------------------------------