fix: post_feedforward_layernorm

Lightning-AI · Jan 8, 2025 · 15f549d · 15f549d
1 parent 3d34921
commit 15f549d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/litgpt/scripts/convert_hf_checkpoint.py b/litgpt/scripts/convert_hf_checkpoint.py
@@ -471,6 +471,7 @@ def copy_weights_olmo2(
         "model.layers.{}.self_attn.rotary_emb.inv_freq": None,
         "model.layers.{}.post_attention_layernorm.weight": "transformer.h.{}.norm_2.weight",
         "model.layers.{}.post_attention_layernorm.bias": "transformer.h.{}.norm_2.bias",
+        "model.layers.{}.post_feedforward_layernorm.weight": "transformer.h.{}.post_mlp_norm.weight",
         "model.norm.weight": "transformer.ln_f.weight",
         "model.norm.bias": "transformer.ln_f.bias",
         "lm_head.weight": "lm_head.weight",

diff --git a/litgpt/scripts/convert_lit_checkpoint.py b/litgpt/scripts/convert_lit_checkpoint.py
@@ -355,6 +355,7 @@ def copy_weights_olmo2(
         "transformer.h.{}.attn.k_norm.weight": "model.layers.{}.self_attn.k_norm.weight",
         "transformer.h.{}.norm_2.weight": "model.layers.{}.post_attention_layernorm.weight",
         "transformer.h.{}.norm_2.bias": "model.layers.{}.post_attention_layernorm.bias",
+        "transformer.h.{}.post_mlp_norm.weight": "model.layers.{}.post_feedforward_layernorm.weight",
         "transformer.ln_f.weight": "model.norm.weight",
         "transformer.ln_f.bias": "model.norm.bias",
         "lm_head.weight": "lm_head.weight",