minor fix

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

Hide file tree

Showing 3 changed files with 2 additions and 4 deletions.
diff --git a/litgpt/config.py b/litgpt/config.py
@@ -933,6 +933,7 @@ def norm_class(self) -> Type:
         intermediate_size=11008,
         rope_base=500000,
         norm_qk=True,
+        post_mlp_norm=True,
     ),
     # https://huggingface.co/allenai/OLMo-2-1124-13B/blob/main/config.json
     dict(
@@ -954,6 +955,7 @@ def norm_class(self) -> Type:
         intermediate_size=13824,
         rope_base=500000,
         norm_qk=True,
+        post_mlp_norm=True,
     ),
 ]
 

diff --git a/litgpt/scripts/convert_hf_checkpoint.py b/litgpt/scripts/convert_hf_checkpoint.py
@@ -460,8 +460,6 @@ def copy_weights_olmo2(
 ) -> None:
     weight_map = {
         "model.embed_tokens.weight": "transformer.wte.weight",
-        "model.layers.{}.input_layernorm.weight": "transformer.h.{}.norm_1.weight",
-        "model.layers.{}.input_layernorm.bias": "transformer.h.{}.norm_1.bias",
         "model.layers.{}.self_attn.q_norm.weight": "transformer.h.{}.attn.q_norm.weight",
         "model.layers.{}.self_attn.q_proj.weight": None,
         "model.layers.{}.self_attn.k_norm.weight": "transformer.h.{}.attn.k_norm.weight",

diff --git a/litgpt/scripts/convert_lit_checkpoint.py b/litgpt/scripts/convert_lit_checkpoint.py
@@ -348,8 +348,6 @@ def copy_weights_olmo2(
 ) -> None:
     weight_map = {
         "transformer.wte.weight": "model.embed_tokens.weight",
-        "transformer.h.{}.norm_1.weight": "model.layers.{}.input_layernorm.weight",
-        "transformer.h.{}.norm_1.bias": "model.layers.{}.input_layernorm.bias",
         "transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight",
         "transformer.h.{}.attn.q_norm.weight": "model.layers.{}.self_attn.q_norm.weight",
         "transformer.h.{}.attn.k_norm.weight": "model.layers.{}.self_attn.k_norm.weight",