diff --git a/nemo_automodel/_transformers/model_init.py b/nemo_automodel/_transformers/model_init.py index 8dfe6532b8..2575d91adf 100644 --- a/nemo_automodel/_transformers/model_init.py +++ b/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 diff --git a/tests/unit_tests/_transformers/test_model_init.py b/tests/unit_tests/_transformers/test_model_init.py index 54def8abdc..553d6585d9 100644 --- a/tests/unit_tests/_transformers/test_model_init.py +++ b/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()