NVIDIA-NeMo · chtruong814 · Jul 16, 2026 · Jul 9, 2026 · Jul 9, 2026 · Jul 9, 2026
@@ -27,6 +27,7 @@
 import gc
 import inspect
 import logging
+import os
 from contextlib import nullcontext
 from typing import TYPE_CHECKING, List, Optional, Union
 
@@ -96,6 +97,7 @@
     no_hf_meta_device,
     resolve_sdpa_method,
 )
+from nemo_automodel.components.models.common.tie_word_embeddings import reject_tie_word_embeddings_flip
 
 if not hasattr(_gen_utils, "NEED_SETUP_CACHE_CLASSES_MAPPING"):
     from transformers.cache_utils import StaticCache
@@ -254,6 +256,42 @@ def _maybe_dequantize_fp8_for_peft(hf_native_quant_cfg, peft_config, pretrained_
     return False
 
 
+def _maybe_reject_tie_word_embeddings_flip(pretrained_model_name_or_path, hf_config, kwargs):
+    """Reject a from_pretrained request that flips tie_word_embeddings from the checkpoint.
+
+    Re-reads the checkpoint's raw config (no user value-overrides) and compares its
+    controlling tie flag to the requested ``hf_config`` via
+    :func:`reject_tie_word_embeddings_flip`. Conservative by design: path-like sources
+    are normalized with :func:`os.fspath`, non-path sources are skipped, and it silently
+    returns if the raw config cannot be re-read, so it never blocks a load except on a
+    genuine flip.
+
+    Args:
+        pretrained_model_name_or_path: The from_pretrained source (``str`` and
+            ``os.PathLike`` are checked; anything else is skipped).
+        hf_config: The resolved config with user overrides applied (the requested value).
+        kwargs: The from_pretrained kwargs (hub-locating keys are reused for the raw load).
+    """
+    if isinstance(pretrained_model_name_or_path, os.PathLike):
+        pretrained_model_name_or_path = os.fspath(pretrained_model_name_or_path)
+    if not isinstance(pretrained_model_name_or_path, str):
+        # Non-path source (e.g. bytes fspath or preloaded object): nothing to re-read.
+        return
+    hub_kwargs = {k: kwargs[k] for k in _AUTO_CONFIG_HUB_KWARG_KEYS if k in kwargs}
+    try:
+        raw_config = AutoConfig.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=kwargs.get("trust_remote_code", resolve_trust_remote_code(pretrained_model_name_or_path)),
+            **hub_kwargs,
+        )
+    except Exception:
+        # Cannot re-read the raw config (offline / custom loader); do not block the load.
+        return
+    architectures = getattr(hf_config, "architectures", None) or []
+    model_class_name = architectures[0] if architectures else type(hf_config).__name__
+    reject_tie_word_embeddings_flip(raw_config, hf_config, model_class_name)
+
+
 class _BaseNeMoAutoModelClass(_BaseAutoModelClass):
     """
     Drop-in replacement for ``_BaseAutoModelClass`` that includes custom-kernels.
@@ -710,6 +748,10 @@ def from_pretrained(
                 raise
         is_hf_model = get_is_hf_model(hf_config, force_hf)
 
+        # Layer 2: reject loading a checkpoint with tie_word_embeddings flipped from the
+        # value it was saved with (the class-level TieSupport policy cannot catch this).
+        _maybe_reject_tie_word_embeddings_flip(pretrained_model_name_or_path, hf_config, kwargs)
+
         sdpa_method = resolve_sdpa_method(sdpa_method, mesh.device_mesh, activation_checkpointing)
 
         return cls._build_model(

@@ -1001,7 +1001,7 @@ def _tie_weights_nemo(model):
     # model is tied. Re-tying an untied model here would alias away the trained
     # ``lm_head.weight`` that ``from_pretrained`` just loaded (see #2941).
     config = getattr(model, "config", None)
-    if config is not None and not checkpoint_utils.get_controlling_tie_word_embeddings(config, type(model).__name__):
+    if config is not None and not checkpoint_utils.is_tied_word_embeddings(model):
         return
 
     def get_module_by_fqn(model, fqn):

@@ -22,6 +22,8 @@
 import torch.nn as nn
 from transformers.modeling_utils import _get_resolved_checkpoint_files, load_state_dict
 
+from nemo_automodel.components.models.common.tie_word_embeddings import TieSupport, get_controlling_tie_word_embeddings
+
 
 def get_rank_safe() -> int:
     """Return the current distributed rank, defaulting to 0 when not initialized."""
@@ -101,137 +103,31 @@ def resolve_trust_remote_code(pretrained_model_name_or_path):
     return not os.path.isdir(pretrained_model_name_or_path) and pretrained_model_name_or_path.startswith("nvidia/")
 
 
-def get_controlling_tie_word_embeddings(config: object, model_class_name: str) -> bool:
-    """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 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``
-            and optionally ``get_text_config``).
-        model_class_name: ``type(model).__name__`` of the owning model class.
-
-    Returns:
-        bool: The controlling ``tie_word_embeddings`` value.
-    """
-    # 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 = (
-        "Mistral3FP8VLMForConditionalGeneration",
-        "Qwen3VLMoeForConditionalGeneration",
-    )
-    if any(name in model_class_name for name in composite_top_level_models):
-        return bool(getattr(config, "tie_word_embeddings", False))
-
-    # General rule: the top-level config wins when it exposes the flag.
-    if hasattr(config, "tie_word_embeddings"):
-        return bool(config.tie_word_embeddings)
-
-    # Fallback only for configs that do not expose a top-level tie flag.
-    text_config = getattr(config, "get_text_config", lambda: None)()
-    return bool(getattr(text_config, "tie_word_embeddings", False))
-
-
 def is_tied_word_embeddings(model: nn.Module) -> bool:
-    """
-    Check if the model's word embeddings are tied.
+    """Check whether the model's word embeddings are tied.
 
-    Delegates to :func:`get_controlling_tie_word_embeddings`, which follows HF's
-    top-level-first tying semantics (replacing the previous ``text_config``-first
-    resolution).
+    A one-direction :class:`TieSupport` policy is authoritative. Only ``BOTH``
+    models and undeclared Hugging Face models need their config flag resolved.
 
     Args:
         model (nn.Module): The model to check.
 
     Returns:
         bool: True if the model's word embeddings are tied, False otherwise.
     """
+    support = getattr(model, "tie_word_embeddings_support", None)
+    # Composite VLM configs can expose an outer tie flag that does not describe
+    # the text head (for example Mistral4). No current BOTH VLM has that exception:
+    # they all honor the outer flag, so a model-owned resolver is not needed yet.
+    if support is TieSupport.TIED_ONLY:
+        return True
+    if support is TieSupport.UNTIED_ONLY:
+        return False
+
     config = getattr(model, "config", None)
     if config is None:
         return False
-    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 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."
-        )
+    return get_controlling_tie_word_embeddings(config)
 
 
 def _normalize_param_name(name: str) -> str:

@@ -752,7 +752,9 @@ def validate_hf_model_for_pipeline_support(model: torch.nn.Module) -> None:
             )
             if weights_tied:
                 issues.append(
-                    "tie_word_embeddings=True is not supported for pipelining. Use separate input/output embeddings."
+                    "Pipeline parallelism does not support tie_word_embeddings=True, and overriding "
+                    "it to tie_word_embeddings=False is not supported either. Train this model with "
+                    "another supported parallelism strategy (e.g., FSDP2) instead."
                 )
         if getattr(config, "is_encoder_decoder", False):
             issues.append("Encoder-Decoder models with cross-attention are not supported yet for pipeline parallelism.")

@@ -50,6 +50,10 @@
     load_bagel_checkpoint_state_dict,
 )
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
+from nemo_automodel.components.models.common.tie_word_embeddings import (
+    TieSupport,
+    reject_unsupported_tie_word_embeddings,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -213,6 +217,9 @@ class BagelForUnifiedMultimodal(HFCheckpointingMixin, nn.Module):
     """
 
     config_class = BagelConfig
+    # BAGEL's served checkpoints are untied; the tie flag lives on the nested
+    # text_config (aliased as llm_config), and the inner Qwen2 LM owns the head.
+    tie_word_embeddings_support: TieSupport = TieSupport.UNTIED_ONLY
 
     @dataclass(frozen=True)
     class ModelCapabilities:
@@ -225,6 +232,10 @@ class ModelCapabilities:
 
     def __init__(self, config: BagelConfig) -> None:
         super().__init__()
+        # Also covers the build_bagel_from_hf_backbones registry-bypass path, which
+        # constructs this class directly. Reads the nested text_config tie flag via
+        # the resolver's get_text_config fallback (BagelConfig has no top-level flag).
+        reject_unsupported_tie_word_embeddings(type(self), config)
         _prepare_config_for_stage(config)
         self.config = config
         self.model = BagelModel(config)

@@ -49,6 +49,10 @@
 
 from nemo_automodel.components.models.baichuan.configuration import BaichuanConfig
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
+from nemo_automodel.components.models.common.tie_word_embeddings import (
+    TieSupport,
+    reject_unsupported_tie_word_embeddings,
+)
 from nemo_automodel.components.models.common.utils import compute_lm_head_logits
 from nemo_automodel.components.models.deprecation import warn_deprecated_model_class
 
@@ -475,7 +479,9 @@ def custom_forward(*inputs):
 # Causal LM head
 # ---------------------------------------------------------------------------
 class BaichuanForCausalLM(HFCheckpointingMixin, BaichuanPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    # lm_head is a weight-normalizing NormHead, so tying it to embed_tokens is
+    # semantically wrong; all shipped Baichuan checkpoints are untied.
+    tie_word_embeddings_support: TieSupport = TieSupport.UNTIED_ONLY
 
     @dataclass(frozen=True)
     class ModelCapabilities:
@@ -488,6 +494,7 @@ class ModelCapabilities:
 
     def __init__(self, config: BaichuanConfig, **model_kwargs):
         warn_deprecated_model_class("BaichuanForCausalLM")
+        reject_unsupported_tie_word_embeddings(type(self), config)
         super().__init__(config)
         self.model = BaichuanModel(config)
         self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False)