use HF standard naming for qkv and mlp layers

Author: zigzagcai

internlm/model/model_implementations/transformers/modeling_internlm2.py

@@ -809,19 +809,34 @@ def unique_kv_index(i):
                         )
                         current_states[name] = torch.index_select(current_states[name], 0, unique_index)
 
+        mlp_layer_fusion = any(".w13." in key for key in current_states.keys())
         fixed_current_states = {}
-        for name in current_states.keys():
-            if "wqkv" in name:
-                new_name = name.replace("wqkv", "qkv_proj")
-            elif ".w1." in name:
-                new_name = name.replace("feed_forward.w1", "feed_forward.gate_proj")
-            elif ".w2." in name:
-                new_name = name.replace("feed_forward.w2", "feed_forward.down_proj")
-            elif ".w3." in name:
-                new_name = name.replace("feed_forward.w3", "feed_forward.up_proj")
+        for key in current_states.keys():
+            if "wqkv" in key:
+                new_key = key.replace("wqkv", "qkv_proj")
+            # elif "wq" in key:
+            #     new_key = key.replace("wq", "q_proj")
+            # elif "wk" in key:
+            #     new_key = key.replace("wk", "k_proj")
+            # elif "wv" in key:
+            #     new_key = key.replace("wv", "v_proj")
+            # elif "wo" in key:
+            #     new_key = key.replace("wo", "o_proj")
+            elif ".w1." in key:
+                new_key = key.replace("feed_forward.w1", "feed_forward.gate_proj")
+            elif ".w2." in key:
+                new_key = (
+                    key.replace("feed_forward.w2", "feed_forward.dense_4h_to_h")
+                    if mlp_layer_fusion
+                    else key.replace("feed_forward.w2", "feed_forward.down_proj")
+                )
+            elif ".w3." in key:
+                new_key = key.replace("feed_forward.w3", "feed_forward.up_proj")
+            elif ".w13." in key:
+                new_key = key.replace("feed_forward.w13", "feed_forward.dense_h_to_4h")
             else:
-                new_name = name
-            fixed_current_states[new_name] = current_states[name]
+                new_key = key
+            fixed_current_states[new_key] = current_states[key]
 
         missing_keys, unexpected_keys = model.load_state_dict(fixed_current_states, strict=False)
 

internlm/model/model_ops/modules/mha.py

@@ -46,27 +46,27 @@ def split_fused_wqkv_weight(wqkv, *args, **kwargs):  # pylint: disable=W0613
 
 
 def _qkv_pre_load_convert(module: "GQA", state_dict, prefix: str, *args, **kwargs) -> None:  # pylint: disable=W0613
-    wq_name, wk_name, wv_name, fused_name = (
+    wq_name, wk_name, wv_name, wqkv_name = (
         f"{prefix}q_proj.weight",
         f"{prefix}k_proj.weight",
         f"{prefix}v_proj.weight",
         f"{prefix}qkv_proj.weight",
     )
 
-    if module.enable_qkv_fusion and fused_name not in state_dict:
+    if module.enable_qkv_fusion and wqkv_name not in state_dict:
         wq, wk, wv = state_dict.pop(wq_name), state_dict.pop(wk_name), state_dict.pop(wv_name)
-        state_dict[fused_name] = torch.cat([wq, wk, wv], dim=0)
+        state_dict[wqkv_name] = torch.cat([wq, wk, wv], dim=0)
 
     if not module.enable_qkv_fusion and (
         wq_name not in state_dict or wk_name not in state_dict or wv_name not in state_dict
     ):
         state_dict[wq_name], state_dict[wk_name], state_dict[wv_name] = split_fused_wqkv_weight(
-            state_dict.pop(fused_name), *args, **kwargs
+            state_dict.pop(wqkv_name), *args, **kwargs
         )
 
 
 def _qkv_save_convert(module: "GQA", state_dict, prefix: str, *args, **kwargs) -> Dict:  # pylint: disable=W0613
-    wq_name, wk_name, wv_name, fused_name = (
+    wq_name, wk_name, wv_name, wqkv_name = (
         f"{prefix}q_proj.weight",
         f"{prefix}k_proj.weight",
         f"{prefix}v_proj.weight",
@@ -75,7 +75,7 @@ def _qkv_save_convert(module: "GQA", state_dict, prefix: str, *args, **kwargs) -
 
     if module.enable_qkv_fusion:
         state_dict[wq_name], state_dict[wk_name], state_dict[wv_name] = split_fused_wqkv_weight(
-            state_dict.pop(fused_name), *args, **kwargs
+            state_dict.pop(wqkv_name), *args, **kwargs
         )
 
     return state_dict
@@ -162,6 +162,10 @@ def __init__(
             self.q_proj = new_linear("wq", embed_dim, embed_dim, bias, **factory_kwargs)
             self.k_proj = new_linear("wk", embed_dim, self.kv_dim, bias, **factory_kwargs)
             self.v_proj = new_linear("wv", embed_dim, self.kv_dim, bias, **factory_kwargs)
+        self.register_checkpoint_compatibility_hooks(
+            partial(_qkv_pre_load_convert, q_dim=self.embed_dim, kv_dim=self.kv_dim),
+            partial(_qkv_save_convert, q_dim=self.embed_dim, kv_dim=self.kv_dim),
+        )
 
         self.inner_attn = SelfAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout)
         self.inner_cross_attn = CrossAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout)
@@ -462,14 +466,14 @@ def __init__(
         if enable_qkv_fusion:
             assert bias is False, "Fuesd wqkv only support bias is False."
             self.qkv_proj = new_linear("wqkv", embed_dim, q_dim + 2 * self.kv_dim, bias, **factory_kwargs)
-            self._register_load_state_dict_pre_hook(
-                partial(_qkv_pre_load_convert, q_dim=q_dim, kv_dim=self.kv_dim), with_module=True
-            )
-            self._register_state_dict_hook(partial(_qkv_save_convert, q_dim=q_dim, kv_dim=self.kv_dim))
         else:
             self.q_proj = new_linear("wq", embed_dim, q_dim, bias, **factory_kwargs)
             self.k_proj = new_linear("wk", embed_dim, self.kv_dim, bias, **factory_kwargs)
             self.v_proj = new_linear("wv", embed_dim, self.kv_dim, bias, **factory_kwargs)
+        self.register_checkpoint_compatibility_hooks(
+            partial(_qkv_pre_load_convert, q_dim=q_dim, kv_dim=self.kv_dim),
+            partial(_qkv_save_convert, q_dim=self.embed_dim, kv_dim=self.kv_dim),
+        )
 
         self.inner_attn = SelfAttention(
             causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout, layer_idx=layer_idx

internlm/model/model_ops/modules/mlp.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional
 
 import torch
 from torch import nn
@@ -22,21 +22,38 @@ def split_fused_mlp_weight(w1_w3):
 def _mlp_pre_load_convert(
     module: "FeedForward", state_dict, prefix: str, *args, **kwargs  # pylint: disable=W0613
 ) -> None:
-    w1_name, w3_name, fused_name = f"{prefix}w1.weight", f"{prefix}w3.weight", f"{prefix}fused_w1_w3.weight"
+    gate_proj_name, up_proj_name, dense_h_to_4h_name = (
+        f"{prefix}gate_proj.weight",
+        f"{prefix}up_proj.weight",
+        f"{prefix}dense_h_to_4h.weight",
+    )
+    down_proj_name, dense_4h_to_h_name = f"{prefix}down_proj.weight", f"{prefix}dense_4h_to_h.weight"
 
-    if module.mlp_layer_fusion and fused_name not in state_dict:
-        w1, w3 = state_dict.pop(w1_name), state_dict.pop(w3_name)
-        state_dict[fused_name] = torch.cat([w1, w3], dim=0)
+    if module.mlp_layer_fusion and dense_h_to_4h_name not in state_dict:
+        gate_proj, up_proj = state_dict.pop(gate_proj_name), state_dict.pop(up_proj_name)
+        state_dict[dense_h_to_4h_name] = torch.cat([gate_proj, up_proj], dim=0)
+        state_dict[dense_4h_to_h_name] = state_dict.pop(down_proj_name)
 
-    if not module.mlp_layer_fusion and (w1_name not in state_dict or w3_name not in state_dict):
-        state_dict[w1_name], state_dict[w3_name] = split_fused_mlp_weight(state_dict.pop(fused_name))
+    if not module.mlp_layer_fusion and (gate_proj_name not in state_dict or up_proj_name not in state_dict):
+        state_dict[gate_proj_name], state_dict[up_proj_name] = split_fused_mlp_weight(
+            state_dict.pop(dense_h_to_4h_name)
+        )
+        state_dict[down_proj_name] = state_dict.pop(dense_4h_to_h_name)
 
 
 def _mlp_save_convert(module: "FeedForward", state_dict, prefix: str, *args, **kwargs) -> Dict:  # pylint: disable=W0613
-    w1_name, w3_name, fused_name = f"{prefix}w1.weight", f"{prefix}w3.weight", f"{prefix}fused_w1_w3.weight"
+    gate_proj_name, up_proj_name, dense_h_to_4h_name = (
+        f"{prefix}gate_proj.weight",
+        f"{prefix}up_proj.weight",
+        f"{prefix}dense_h_to_4h.weight",
+    )
+    down_proj_name, dense_4h_to_h_name = f"{prefix}down_proj.weight", f"{prefix}dense_4h_to_h.weight"
 
     if module.mlp_layer_fusion:
-        state_dict[w1_name], state_dict[w3_name] = split_fused_mlp_weight(state_dict.pop(fused_name))
+        state_dict[gate_proj_name], state_dict[up_proj_name] = split_fused_mlp_weight(
+            state_dict.pop(dense_h_to_4h_name)
+        )
+        state_dict[down_proj_name] = state_dict.pop(dense_4h_to_h_name)
 
     return state_dict
 
@@ -92,9 +109,6 @@ def __init__(
             self.dense_4h_to_h = new_linear(
                 "w2", hidden_features, out_features, bias, device=device, dtype=dtype, is_expert=is_expert
             )
-
-            self._register_load_state_dict_pre_hook(_mlp_pre_load_convert, with_module=True)
-            self._register_state_dict_hook(_mlp_save_convert)
         else:
             self.gate_proj = new_linear(
                 "w1", in_features, hidden_features, bias, device=device, dtype=dtype, is_expert=is_expert
@@ -105,6 +119,16 @@ def __init__(
             self.down_proj = new_linear(
                 "w2", hidden_features, out_features, bias, device=device, dtype=dtype, is_expert=is_expert
             )
+        self.register_checkpoint_compatibility_hooks(_mlp_pre_load_convert, _mlp_save_convert)
+
+    def register_checkpoint_compatibility_hooks(
+        self, pre_load_hook: Optional[Callable] = None, pre_save_hook: Optional[Callable] = None
+    ):
+        # Here we explicitly expose the checkpoint compatibility interface of the module,
+        # hoping that model developers will make good use of it when adapting.
+        # Is this interface already meeting all reasonable requirements?
+        self._register_load_state_dict_pre_hook(pre_load_hook, with_module=True)
+        self._register_state_dict_hook(pre_save_hook)
 
     def forward(self, x):
         if not self.mlp_layer_fusion:

tests/test_training/train_CI.py

@@ -22,26 +22,26 @@
 from internlm.checkpoint import CheckpointManager  # noqa: E402
 from internlm.core.context import ParallelMode  # noqa: E402
 from internlm.core.context import global_context as gpc  # noqa: E402
+from internlm.core.trainer import record_current_batch_training_metrics  # noqa: E402
 from internlm.core.trainer import Trainer, TrainState  # noqa: E402
 from internlm.data import (  # noqa: E402
     build_train_loader_with_data_type,
     build_valid_loader_with_data_type,
 )
 from internlm.eval import evaluate_on_val_dls  # noqa: E402
 from internlm.initialize import initialize_launcher  # noqa: E402
+from internlm.initialize import initialize_trainer  # noqa: E402
 from internlm.initialize.initialize_model import (  # noqa: E402
     initialize_model_and_parallel_communicator,
 )
-from internlm.initialize import initialize_trainer  # noqa: E402
 from internlm.model.model_ops.losses import InternLoss  # noqa: E402
 from internlm.model.model_ops.metrics import (  # noqa: E402
     AccPerplex,
     SchedulerMetricHook,
 )
 from internlm.monitor import initialize_monitor_manager  # noqa: E402
-from internlm.monitor import monitor_manager as mm  # noqa: E402
 from internlm.monitor import send_alert_message  # noqa: E402
-from internlm.core.trainer import record_current_batch_training_metrics  # noqa: E402
+from internlm.monitor import monitor_manager as mm  # noqa: E402
 from internlm.utils.common import (  # noqa: E402
     BatchSkipper,
     get_current_device,
@@ -65,18 +65,38 @@ def check_model_weights(model, ckpt_path, total_equal=False):
     model1_dict = torch.load(ckpt_path, map_location="cuda")
     model2_dict = model.state_dict()
 
-    copy_of_ordered_dict = model2_dict.copy()
-
-    for key in copy_of_ordered_dict.keys():
+    mlp_layer_fusion = any(".w13." in key for key in model1_dict.keys())
+    fixed_model1_dict = {}
+    for key in model1_dict.keys():
         if "wqkv" in key:
-            model2_dict[key.replace("wqkv", "Wqkv")] = model2_dict.pop(key)
-            key = key.replace("wqkv", "Wqkv")
-        if key not in model1_dict:
-            assert False, f"Error: The key {key} for current model dose not exist in standard ckpt!"
+            new_key = key.replace("wqkv", "qkv_proj")
+        elif "wq" in key:
+            new_key = key.replace("wq", "q_proj")
+        elif "wk" in key:
+            new_key = key.replace("wk", "k_proj")
+        elif "wv" in key:
+            new_key = key.replace("wv", "v_proj")
+        # elif "wo" in key:
+        #     new_key = key.replace("wo", "o_proj")
+        elif ".w1." in key:
+            new_key = key.replace("feed_forward.w1", "feed_forward.gate_proj")
+        elif ".w2." in key:
+            new_key = (
+                key.replace("feed_forward.w2", "feed_forward.dense_4h_to_h")
+                if mlp_layer_fusion
+                else key.replace("feed_forward.w2", "feed_forward.down_proj")
+            )
+        elif ".w3." in key:
+            new_key = key.replace("feed_forward.w3", "feed_forward.up_proj")
+        elif ".w13." in key:
+            new_key = key.replace("feed_forward.w13", "feed_forward.dense_h_to_4h")
+        else:
+            new_key = key
+        fixed_model1_dict[new_key] = model1_dict[key]
 
-    for key in model1_dict.keys():
+    for key in fixed_model1_dict.keys():
         if key in model2_dict:
-            tensor1 = model1_dict[key]
+            tensor1 = fixed_model1_dict[key]
             tensor2 = model2_dict[key]
             if total_equal:
                 assert torch.equal(tensor1, tensor2), "model weights are not equal"