From 55b39570884b0fc0fa93adf2ecf9ec1f1fb8aaaf Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Fri, 3 Jul 2026 14:09:13 +0400 Subject: [PATCH 1/2] feat(models): reject unsupported tie_word_embeddings flips for remaining families Closes the tie-audit follow-up. Reject tie=True on the untied-default families (mistral3_vlm, step3p5, step3p7, nemotron_omni), verified against their published configs; reject tie=False on the tied-default gemma4_moe and flip the #2601 untied tests to expect the raise. Adds the symmetric reject_unsupported_untied_word_embeddings helper. Refs #2512 Signed-off-by: Achyuthan Sivasankar --- nemo_automodel/components/checkpoint/utils.py | 27 +++++++++++++++++++ .../components/models/gemma4_moe/model.py | 4 +++ .../components/models/mistral3_vlm/model.py | 4 +++ .../components/models/nemotron_omni/model.py | 2 ++ .../components/models/step3p5/model.py | 2 ++ .../components/models/step3p7/model.py | 2 ++ .../test_gemma4_moe_tied_weights.py | 25 +++++++---------- .../unit_tests/utils/test_checkpoint_utils.py | 17 +++++++++--- 8 files changed, 64 insertions(+), 19 deletions(-) diff --git a/nemo_automodel/components/checkpoint/utils.py b/nemo_automodel/components/checkpoint/utils.py index da2cef5cc6..45f3b7b8ba 100644 --- a/nemo_automodel/components/checkpoint/utils.py +++ b/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 is unset. + """ + 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.", "") diff --git a/nemo_automodel/components/models/gemma4_moe/model.py b/nemo_automodel/components/models/gemma4_moe/model.py index bbf439f18c..7c057af47a 100644 --- a/nemo_automodel/components/models/gemma4_moe/model.py +++ b/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 diff --git a/nemo_automodel/components/models/mistral3_vlm/model.py b/nemo_automodel/components/models/mistral3_vlm/model.py index 06253e82ed..142675058b 100644 --- a/nemo_automodel/components/models/mistral3_vlm/model.py +++ b/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 diff --git a/nemo_automodel/components/models/nemotron_omni/model.py b/nemo_automodel/components/models/nemotron_omni/model.py index 3b79756b13..3aeac5ee30 100644 --- a/nemo_automodel/components/models/nemotron_omni/model.py +++ b/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() # --------------------------------------------------------------- diff --git a/nemo_automodel/components/models/step3p5/model.py b/nemo_automodel/components/models/step3p5/model.py index b8a23fef74..c5b9146890 100644 --- a/nemo_automodel/components/models/step3p5/model.py +++ b/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) diff --git a/nemo_automodel/components/models/step3p7/model.py b/nemo_automodel/components/models/step3p7/model.py index 879005e713..fc6a37e3d1 100644 --- a/nemo_automodel/components/models/step3p7/model.py +++ b/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) diff --git a/tests/unit_tests/models/gemma4_moe/test_gemma4_moe_tied_weights.py b/tests/unit_tests/models/gemma4_moe/test_gemma4_moe_tied_weights.py index 8a97ca575d..394be43358 100644 --- a/tests/unit_tests/models/gemma4_moe/test_gemma4_moe_tied_weights.py +++ b/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) diff --git a/tests/unit_tests/utils/test_checkpoint_utils.py b/tests/unit_tests/utils/test_checkpoint_utils.py index decd7aa819..a28d7ec2bb 100644 --- a/tests/unit_tests/utils/test_checkpoint_utils.py +++ b/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. From b37ec0f5c9e189bea288fab81b7c168e494e6203 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Fri, 3 Jul 2026 19:35:44 +0400 Subject: [PATCH 2/2] docs(checkpoint): clarify untie-guard raises when the flag is False (review) Signed-off-by: Achyuthan Sivasankar --- nemo_automodel/components/checkpoint/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_automodel/components/checkpoint/utils.py b/nemo_automodel/components/checkpoint/utils.py index 45f3b7b8ba..221b1c53b2 100644 --- a/nemo_automodel/components/checkpoint/utils.py +++ b/nemo_automodel/components/checkpoint/utils.py @@ -224,7 +224,7 @@ def reject_unsupported_untied_word_embeddings(config: object, model_class_name: model_class_name: ``type(self).__name__`` of the constructing model. Raises: - NotImplementedError: if the controlling ``tie_word_embeddings`` flag is unset. + NotImplementedError: if the controlling ``tie_word_embeddings`` flag evaluates to ``False``. """ if not get_controlling_tie_word_embeddings(config, model_class_name): raise NotImplementedError(