feat(models): complete tie_word_embeddings guards for #2512 families

nemo_automodel/components/checkpoint/utils.py

@@ -207,6 +207,33 @@ def reject_unsupported_tied_word_embeddings(config: object, model_class_name: st
         )
 
 
+def reject_unsupported_untied_word_embeddings(config: object, model_class_name: str) -> None:
+    """Reject ``tie_word_embeddings=False`` for models whose HF default is tied.
+
+    Tied-by-default architectures share ``lm_head`` with the input embedding and
+    ship checkpoints without a separate ``lm_head.weight``. Honoring
+    ``tie_word_embeddings=False`` would require materializing a distinct
+    ``lm_head`` NeMo does not build (and a tied checkpoint has no weights for it),
+    so reject it explicitly instead of running with a randomly-initialized head.
+
+    The mirror of :func:`reject_unsupported_tied_word_embeddings`; both read the
+    controlling flag via :func:`get_controlling_tie_word_embeddings`.
+
+    Args:
+        config: The model's config.
+        model_class_name: ``type(self).__name__`` of the constructing model.
+
+    Raises:
+        NotImplementedError: if the controlling ``tie_word_embeddings`` flag evaluates to ``False``.
+    """
+    if not get_controlling_tie_word_embeddings(config, model_class_name):
+        raise NotImplementedError(
+            f"{model_class_name} ties its input and output embeddings and does not "
+            f"support tie_word_embeddings=False. The Hugging Face default for this "
+            f"architecture is tied; set tie_word_embeddings=True."
+        )
+
+
 def _normalize_param_name(name: str) -> str:
     """Strip wrapper-specific prefixes from a parameter name."""
     return name.replace("_orig_mod.", "")

nemo_automodel/components/models/gemma4_moe/model.py

@@ -78,6 +78,7 @@ def _make_missing(name: str):
     CausalLMOutputWithPast = _make_missing("CausalLMOutputWithPast")
 
 from nemo_automodel._transformers.model_capabilities import ModelCapabilities
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_untied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig, compute_lm_head_logits
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.utils import cast_model_to_dtype
@@ -925,6 +926,9 @@ def __init__(
     ):
         if not _GEMMA4_HF_AVAILABLE:
             raise UnavailableError("transformers.models.gemma4 is not available.")
+        # Gemma4 is tied by default; untying would need a materialized separate
+        # lm_head NeMo doesn't build, so reject tie_word_embeddings=False up front.
+        reject_unsupported_untied_word_embeddings(config, type(self).__name__)
         backend = backend or BackendConfig()
 
         # Merge text_config overrides (e.g. from YAML) into the proper config

nemo_automodel/components/models/mistral3_vlm/model.py

@@ -43,6 +43,7 @@
     Mistral3ForConditionalGeneration as _HFMistral3ForConditionalGeneration,
 )
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common.utils import compute_lm_head_logits
 from nemo_automodel.components.models.mistral3_vlm.state_dict_adapter import (
     Mistral3FP8StateDictAdapter,
@@ -130,6 +131,9 @@ class ModelCapabilities:
         supports_ep: bool = False
 
     def __init__(self, config: PretrainedConfig):
+        # The supported Mistral3 checkpoint (mistralai/Mistral-Medium-3.5-128B) is
+        # untied (tie_word_embeddings=False), so reject tie_word_embeddings=True.
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         # HF's Mistral3ForConditionalGeneration.__init__ consults
         # ``config.quantization_config`` and swaps nn.Linear → FP8Linear for
         # every language_model Linear. FP8Linear registers a 0-d

nemo_automodel/components/models/nemotron_omni/model.py

@@ -34,6 +34,7 @@
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.utils import cast_model_to_dtype
@@ -303,6 +304,7 @@ def __init__(
         """
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
 
         # ---------------------------------------------------------------

nemo_automodel/components/models/step3p5/model.py

@@ -19,6 +19,7 @@
 import torch.nn as nn
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig, get_rope_config, initialize_linear_module
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.utils import cast_model_to_dtype, compute_lm_head_logits
@@ -450,6 +451,7 @@ def __init__(
     ) -> None:
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
         moe_overrides = kwargs.pop("moe_overrides", None)
         self.model = Step3p5Model(config, backend=self.backend, moe_config=moe_config, moe_overrides=moe_overrides)

nemo_automodel/components/models/step3p7/model.py

@@ -23,6 +23,7 @@
 import torch.nn as nn
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig, initialize_linear_module
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.mtp import roll_tensor
@@ -355,6 +356,7 @@ def __init__(
     ) -> None:
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
         moe_overrides = kwargs.pop("moe_overrides", None)
         mtp_loss_scaling_factor = kwargs.pop("mtp_loss_scaling_factor", 0.1)

tests/unit_tests/models/gemma4_moe/test_gemma4_moe_tied_weights.py

@@ -24,6 +24,7 @@
 Runs on CPU (no CUDA / TE / DeepEP required).
 """
 
+import pytest
 import torch
 import torch.nn as nn
 from transformers.models.gemma4.configuration_gemma4 import Gemma4Config, Gemma4TextConfig
@@ -104,11 +105,10 @@ def test_tied_lm_head_survives_initialize_weights():
     assert lm_head.dtype == torch.bfloat16
 
 
-def test_untied_lm_head_is_separate():
-    """tie_word_embeddings=False: lm_head must keep its own storage."""
-    model = _build(tie_word_embeddings=False)
-    assert model.lm_head.weight is not model.model.language_model.embed_tokens.weight
-    assert model.lm_head.weight.data_ptr() != model.model.language_model.embed_tokens.weight.data_ptr()
+def test_untied_construction_is_rejected():
+    """Gemma4 is tied by default; tie_word_embeddings=False is rejected at construction."""
+    with pytest.raises(NotImplementedError, match="does not support tie_word_embeddings=False"):
+        _build(tie_word_embeddings=False)
 
 
 def test_tie_weights_hook_reties_to_active_embedding():
@@ -127,18 +127,11 @@ def test_tie_weights_hook_reties_to_active_embedding():
     assert model.lm_head.weight is model.model.language_model.embed_tokens.weight
 
 
-def test_tie_weights_hook_is_noop_when_untied():
-    """tie_weights() must not alias storage when the config requests untied embeddings."""
-    model = _build(tie_word_embeddings=False)
-    model.tie_weights()
-    assert model.lm_head.weight is not model.model.language_model.embed_tokens.weight
-
-
 def test_top_level_flag_controls_tie_when_flags_disagree():
     """The controlling flag is top-level Gemma4Config.tie_word_embeddings, not text_config (matches HF)."""
-    # top-level True wins over text_config False -> tied
+    # top-level True wins over text_config False -> tied (constructs, shares storage)
     tied = _build(tie_word_embeddings=True, text_tie=False)
     assert tied.lm_head.weight is tied.model.language_model.embed_tokens.weight
-    # top-level False wins over text_config True -> untied
-    untied = _build(tie_word_embeddings=False, text_tie=True)
-    assert untied.lm_head.weight is not untied.model.language_model.embed_tokens.weight
+    # top-level False is the controlling flag -> untie is rejected even when text_config is True
+    with pytest.raises(NotImplementedError, match="does not support tie_word_embeddings=False"):
+        _build(tie_word_embeddings=False, text_tie=True)

tests/unit_tests/utils/test_checkpoint_utils.py

@@ -178,15 +178,26 @@ def test_reject_unsupported_tied_word_embeddings_omni_wrapper_path():
     """The guard raises for a full Omni wrapper whose thinker_config requests tying."""
     wrapper = SimpleNamespace(thinker_config=SimpleNamespace(tie_word_embeddings=True))
     with pytest.raises(NotImplementedError):
-        checkpoint_utils.reject_unsupported_tied_word_embeddings(
-            wrapper, "Qwen2_5OmniThinkerForConditionalGeneration"
-        )
+        checkpoint_utils.reject_unsupported_tied_word_embeddings(wrapper, "Qwen2_5OmniThinkerForConditionalGeneration")
     wrapper_untied = SimpleNamespace(thinker_config=SimpleNamespace(tie_word_embeddings=False))
     checkpoint_utils.reject_unsupported_tied_word_embeddings(
         wrapper_untied, "Qwen3OmniMoeThinkerForConditionalGeneration"
     )  # no raise
 
 
+def test_reject_unsupported_untied_word_embeddings_raises_when_untied():
+    """A tied-default model with tie_word_embeddings=False is rejected."""
+    config = SimpleNamespace(tie_word_embeddings=False)
+    with pytest.raises(NotImplementedError, match="does not support tie_word_embeddings=False"):
+        checkpoint_utils.reject_unsupported_untied_word_embeddings(config, "Gemma4ForConditionalGeneration")
+
+
+def test_reject_unsupported_untied_word_embeddings_noop_when_tied():
+    """The default (tied) config passes the untie guard without raising."""
+    config = SimpleNamespace(tie_word_embeddings=True)
+    checkpoint_utils.reject_unsupported_untied_word_embeddings(config, "Gemma4ForConditionalGeneration")  # no raise
+
+
 class _DraftLikeModel(nn.Module):
     """Minimal stand-in for an EAGLE-3 draft model.