Update conversion script

bigscience-workshop · Jul 5, 2022 · 2c9bea5 · 2c9bea5
1 parent cc465ca
commit 2c9bea5
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/evaluation/results/opt/convert_transformers_checkpoint_to_meg_ds.py b/evaluation/results/opt/convert_transformers_checkpoint_to_meg_ds.py
@@ -62,14 +62,18 @@ def compute_meg_ds_weight_names(num_layers: int):
         f"layer_{str(num_layers + 4).zfill(2)}-model_00-model_states.pt": [
             "weight",
             "bias"
+        ],
+        f"layer_{str(num_layers + 5).zfill(2)}-model_00-model_states.pt": [
+            "lm_head.weight",
         ]
     }
 
 NON_TRANSFORMERS_BLOCK_WEIGHTS = {
     "word_embeddings.weight": "decoder.embed_tokens.weight",
     "position_embeddings.weight": "decoder.embed_positions.weight",
     "weight": "decoder.final_layer_norm.weight",
-    "bias": "decoder.final_layer_norm.bias"
+    "bias": "decoder.final_layer_norm.bias",
+    "lm_head.weight": "lm_head.weight"
 }
 TRANSFORMERS_BLOCK_WEIGHTS = {
     "input_layernorm.weight": ["self_attn_layer_norm.weight"],