[Mamba] from pretrained issue with self.embeddings (#29851)

* nit * update * oups * Update src/transformers/models/mamba/modeling_mamba.py Co-authored-by: Lysandre Debut <[email protected]> --------- Co-authored-by: Lysandre Debut <[email protected]>
huggingface · May 14, 2024 · d6453e7 · d6453e7
1 parent c37849c
commit d6453e7
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
@@ -501,8 +501,15 @@ def __init__(self, config):
         self.gradient_checkpointing = False
         self.norm_f = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
         # Initialize weights and apply final processing
+        self._register_load_state_dict_pre_hook(self.load_hook)
         self.post_init()
 
+    def load_hook(self, state_dict, prefix, *args):
+        for k in state_dict:
+            if "embedding." in k:
+                state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k)
+                break
+
     def get_input_embeddings(self):
         return self.embeddings