Fix register_buffer parameter (#1036)

* Fix register_buffer parameter * Fix register_buffer parameter
EleutherAI · Sep 20, 2023 · 0ce77ab · 0ce77ab
1 parent 8903a96
commit 0ce77ab
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/tools/convert_module_to_hf.py b/tools/convert_module_to_hf.py
@@ -227,10 +227,12 @@ def convert(input_checkpoint_path, loaded_config, output_checkpoint_path):
         state_dict["attention.rotary_emb.inv_freq"] = loaded_tp_ranks[0][
             "attention.rotary_emb.inv_freq"
         ]
-        state_dict["attention.bias"] = hf_layer.state_dict()["attention.bias"]
-        state_dict["attention.masked_bias"] = hf_layer.state_dict()[
-            "attention.masked_bias"
-        ]
+        if "attention.bias" in hf_layer.state_dict():
+            state_dict["attention.bias"] = hf_layer.state_dict()["attention.bias"]
+        if "attention.masked_bias" in hf_layer.state_dict():
+            state_dict["attention.masked_bias"] = hf_layer.state_dict()[
+                "attention.masked_bias"
+            ]
 
         # load state_dict into layer
         hf_layer.load_state_dict(state_dict)

diff --git a/tools/convert_sequential_to_hf.py b/tools/convert_sequential_to_hf.py
@@ -247,10 +247,12 @@ def convert(input_checkpoint_path, loaded_config, output_checkpoint_path):
             loaded_tp_ranks, "attention.rotary_emb.inv_freq", layer_i + 2
         )[0]
 
-        state_dict["attention.bias"] = hf_layer.state_dict()["attention.bias"]
-        state_dict["attention.masked_bias"] = hf_layer.state_dict()[
-            "attention.masked_bias"
-        ]
+        if "attention.bias" in hf_layer.state_dict():
+            state_dict["attention.bias"] = hf_layer.state_dict()["attention.bias"]
+        if "attention.masked_bias" in hf_layer.state_dict():
+            state_dict["attention.masked_bias"] = hf_layer.state_dict()[
+                "attention.masked_bias"
+            ]
 
         # load state_dict into layer
         hf_layer.load_state_dict(state_dict)