From 1e0fcca94875bc46157e0747b81b9ed3ccaef732 Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> Date: Tue, 14 Apr 2026 01:01:55 -0700 Subject: [PATCH] fix: tie weights outside _init_model (#1817) * tie weights outside _init_model Signed-off-by: Alexandros Koumparoulis * update test Signed-off-by: Alexandros Koumparoulis * undo test Signed-off-by: Alexandros Koumparoulis * add typecasting back Signed-off-by: Alexandros Koumparoulis --------- Signed-off-by: Alexandros Koumparoulis Signed-off-by: NeMo Bot --- nemo_automodel/_transformers/model_init.py | 53 ++++++++++++++++++- nemo_automodel/_transformers/utils.py | 13 +++-- .../_transformers/test_transformers_utils.py | 8 +-- 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/nemo_automodel/_transformers/model_init.py b/nemo_automodel/_transformers/model_init.py index e60c7ba9ac..71948335e1 100644 --- a/nemo_automodel/_transformers/model_init.py +++ b/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. diff --git a/nemo_automodel/_transformers/utils.py b/nemo_automodel/_transformers/utils.py index 0f122976e2..f0525d4ebc 100644 --- a/nemo_automodel/_transformers/utils.py +++ b/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] diff --git a/tests/unit_tests/_transformers/test_transformers_utils.py b/tests/unit_tests/_transformers/test_transformers_utils.py index 987d10f72e..6b3c6ee4f6 100644 --- a/tests/unit_tests/_transformers/test_transformers_utils.py +++ b/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."""