fix(transformers): gate _tie_weights_nemo on tie_word_embeddings flag

nemo_automodel/_transformers/model_init.py

@@ -995,6 +995,15 @@ def _tie_weights_nemo(model):
     if not hasattr(model, "_nemo_tied_weights_keys"):
         return
 
+    # Only re-tie when the controlling config flag actually requests tied
+    # embeddings. ``_nemo_tied_weights_keys`` names the *candidate* tied keys
+    # (pre-v5 list-form ``_tied_weights_keys`` semantics); it does not mean the
+    # 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__):
+        return
+
     def get_module_by_fqn(model, fqn):
         from functools import reduce
 

tests/unit_tests/_transformers/test_model_init.py

@@ -703,3 +703,85 @@ def test_returns_none_when_trust_remote_code_not_set(self):
         cfg.auto_map = {"AutoModelForCausalLM": "modeling.MyModel"}
         result = _try_get_remote_code_model_cls(cfg, "/some/path", "AutoModelForCausalLM", {})
         assert result is None
+
+
+class TestTieWeightsNemoConfigGate:
+    """_tie_weights_nemo must honor the controlling tie_word_embeddings flag (#2941).
+
+    ``_nemo_tied_weights_keys`` names the candidate tied keys (pre-v5 list-form
+    ``_tied_weights_keys`` semantics); it does not mean the model is tied.
+    Re-tying an untied model aliases away the trained ``lm_head.weight`` that
+    ``from_pretrained`` just loaded.
+    """
+
+    @staticmethod
+    def _make_model(tie: bool | None) -> nn.Module:
+        from transformers import PretrainedConfig
+
+        class _TinyModel(nn.Module):
+            _nemo_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+            def __init__(self):
+                super().__init__()
+                self.model = nn.Module()
+                self.model.embed_tokens = nn.Embedding(10, 4)
+                self.lm_head = nn.Linear(4, 10, bias=False)
+
+        model = _TinyModel()
+        if tie is not None:
+            model.config = PretrainedConfig(tie_word_embeddings=tie)
+        return model
+
+    def test_untied_config_keeps_separate_lm_head(self):
+        """tie_word_embeddings=False: the loaded lm_head must not be aliased away."""
+        from nemo_automodel._transformers.model_init import _tie_weights_nemo
+
+        model = self._make_model(tie=False)
+        lm_head_before = model.lm_head.weight.detach().clone()
+
+        _tie_weights_nemo(model)
+
+        assert model.lm_head.weight is not model.model.embed_tokens.weight
+        assert model.lm_head.weight.data_ptr() != model.model.embed_tokens.weight.data_ptr()
+        torch.testing.assert_close(model.lm_head.weight, lm_head_before)
+
+    def test_tied_config_reties(self):
+        """tie_word_embeddings=True: keep the #1817 re-tie behavior."""
+        from nemo_automodel._transformers.model_init import _tie_weights_nemo
+
+        model = self._make_model(tie=True)
+        _tie_weights_nemo(model)
+
+        assert model.lm_head.weight is model.model.embed_tokens.weight
+
+    def test_missing_config_still_reties(self):
+        """No config attribute: fall back to the conservative #1817 re-tie."""
+        from nemo_automodel._transformers.model_init import _tie_weights_nemo
+
+        model = self._make_model(tie=None)
+        _tie_weights_nemo(model)
+
+        assert model.lm_head.weight is model.model.embed_tokens.weight
+
+    def test_untied_state_dict_roundtrip_is_lossless(self):
+        """Resume scenario: distinct lm_head/embed weights must survive construct-then-load.
+
+        Before the fix, construction aliased both params to one storage, so
+        ``load_state_dict`` wrote both checkpoint tensors into the same memory
+        (last writer wins) — corrupting embeddings and/or head on resume.
+        """
+        from nemo_automodel._transformers.model_init import _tie_weights_nemo
+
+        source = self._make_model(tie=False)
+        with torch.no_grad():
+            source.model.embed_tokens.weight.uniform_(-1.0, 1.0)
+            source.lm_head.weight.uniform_(-1.0, 1.0)
+        checkpoint = {k: v.detach().clone() for k, v in source.state_dict().items()}
+
+        resumed = self._make_model(tie=False)
+        _tie_weights_nemo(resumed)  # runs at the end of _init_model, before checkpoint load
+        resumed.load_state_dict(checkpoint)
+
+        torch.testing.assert_close(resumed.lm_head.weight, checkpoint["lm_head.weight"])
+        torch.testing.assert_close(resumed.model.embed_tokens.weight, checkpoint["model.embed_tokens.weight"])
+        assert resumed.lm_head.weight.data_ptr() != resumed.model.embed_tokens.weight.data_ptr()