cp: fix: tie weights outside _init_model (1817) into r0.4.0

nemo_automodel/_transformers/model_init.py

@@ -313,7 +313,7 @@ def _restore_loaded_model_dtype(
         logger.info("Restored checkpoint dtypes for %d tensors from %s", restored_count, pretrained_model_name_or_path)
 
 
-def _init_model(
+def __init_model(
     cls,
     pretrained_model_name_or_path_or_config,
     attn_implementation,
@@ -421,6 +421,57 @@ def _init_model(
     return False, model
 
 
+def _tie_weights_nemo(model):
+    if not hasattr(model, "_nemo_tied_weights_keys"):
+        return
+
+    def get_module_by_fqn(model, fqn):
+        from functools import reduce
+
+        fqn = fqn.split(".")
+        if fqn[-1] == "weight":
+            fqn = fqn[:-1]
+        return reduce(getattr, fqn, model)
+
+    for k, v in model._nemo_tied_weights_keys.items():
+        get_module_by_fqn(model, k).weight = get_module_by_fqn(model, v).weight
+
+
+def _init_model(
+    cls,
+    pretrained_model_name_or_path_or_config,
+    attn_implementation,
+    torch_dtype,
+    quantization_config,
+    force_hf,
+    *model_args,
+    **kwargs,
+):
+    is_custom_model, model = __init_model(
+        cls,
+        pretrained_model_name_or_path_or_config,
+        attn_implementation,
+        torch_dtype,
+        quantization_config,
+        force_hf,
+        *model_args,
+        **kwargs,
+    )
+    # https://github.com/NVIDIA-NeMo/Automodel/blob/a3a57176f68add7917faaa32f19228f49fcbb1ba/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad.yaml#L41
+    # this happens in nemotron_flash, where we load using force_hf, and the model is pre 5.x
+    #
+    # for safety, we tied weights after _model_init. We could do the tying in post_init, but it could be overwritten.
+    # So the sequence is roughly:
+    #   1. HF constructs NemotronFlashForCausalLM(config).
+    #   2. Inside that constructor, self.post_init() runs.
+    #   3. Only after construction returns does from_pretrained() finish loading/applying checkpoint weights.
+    #   4. That later load can assign lm_head.weight and model.embed_tokens.weight separately, which breaks any alias we create inside post_init().
+
+    if hasattr(model, "_nemo_tied_weights_keys"):
+        _tie_weights_nemo(model)
+    return is_custom_model, model
+
+
 def get_architectures(hf_config):
     """
     Get the architectures from the HF config.

nemo_automodel/_transformers/utils.py

@@ -214,14 +214,21 @@ def _find_embedding_source(model):
             for name, module in model.named_modules():
                 if isinstance(module, torch.nn.Embedding):
                     return f"{name}.weight"
-            return "model.embed_tokens.weight"
+            return None
 
         def _patched_post_init(self):
             tied = getattr(self, "_tied_weights_keys", None)
+            # if tied is list -> model is pre 5.x -> we will tie the weights after _model_init.
+            # between post_init and returned value of _model_init, there's code we don't control or can test for regressions,
+            # thus seems safer to tie weights after _model_init.
             if isinstance(tied, list):
                 source = _find_embedding_source(self)
-                self._tied_weights_keys = {k: source for k in tied}
-            return _orig_post_init(self)
+                if source is None:
+                    raise ValueError("Could not find the source of the embedding layer")
+                self._nemo_tied_weights_keys = {k: source for k in tied}
+                self._tied_weights_keys = {}
+            # call orig post init
+            _orig_post_init(self)
 
         mu.PreTrainedModel.post_init = _patched_post_init
         mu.PreTrainedModel.post_init._nemo_tied_keys_patched = True  # type: ignore[attr-defined]

tests/unit_tests/_transformers/test_transformers_utils.py

@@ -385,10 +385,11 @@ def __init__(self, config):
 
         config = FakePhi4mmCfg(tie_word_embeddings=True)
         model = FakePhi4mmModel(config)
-        tied = model._tied_weights_keys
+        tied = model._nemo_tied_weights_keys
         assert isinstance(tied, dict)
         assert "lm_head.weight" in tied
         assert tied["lm_head.weight"] == "model.embed_tokens.weight"
+        assert model._tied_weights_keys == {}
 
     def test_tied_weights_keys_patch_converts_any_model(self):
         """The post_init patch should convert _tied_weights_keys for any model, not just phi4mm."""
@@ -413,10 +414,11 @@ def __init__(self, config):
 
         config = FakeOtherCfg(tie_word_embeddings=True)
         model = FakeOtherModel(config)
-        tied = model._tied_weights_keys
+        tied = model._nemo_tied_weights_keys
         assert isinstance(tied, dict)
-        assert "lm_head.weight" in tied
+        assert "lm_head.weight" in model._nemo_tied_weights_keys
         assert tied["lm_head.weight"] == "model.embed_tokens.weight"
+        assert model._tied_weights_keys == {}
 
     def test_patches_peft_prepare_inputs(self):
         """PeftModelForCausalLM.__init__ should be patched for missing prepare_inputs_for_generation."""