Upload MyLLaMa

configure_for_hf.py

@@ -41,6 +41,20 @@ class MyLLaMa(PreTrainedModel):
             n_chckpnt_segments=config.n_chckpnt_segments,
         )
 
+    def load_state_dict(self, state_dict, **kwargs):
+        for key in list(state_dict.keys()):
+            if "rmsnorm1.weight" in key:
+                new_key = key.replace("rmsnorm1.weight", "rmsnorm1.gamma")
+                state_dict[new_key] = state_dict.pop(key)
+            elif "rmsnorm2.weight" in key:
+                new_key = key.replace("rmsnorm2.weight", "rmsnorm2.gamma")
+                state_dict[new_key] = state_dict.pop(key)
+            elif "rmsnorm.weight" in key:
+                new_key = key.replace("rmsnorm.weight", "rmsnorm.gamma")
+                state_dict[new_key] = state_dict.pop(key)
+
+        super().load_state_dict(state_dict, **kwargs)
+
     def forward(self, tensor, labels=None):
         logits = self.model(tensor)["logits"]
         if labels is not None: