refine hf load funcs

InternLM · Aug 26, 2024 · 0752f65 · 0752f65
1 parent d712899
commit 0752f65
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/internlm/model/modeling_internlm.py b/internlm/model/modeling_internlm.py
@@ -486,9 +486,11 @@ def load_hf_weights(folder: str, model: nn.Module):
                 split_size,
                 dim=0,
             )[local_rank]
-            new_state_dict["norm.weight"] = state_dict["model.norm.weight"]
+            new_state_dict["norm.weight"] = state_dict.pop("model.norm.weight")
 
         missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)
+        if len(state_dict) == 0:
+            logger.warning(f"Be cautious, checkpoint state_dict keys={state_dict.keys()} have not beed loaded.")
 
         if gpc.get_local_rank(ParallelMode.DATA) == 0:
             pp_rank = 0 if not gpc.is_initialized(ParallelMode.PIPELINE) else gpc.get_local_rank(ParallelMode.PIPELINE)

diff --git a/internlm/model/modeling_internlm2.py b/internlm/model/modeling_internlm2.py
@@ -551,9 +551,11 @@ def load_hf_weights(folder: str, model: nn.Module):
                 split_size,
                 dim=0,
             )[local_rank]
-            new_state_dict["norm.weight"] = state_dict["model.norm.weight"]
+            new_state_dict["norm.weight"] = state_dict.pop("model.norm.weight")
 
         missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)
+        if len(state_dict) == 0:
+            logger.warning(f"Be cautious, checkpoint state_dict keys={state_dict.keys()} have not beed loaded.")
 
         if gpc.get_local_rank(ParallelMode.DATA) == 0:
             pp_rank = 0 if not gpc.is_initialized(ParallelMode.PIPELINE) else gpc.get_local_rank(ParallelMode.PIPELINE)

diff --git a/internlm/model/modeling_llama.py b/internlm/model/modeling_llama.py
@@ -556,9 +556,11 @@ def load_hf_weights(folder: str, model: nn.Module):
                 split_size,
                 dim=0,
             )[local_rank]
-            new_state_dict["norm.weight"] = state_dict["model.norm.weight"]
+            new_state_dict["norm.weight"] = state_dict.pop("model.norm.weight")
 
         missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)
+        if len(state_dict) == 0:
+            logger.warning(f"Be cautious, checkpoint state_dict keys={state_dict.keys()} have not beed loaded.")
 
         if gpc.get_local_rank(ParallelMode.DATA) == 0:
             pp_rank = 0 if not gpc.is_initialized(ParallelMode.PIPELINE) else gpc.get_local_rank(ParallelMode.PIPELINE)