From 702b9c5e44d97b93c486767fb74df55a515da20d Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Fri, 26 Jun 2026 22:30:21 +0400
Subject: [PATCH 1/3] feat(models): reject tie_word_embeddings=True on
 separate-head model families

Add reject_unsupported_tied_word_embeddings() (built on the #2732 resolver) and wire it into the __init__ of the 25 verified untied-default model classes, so setting tie_word_embeddings=True raises a clear error instead of a silently-untied head. Excludes mistral3_vlm (HF default tied) and step3p5/step3p7/nemotron_omni (pending hub verification).

Refs #2512

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 nemo_automodel/components/checkpoint/utils.py | 28 ++++++++++
 .../components/models/deepseek_v3/model.py    |  2 +
 .../components/models/deepseek_v32/model.py   |  2 +
 .../components/models/deepseek_v4/model.py    |  2 +
 .../components/models/glm4_moe/model.py       |  2 +
 .../components/models/glm4_moe_lite/model.py  |  2 +
 .../components/models/glm_moe_dsa/model.py    |  2 +
 .../components/models/gpt_oss/model.py        |  2 +
 .../components/models/hy_mt2/model.py         |  2 +
 .../components/models/hy_v3/model.py          |  2 +
 .../components/models/kimi_k25_vl/model.py    |  2 +
 .../components/models/kimivl/model.py         |  2 +
 .../components/models/ling_v2/model.py        |  2 +
 .../models/llava_onevision/model.py           |  2 +
 .../components/models/mimo_v2_flash/model.py  |  2 +
 .../components/models/minimax_m2/model.py     |  2 +
 .../components/models/minimax_m3_vl/model.py  |  3 +
 .../components/models/mistral4/model.py       |  2 +
 .../components/models/nemotron_parse/model.py |  2 +
 .../components/models/nemotron_v3/model.py    |  2 +
 .../components/models/qwen2_5_omni/model.py   |  2 +
 .../components/models/qwen3_5_moe/model.py    |  2 +
 .../components/models/qwen3_moe/model.py      |  2 +
 .../components/models/qwen3_next/model.py     |  2 +
 .../components/models/qwen3_omni_moe/model.py |  2 +
 .../components/models/qwen3_vl_moe/model.py   |  2 +
 .../qwen3_moe/test_qwen3_moe_tie_guard.py     | 56 +++++++++++++++++++
 .../unit_tests/utils/test_checkpoint_utils.py | 27 +++++++++
 28 files changed, 162 insertions(+)
 create mode 100644 tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py

diff --git a/nemo_automodel/components/checkpoint/utils.py b/nemo_automodel/components/checkpoint/utils.py
index 0001a274b5..4f542627d3 100644
--- a/nemo_automodel/components/checkpoint/utils.py
+++ b/nemo_automodel/components/checkpoint/utils.py
@@ -163,6 +163,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.", "")
diff --git a/nemo_automodel/components/models/deepseek_v3/model.py b/nemo_automodel/components/models/deepseek_v3/model.py
index 128710777d..fd1b4e76ed 100644
--- a/nemo_automodel/components/models/deepseek_v3/model.py
+++ b/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,
diff --git a/nemo_automodel/components/models/deepseek_v32/model.py b/nemo_automodel/components/models/deepseek_v32/model.py
index ea52443f49..0f8f437c2f 100644
--- a/nemo_automodel/components/models/deepseek_v32/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py
index f8dcd73d29..faeb985317 100644
--- a/nemo_automodel/components/models/deepseek_v4/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/glm4_moe/model.py b/nemo_automodel/components/models/glm4_moe/model.py
index f5b55a937a..8db52501f0 100644
--- a/nemo_automodel/components/models/glm4_moe/model.py
+++ b/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(
diff --git a/nemo_automodel/components/models/glm4_moe_lite/model.py b/nemo_automodel/components/models/glm4_moe_lite/model.py
index 4a15238bfe..20be465d45 100644
--- a/nemo_automodel/components/models/glm4_moe_lite/model.py
+++ b/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(
diff --git a/nemo_automodel/components/models/glm_moe_dsa/model.py b/nemo_automodel/components/models/glm_moe_dsa/model.py
index 96b9b9b89b..7cc2b498e3 100644
--- a/nemo_automodel/components/models/glm_moe_dsa/model.py
+++ b/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(
diff --git a/nemo_automodel/components/models/gpt_oss/model.py b/nemo_automodel/components/models/gpt_oss/model.py
index d108e02e38..73d967e915 100644
--- a/nemo_automodel/components/models/gpt_oss/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/hy_mt2/model.py b/nemo_automodel/components/models/hy_mt2/model.py
index 9816d27da4..c754bd897e 100644
--- a/nemo_automodel/components/models/hy_mt2/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/hy_v3/model.py b/nemo_automodel/components/models/hy_v3/model.py
index 2dc1a5171c..09d49e1860 100644
--- a/nemo_automodel/components/models/hy_v3/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/kimi_k25_vl/model.py b/nemo_automodel/components/models/kimi_k25_vl/model.py
index bf07c3d03e..5fd5d95dcf 100644
--- a/nemo_automodel/components/models/kimi_k25_vl/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/kimivl/model.py b/nemo_automodel/components/models/kimivl/model.py
index f2c50c1611..15177f25e8 100644
--- a/nemo_automodel/components/models/kimivl/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/ling_v2/model.py b/nemo_automodel/components/models/ling_v2/model.py
index 06b33b8543..862acc9677 100644
--- a/nemo_automodel/components/models/ling_v2/model.py
+++ b/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(
diff --git a/nemo_automodel/components/models/llava_onevision/model.py b/nemo_automodel/components/models/llava_onevision/model.py
index 1562997b6f..fb483d1a68 100644
--- a/nemo_automodel/components/models/llava_onevision/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/mimo_v2_flash/model.py b/nemo_automodel/components/models/mimo_v2_flash/model.py
index bccd125048..d4a3278d9c 100644
--- a/nemo_automodel/components/models/mimo_v2_flash/model.py
+++ b/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(
diff --git a/nemo_automodel/components/models/minimax_m2/model.py b/nemo_automodel/components/models/minimax_m2/model.py
index 25df209d1a..2b48ba5e8c 100644
--- a/nemo_automodel/components/models/minimax_m2/model.py
+++ b/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(
diff --git a/nemo_automodel/components/models/minimax_m3_vl/model.py b/nemo_automodel/components/models/minimax_m3_vl/model.py
index 20b99ef7dd..1ff4610139 100644
--- a/nemo_automodel/components/models/minimax_m3_vl/model.py
+++ b/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)
diff --git a/nemo_automodel/components/models/mistral4/model.py b/nemo_automodel/components/models/mistral4/model.py
index 291dbb8e67..77b40f80f9 100644
--- a/nemo_automodel/components/models/mistral4/model.py
+++ b/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,
@@ -348,6 +349,7 @@ def __init__(
         # Extract text_config if this is a multimodal wrapper config
         config = getattr(config, "text_config", config)
         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 = Mistral4Model(
diff --git a/nemo_automodel/components/models/nemotron_parse/model.py b/nemo_automodel/components/models/nemotron_parse/model.py
index ad2f49eba4..8efb878528 100644
--- a/nemo_automodel/components/models/nemotron_parse/model.py
+++ b/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
 
@@ -445,6 +446,7 @@ class ModelCapabilities:
 
     def __init__(self, config: NemotronParseConfig, loss_fn=None, **kwargs):
         super().__init__(config)
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.loss_fn = loss_fn
 
         self.encoder = RadioWithNeck(config.encoder)
diff --git a/nemo_automodel/components/models/nemotron_v3/model.py b/nemo_automodel/components/models/nemotron_v3/model.py
index c6df69aa50..94c078d3fb 100644
--- a/nemo_automodel/components/models/nemotron_v3/model.py
+++ b/nemo_automodel/components/models/nemotron_v3/model.py
@@ -22,6 +22,7 @@
 from transformers.generation import GenerationConfig, GenerationMixin
 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,
     HFCheckpointingMixin,
@@ -370,6 +371,7 @@ def __init__(
         """
         super().__init__()
         self.config = config
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
 
         # Base model
diff --git a/nemo_automodel/components/models/qwen2_5_omni/model.py b/nemo_automodel/components/models/qwen2_5_omni/model.py
index 5688aeff7f..efab0d3033 100644
--- a/nemo_automodel/components/models/qwen2_5_omni/model.py
+++ b/nemo_automodel/components/models/qwen2_5_omni/model.py
@@ -44,6 +44,7 @@
     Qwen2_5OmniThinkerForConditionalGeneration as HFQwen2_5OmniThinkerForConditionalGeneration,
 )
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_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.qwen2_5_omni.state_dict_adapter import Qwen2_5OmniStateDictAdapter
@@ -102,6 +103,7 @@ def __init__(
     ):
         thinker_config = _resolve_thinker_config(config)
         super().__init__(thinker_config)
+        reject_unsupported_tied_word_embeddings(self.config, type(self).__name__)
 
         # HF Qwen2.5-Omni declares ``audio_tower.audio_bos_eos_token`` as an
         # ``nn.Embedding(2, output_dim)`` (modeling_qwen2_5_omni.py:751) but it
diff --git a/nemo_automodel/components/models/qwen3_5_moe/model.py b/nemo_automodel/components/models/qwen3_5_moe/model.py
index 389417fe06..31abb19911 100644
--- a/nemo_automodel/components/models/qwen3_5_moe/model.py
+++ b/nemo_automodel/components/models/qwen3_5_moe/model.py
@@ -60,6 +60,7 @@ def _make_missing(name: str):
     Qwen3_5MoeVisionRotaryEmbedding = _make_missing("Qwen3_5MoeVisionRotaryEmbedding")
     HFQwen3_5MoeModel = _make_missing("Qwen3_5MoeModel")
 
+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 MTPConfig, MTPModule, roll_tensor
@@ -755,6 +756,7 @@ def __init__(
 
         # Initialize HF parent (creates self.model, self.lm_head, vision encoder, etc.)
         super().__init__(config)
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
 
         self.backend = backend
 
diff --git a/nemo_automodel/components/models/qwen3_moe/model.py b/nemo_automodel/components/models/qwen3_moe/model.py
index 3c7845af3c..d2fb4348d4 100644
--- a/nemo_automodel/components/models/qwen3_moe/model.py
+++ b/nemo_automodel/components/models/qwen3_moe/model.py
@@ -20,6 +20,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import (
     BackendConfig,
     get_rope_config,
@@ -280,6 +281,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 = Qwen3MoeModel(config, backend=self.backend, moe_config=moe_config, moe_overrides=moe_overrides)
diff --git a/nemo_automodel/components/models/qwen3_next/model.py b/nemo_automodel/components/models/qwen3_next/model.py
index d9e8c6be15..dc044af07d 100644
--- a/nemo_automodel/components/models/qwen3_next/model.py
+++ b/nemo_automodel/components/models/qwen3_next/model.py
@@ -20,6 +20,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.qwen3_next.configuration_qwen3_next import Qwen3NextConfig
 
+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()
         moe_overrides = kwargs.pop("moe_overrides", None)
         self.model = Qwen3NextModel(config, backend=self.backend, moe_config=moe_config, moe_overrides=moe_overrides)
diff --git a/nemo_automodel/components/models/qwen3_omni_moe/model.py b/nemo_automodel/components/models/qwen3_omni_moe/model.py
index 9749493672..dc65fc39f9 100644
--- a/nemo_automodel/components/models/qwen3_omni_moe/model.py
+++ b/nemo_automodel/components/models/qwen3_omni_moe/model.py
@@ -29,6 +29,7 @@
     Qwen3OmniMoeThinkerTextRotaryEmbedding as HFQwen3OmniMoeThinkerTextRotaryEmbedding,
 )
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig, initialize_linear_module, initialize_rms_norm_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
@@ -257,6 +258,7 @@ def __init__(
     ):
         base_config = config.thinker_config if hasattr(config, "thinker_config") else config
         backend = backend or BackendConfig()
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
 
         # _init_model() only overrides the top-level hf_config.torch_dtype; for
         # Omni configs the real params live under thinker_config.text_config /
diff --git a/nemo_automodel/components/models/qwen3_vl_moe/model.py b/nemo_automodel/components/models/qwen3_vl_moe/model.py
index 6058110d91..922cf8349c 100644
--- a/nemo_automodel/components/models/qwen3_vl_moe/model.py
+++ b/nemo_automodel/components/models/qwen3_vl_moe/model.py
@@ -30,6 +30,7 @@
     Qwen3VLMoeVisionRotaryEmbedding,
 )
 
+from nemo_automodel.components.checkpoint.utils import reject_unsupported_tied_word_embeddings
 from nemo_automodel.components.models.common import BackendConfig, initialize_linear_module, initialize_rms_norm_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
@@ -500,6 +501,7 @@ def __init__(
                     sub_cfg.torch_dtype = top_dtype
 
         super().__init__(config)
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
 
         self.backend = backend
         self.model.__class__ = Qwen3VLMoeModel
diff --git a/tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py b/tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py
new file mode 100644
index 0000000000..616b2fdc74
--- /dev/null
+++ b/tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py
@@ -0,0 +1,56 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""CPU guard test: separate-head models reject tie_word_embeddings=True.
+
+The reject guard runs at the very start of ``__init__`` (before any device- or
+kernel-dependent construction), so this is CPU-safe even though the full
+qwen3_moe model build requires a GPU. qwen3_moe stands in for the whole
+untied-default family wired through ``reject_unsupported_tied_word_embeddings``.
+"""
+
+import pytest
+from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
+
+from nemo_automodel.components.models.common import BackendConfig
+from nemo_automodel.components.models.qwen3_moe.model import Qwen3MoeForCausalLM
+
+
+def _tiny_config(tie_word_embeddings: bool) -> Qwen3MoeConfig:
+    return Qwen3MoeConfig(
+        vocab_size=256,
+        hidden_size=64,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        num_hidden_layers=2,
+        intermediate_size=128,
+        moe_intermediate_size=64,
+        num_experts=4,
+        num_experts_per_tok=2,
+        decoder_sparse_step=1,
+        max_position_embeddings=256,
+        rms_norm_eps=1e-6,
+        rope_theta=5000.0,
+        router_aux_loss_coef=0.01,
+        use_sliding_window=False,
+        tie_word_embeddings=tie_word_embeddings,
+    )
+
+
+def test_qwen3_moe_rejects_tied_word_embeddings():
+    """Constructing with tie_word_embeddings=True raises a clear error before model build."""
+    backend = BackendConfig(linear="torch", attn="sdpa", rms_norm="torch", experts="torch", dispatcher="torch")
+    with pytest.raises(NotImplementedError, match="does not support tie_word_embeddings=True"):
+        Qwen3MoeForCausalLM(_tiny_config(tie_word_embeddings=True), backend=backend)
diff --git a/tests/unit_tests/utils/test_checkpoint_utils.py b/tests/unit_tests/utils/test_checkpoint_utils.py
index 9ef3aad092..c9f7018c47 100644
--- a/tests/unit_tests/utils/test_checkpoint_utils.py
+++ b/tests/unit_tests/utils/test_checkpoint_utils.py
@@ -14,6 +14,7 @@
 
 from types import SimpleNamespace
 
+import pytest
 import torch.nn as nn
 
 import nemo_automodel.components.checkpoint.utils as checkpoint_utils
@@ -125,6 +126,32 @@ def get_text_config(self):
     assert checkpoint_utils.get_controlling_tie_word_embeddings(_NoTopFlag(), "SomeForCausalLM") is True
 
 
+def test_reject_unsupported_tied_word_embeddings_raises_when_tied():
+    """A separate-head model with tie_word_embeddings=True is rejected."""
+    config = SimpleNamespace(tie_word_embeddings=True)
+    with pytest.raises(NotImplementedError, match="does not support tie_word_embeddings=True"):
+        checkpoint_utils.reject_unsupported_tied_word_embeddings(config, "Qwen3MoeForCausalLM")
+
+
+def test_reject_unsupported_tied_word_embeddings_noop_when_untied():
+    """The default (untied) config passes the guard without raising."""
+    config = SimpleNamespace(tie_word_embeddings=False)
+    checkpoint_utils.reject_unsupported_tied_word_embeddings(config, "Qwen3MoeForCausalLM")  # no raise
+
+
+def test_reject_unsupported_tied_word_embeddings_uses_top_level_for_composite():
+    """Composite VLM/omni configs read the controlling top-level flag, not nested text_config."""
+    # top-level False (even with nested text True) -> not tied -> no raise
+    untied = SimpleNamespace(
+        tie_word_embeddings=False, get_text_config=lambda: SimpleNamespace(tie_word_embeddings=True)
+    )
+    checkpoint_utils.reject_unsupported_tied_word_embeddings(untied, "Qwen3VLMoeForConditionalGeneration")
+    # top-level True -> tied -> raise
+    tied = SimpleNamespace(tie_word_embeddings=True, get_text_config=lambda: SimpleNamespace(tie_word_embeddings=False))
+    with pytest.raises(NotImplementedError):
+        checkpoint_utils.reject_unsupported_tied_word_embeddings(tied, "Qwen3VLMoeForConditionalGeneration")
+
+
 class _DraftLikeModel(nn.Module):
     """Minimal stand-in for an EAGLE-3 draft model.
 

From 01b18b02b23117377dc9781dbfc1a9aa6cab25fe Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Fri, 26 Jun 2026 22:49:28 +0400
Subject: [PATCH 2/3] fix(models): run tie guard on top-level config before
 unwrap/super (review)

Per review: in composite models the guard was reading the unwrapped text_config/thinker_config, which could miss a top-level tie_word_embeddings=True. Run it on the original config before any unwrap or super().__init__(), which also fails fast. Refs #2512

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 nemo_automodel/components/models/mistral4/model.py       | 4 +++-
 nemo_automodel/components/models/nemotron_parse/model.py | 2 +-
 nemo_automodel/components/models/qwen2_5_omni/model.py   | 4 +++-
 nemo_automodel/components/models/qwen3_5_moe/model.py    | 2 +-
 nemo_automodel/components/models/qwen3_vl_moe/model.py   | 2 +-
 5 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/nemo_automodel/components/models/mistral4/model.py b/nemo_automodel/components/models/mistral4/model.py
index 77b40f80f9..bffa30907c 100644
--- a/nemo_automodel/components/models/mistral4/model.py
+++ b/nemo_automodel/components/models/mistral4/model.py
@@ -346,10 +346,12 @@ 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
-        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         self.backend = backend or BackendConfig()
         moe_overrides = kwargs.pop("moe_overrides", None)
         self.model = Mistral4Model(
diff --git a/nemo_automodel/components/models/nemotron_parse/model.py b/nemo_automodel/components/models/nemotron_parse/model.py
index 8efb878528..6a06015677 100644
--- a/nemo_automodel/components/models/nemotron_parse/model.py
+++ b/nemo_automodel/components/models/nemotron_parse/model.py
@@ -445,8 +445,8 @@ class ModelCapabilities:
         supports_ep: bool = False
 
     def __init__(self, config: NemotronParseConfig, loss_fn=None, **kwargs):
-        super().__init__(config)
         reject_unsupported_tied_word_embeddings(config, type(self).__name__)
+        super().__init__(config)
         self.loss_fn = loss_fn
 
         self.encoder = RadioWithNeck(config.encoder)
diff --git a/nemo_automodel/components/models/qwen2_5_omni/model.py b/nemo_automodel/components/models/qwen2_5_omni/model.py
index efab0d3033..78cb5ef5e8 100644
--- a/nemo_automodel/components/models/qwen2_5_omni/model.py
+++ b/nemo_automodel/components/models/qwen2_5_omni/model.py
@@ -101,9 +101,11 @@ def __init__(
         backend: BackendConfig | None = None,
         **kwargs,
     ):
+        # Check the controlling top-level flag on the original config before
+        # resolving to thinker_config and building the HF parent.
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         thinker_config = _resolve_thinker_config(config)
         super().__init__(thinker_config)
-        reject_unsupported_tied_word_embeddings(self.config, type(self).__name__)
 
         # HF Qwen2.5-Omni declares ``audio_tower.audio_bos_eos_token`` as an
         # ``nn.Embedding(2, output_dim)`` (modeling_qwen2_5_omni.py:751) but it
diff --git a/nemo_automodel/components/models/qwen3_5_moe/model.py b/nemo_automodel/components/models/qwen3_5_moe/model.py
index 31abb19911..0d588be584 100644
--- a/nemo_automodel/components/models/qwen3_5_moe/model.py
+++ b/nemo_automodel/components/models/qwen3_5_moe/model.py
@@ -754,9 +754,9 @@ def __init__(
                 if sub_cfg is not config and hasattr(sub_cfg, "torch_dtype"):
                     sub_cfg.torch_dtype = top_dtype
 
+        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
         # Initialize HF parent (creates self.model, self.lm_head, vision encoder, etc.)
         super().__init__(config)
-        reject_unsupported_tied_word_embeddings(config, type(self).__name__)
 
         self.backend = backend
 
diff --git a/nemo_automodel/components/models/qwen3_vl_moe/model.py b/nemo_automodel/components/models/qwen3_vl_moe/model.py
index 922cf8349c..514b82716c 100644
--- a/nemo_automodel/components/models/qwen3_vl_moe/model.py
+++ b/nemo_automodel/components/models/qwen3_vl_moe/model.py
@@ -500,8 +500,8 @@ def __init__(
                 if sub_cfg is not config and hasattr(sub_cfg, "torch_dtype"):
                     sub_cfg.torch_dtype = top_dtype
 
-        super().__init__(config)
         reject_unsupported_tied_word_embeddings(config, type(self).__name__)
+        super().__init__(config)
 
         self.backend = backend
         self.model.__class__ = Qwen3VLMoeModel

From ad7ed14b05776fa46206e4f767c5bb3352ee3d33 Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Sat, 27 Jun 2026 00:11:33 +0400
Subject: [PATCH 3/3] fix(checkpoint): resolve Omni wrapper tie flag from
 thinker_config (review)

Qwen2_5OmniConfig/Qwen3OmniMoeConfig don't expose tie_word_embeddings at the top level; the controlling flag is on config.thinker_config. Unwrap to it in the resolver so the constructor guard catches a tied request via the full wrapper config. Add wrapper-path unit coverage, clarify the resolver docstring, and fix the copyright year to 2026. Refs #2512

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 nemo_automodel/components/checkpoint/utils.py | 32 ++++++++++++-----
 .../qwen3_moe/test_qwen3_moe_tie_guard.py     |  2 +-
 .../unit_tests/utils/test_checkpoint_utils.py | 35 +++++++++++++++++++
 3 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/nemo_automodel/components/checkpoint/utils.py b/nemo_automodel/components/checkpoint/utils.py
index 4f542627d3..da2cef5cc6 100644
--- a/nemo_automodel/components/checkpoint/utils.py
+++ b/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))
diff --git a/tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py b/tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py
index 616b2fdc74..d8b1be33dc 100644
--- a/tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py
+++ b/tests/unit_tests/models/qwen3_moe/test_qwen3_moe_tie_guard.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/tests/unit_tests/utils/test_checkpoint_utils.py b/tests/unit_tests/utils/test_checkpoint_utils.py
index c9f7018c47..decd7aa819 100644
--- a/tests/unit_tests/utils/test_checkpoint_utils.py
+++ b/tests/unit_tests/utils/test_checkpoint_utils.py
@@ -152,6 +152,41 @@ def test_reject_unsupported_tied_word_embeddings_uses_top_level_for_composite():
         checkpoint_utils.reject_unsupported_tied_word_embeddings(tied, "Qwen3VLMoeForConditionalGeneration")
 
 
+def test_get_controlling_tie_word_embeddings_omni_wrapper_reads_thinker_config():
+    """Full Omni wrapper config nests the controlling flag under thinker_config.
+
+    Qwen2_5OmniConfig / Qwen3OmniMoeConfig do not expose tie_word_embeddings at the
+    top level; the controlling flag lives on config.thinker_config.
+    """
+    wrapper_tied = SimpleNamespace(thinker_config=SimpleNamespace(tie_word_embeddings=True))
+    wrapper_untied = SimpleNamespace(thinker_config=SimpleNamespace(tie_word_embeddings=False))
+    for cls in (
+        "Qwen2_5OmniThinkerForConditionalGeneration",
+        "Qwen3OmniMoeThinkerForConditionalGeneration",
+    ):
+        assert checkpoint_utils.get_controlling_tie_word_embeddings(wrapper_tied, cls) is True
+        assert checkpoint_utils.get_controlling_tie_word_embeddings(wrapper_untied, cls) is False
+    # When the thinker config itself is passed (no nested thinker_config), read its own flag.
+    direct = SimpleNamespace(tie_word_embeddings=True)
+    assert (
+        checkpoint_utils.get_controlling_tie_word_embeddings(direct, "Qwen2_5OmniThinkerForConditionalGeneration")
+        is True
+    )
+
+
+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"
+        )
+    wrapper_untied = SimpleNamespace(thinker_config=SimpleNamespace(tie_word_embeddings=False))
+    checkpoint_utils.reject_unsupported_tied_word_embeddings(
+        wrapper_untied, "Qwen3OmniMoeThinkerForConditionalGeneration"
+    )  # no raise
+
+
 class _DraftLikeModel(nn.Module):
     """Minimal stand-in for an EAGLE-3 draft model.