feat(models): reject tie_word_embeddings=True on separate-head model families

nemo_automodel/components/checkpoint/utils.py

@@ -105,11 +105,15 @@ def get_controlling_tie_word_embeddings(config: object, model_class_name: str) -
     """Resolve the ``tie_word_embeddings`` flag that actually controls lm_head tying.
 
     HF ties ``lm_head`` based on the *top-level* config flag, not a nested
-    ``text_config`` (verified by construction for Gemma4, Mistral3, and
-    Qwen2.5-Omni under transformers 5.8.1: the top-level flag decides tying
-    regardless of the nested value). So prefer the top-level flag, and only fall
-    back to ``text_config`` for configs that don't expose a top-level
-    ``tie_word_embeddings``.
+    ``text_config`` (verified by construction for Gemma4 and Mistral3 under
+    transformers 5.8.1: the top-level flag decides tying regardless of the nested
+    value). So prefer the top-level flag, and only fall back to ``text_config``
+    for configs that don't expose a top-level ``tie_word_embeddings``.
+
+    Omni "thinker" models are the exception: the full wrapper config
+    (``Qwen2_5OmniConfig`` / ``Qwen3OmniMoeConfig``) does not expose
+    ``tie_word_embeddings`` at the top level at all -- the controlling flag lives
+    on ``config.thinker_config`` -- so unwrap to it for those classes.
 
     Args:
         config: The model's config (or anything exposing ``tie_word_embeddings``
@@ -119,17 +123,29 @@ def get_controlling_tie_word_embeddings(config: object, model_class_name: str) -
     Returns:
         bool: The controlling ``tie_word_embeddings`` value.
     """
-    # Composite models whose top-level / thinker config owns the lm_head tying
+    # Omni "thinker" models: the controlling flag lives on the thinker config.
+    # A full wrapper config (e.g. ``Qwen2_5OmniConfig`` / ``Qwen3OmniMoeConfig``)
+    # does not expose ``tie_word_embeddings`` at the top level at all -- it nests
+    # under ``config.thinker_config`` -- so unwrap to it when present. When the
+    # thinker config itself is passed, ``thinker_config`` is absent and we read
+    # its own top-level flag.
+    omni_thinker_models = (
+        "Qwen2_5OmniThinkerForConditionalGeneration",
+        "Qwen3OmniMoeThinkerForConditionalGeneration",
+    )
+    if any(name in model_class_name for name in omni_thinker_models):
+        thinker_config = getattr(config, "thinker_config", config)
+        return bool(getattr(thinker_config, "tie_word_embeddings", False))
+
+    # Other composite models whose top-level config owns the lm_head tying
     # decision; their nested ``text_config`` flag can disagree and must be ignored.
     # Return the top-level flag (not a forced ``False``) so a constructor guard can
     # still see and reject an unsupported ``top-level=True``. The checkpoint save
     # path stays safe through the storage-based ``has_local_tied_lm_head()`` check,
     # which only drops ``lm_head.weight`` when the tensors actually share storage.
     composite_top_level_models = (
-        "Qwen2_5OmniThinkerForConditionalGeneration",
         "Mistral3FP8VLMForConditionalGeneration",
         "Qwen3VLMoeForConditionalGeneration",
-        "Qwen3OmniMoeThinkerForConditionalGeneration",
     )
     if any(name in model_class_name for name in composite_top_level_models):
         return bool(getattr(config, "tie_word_embeddings", False))
@@ -163,6 +179,34 @@ def is_tied_word_embeddings(model: nn.Module) -> bool:
     return get_controlling_tie_word_embeddings(config, type(model).__name__)
 
 
+def reject_unsupported_tied_word_embeddings(config: object, model_class_name: str) -> None:
+    """Reject ``tie_word_embeddings=True`` for models whose HF default is untied.
+
+    Separate-head architectures (HF default: distinct input/output embeddings)
+    don't build a shared ``lm_head``, so honoring ``tie_word_embeddings=True``
+    would silently leave a randomly-initialized head or require materializing a
+    tied weight NeMo does not support. Reject it explicitly with a clear message
+    instead of pretending to support it.
+
+    Uses :func:`get_controlling_tie_word_embeddings`, so composite VLM/omni configs
+    are read from the controlling top-level flag rather than a nested
+    ``text_config``.
+
+    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 is set.
+    """
+    if get_controlling_tie_word_embeddings(config, model_class_name):
+        raise NotImplementedError(
+            f"{model_class_name} has separate input and output embeddings and does not "
+            f"support tie_word_embeddings=True. The Hugging Face default for this "
+            f"architecture is untied; set tie_word_embeddings=False."
+        )
+
+
 def _normalize_param_name(name: str) -> str:
     """Strip wrapper-specific prefixes from a parameter name."""
     return name.replace("_orig_mod.", "")

nemo_automodel/components/models/deepseek_v3/model.py

@@ -20,6 +20,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import (
     BackendConfig,
     get_rope_config,
@@ -299,6 +300,7 @@ def __init__(
     ):
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
         # The HF DeepSeek-V3 reference computes router scoring in fp32; routing is highly
         # precision-sensitive (small bf16 errors flip expert selection) and the gate is tiny,

nemo_automodel/components/models/deepseek_v32/model.py

@@ -26,6 +26,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,
     compute_lm_head_logits,
@@ -208,6 +209,7 @@ def __init__(
         from nemo_automodel.components.models.common import initialize_linear_module
 
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
         # Use V3.2 Model instead of V3 Model
         moe_overrides = kwargs.pop("moe_overrides", None)

nemo_automodel/components/models/deepseek_v4/model.py

@@ -49,6 +49,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,
@@ -648,6 +649,7 @@ def __init__(
     ):
         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)

nemo_automodel/components/models/glm4_moe/model.py

@@ -20,6 +20,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.glm4_moe.configuration_glm4_moe import Glm4MoeConfig
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import (
     BackendConfig,
     get_rope_config,
@@ -272,6 +273,7 @@ def __init__(
     ):
         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 = Glm4MoeModel(

nemo_automodel/components/models/glm4_moe_lite/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.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.utils import (
     BackendConfig,
@@ -274,6 +275,7 @@ def __init__(
     ):
         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 = Glm4MoeLiteModel(

nemo_automodel/components/models/glm_moe_dsa/model.py

@@ -20,6 +20,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import (
     BackendConfig,
     compute_lm_head_logits,
@@ -294,6 +295,7 @@ def __init__(
     ):
         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 = GlmMoeDsaModel(

nemo_automodel/components/models/gpt_oss/model.py

@@ -21,6 +21,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import (
     BackendConfig,
     get_rope_config,
@@ -260,6 +261,7 @@ def __init__(
     ):
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig(attn="flex")
         moe_overrides = kwargs.pop("moe_overrides", None)
         self.model = GptOssModel(config, backend=self.backend, moe_config=moe_config, moe_overrides=moe_overrides)

nemo_automodel/components/models/hy_mt2/model.py

@@ -43,6 +43,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,
@@ -325,6 +326,7 @@ def __init__(
     ):
         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 = HyMT2Model(config, backend=self.backend, moe_config=moe_config, moe_overrides=moe_overrides)

nemo_automodel/components/models/hy_v3/model.py

@@ -30,6 +30,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,
@@ -262,6 +263,7 @@ def __init__(
     ):
         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 = HYV3Model(config, backend=self.backend, moe_config=moe_config, moe_overrides=moe_overrides)

nemo_automodel/components/models/kimi_k25_vl/model.py

@@ -145,6 +145,7 @@ def to_dict(self) -> Dict[str, Any]:
         return output
 
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig, compute_lm_head_logits, initialize_linear_module
 from nemo_automodel.components.models.deepseek_v3.model import DeepseekV3Model
 from nemo_automodel.components.models.deepseek_v3.rope_utils import freqs_cis_from_position_ids
@@ -951,6 +952,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *model_args, **kwar
     def __init__(self, config, moe_config: MoEConfig | None = None, backend: BackendConfig | None = None, **kwargs):
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
 
         self.model = KimiK25VLModel(config, moe_config=moe_config, backend=self.backend)

nemo_automodel/components/models/kimivl/model.py

@@ -107,6 +107,7 @@ def to_dict(self) -> Dict[str, Any]:
         return output
 
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig, compute_lm_head_logits, initialize_linear_module
 from nemo_automodel.components.models.deepseek_v3.model import DeepseekV3Model
 from nemo_automodel.components.models.deepseek_v3.rope_utils import freqs_cis_from_position_ids
@@ -657,6 +658,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *model_args, **kwar
     def __init__(self, config, moe_config: MoEConfig | None = None, backend: BackendConfig | None = None, **kwargs):
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
 
         self.model = KimiVLModel(config, moe_config=moe_config, backend=self.backend)

nemo_automodel/components/models/ling_v2/model.py

@@ -42,6 +42,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from nemo_automodel._transformers.model_capabilities import ModelCapabilities
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import (
     BackendConfig,
     initialize_linear_module,
@@ -342,6 +343,7 @@ def __init__(
     ):
         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 = BailingMoeV2Model(

nemo_automodel/components/models/llava_onevision/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.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.utils import compute_lm_head_logits
 from nemo_automodel.components.models.llava_onevision.rice_vit import RiceTransformer
@@ -313,6 +314,7 @@ def __init__(
     ):
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         if attn_implementation is None:
             attn_implementation = getattr(config, "_attn_implementation", None) or "eager"
         self.model = LLaVAOneVision1_5_Model(config, attn_implementation=attn_implementation)

nemo_automodel/components/models/mimo_v2_flash/model.py

@@ -23,6 +23,7 @@
 from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
 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,
@@ -630,6 +631,7 @@ def __init__(
     ):
         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 = MiMoV2FlashModel(

nemo_automodel/components/models/minimax_m2/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,
@@ -272,6 +273,7 @@ def __init__(
     ):
         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 = MiniMaxM2Model(

nemo_automodel/components/models/minimax_m3_vl/model.py

@@ -26,6 +26,7 @@
 import torch
 import torch.nn as nn
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import (
     BackendConfig,
     get_rope_config,
@@ -278,6 +279,7 @@ def __init__(
     ):
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
         self.model = MiniMaxM3TextModel(config, backend=self.backend, moe_config=moe_config)
         self.lm_head = initialize_linear_module(self.backend.linear, config.hidden_size, config.vocab_size, bias=False)
@@ -423,6 +425,7 @@ def __init__(
     ):
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         text_config = config.text_config
         self.backend = backend or BackendConfig()
         self.model = MiniMaxM3TextModel(text_config, backend=self.backend, moe_config=moe_config)

nemo_automodel/components/models/mistral4/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,
     compute_lm_head_logits,
@@ -345,6 +346,9 @@ def __init__(
         **kwargs,
     ):
         super().__init__()
+        # Reject an unsupported tied request on the controlling top-level flag
+        # before unwrapping to text_config below.
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         # Extract text_config if this is a multimodal wrapper config
         config = getattr(config, "text_config", config)
         self.config = config

nemo_automodel/components/models/nemotron_parse/model.py

@@ -38,6 +38,7 @@
     MBartScaledWordEmbedding,
 )
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.utils import compute_lm_head_logits
 
@@ -444,6 +445,7 @@ class ModelCapabilities:
         supports_ep: bool = False
 
     def __init__(self, config: NemotronParseConfig, loss_fn=None, **kwargs):
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         super().__init__(config)
         self.loss_fn = loss_fn