remove model.model in ckpt loading

Author: Gasoonjia

install/install_requirements.sh

@@ -53,7 +53,7 @@ PYTORCH_NIGHTLY_VERSION=dev20240814
 VISION_NIGHTLY_VERSION=dev20240814
 
 # Nightly version for torchtune
-TUNE_NIGHTLY_VERSION=dev20240916
+TUNE_NIGHTLY_VERSION=dev20240918
 
 
 # Uninstall triton, as nightly will depend on pytorch-triton, which is one and the same

torchchat/cli/builder.py

@@ -388,11 +388,9 @@ def _load_model_default(builder_args, only_config=False):
             builder_args.device
         ):
             model = Model.from_params(builder_args.params_path)
-        state_dict = flamingo_meta_to_tune(checkpoint)
-        model.model.load_state_dict(state_dict)
-    else:
-        checkpoint = {"model." + k: v for k, v in checkpoint.items()}
-        model.load_state_dict(checkpoint, assign=True, strict=True)
+        checkpoint = flamingo_meta_to_tune(checkpoint)
+
+    model.load_state_dict(checkpoint, assign=True, strict=True)
 
     return model
 

torchchat/model.py

@@ -442,6 +442,8 @@ def __init__(self, config: ModelArgs) -> None:
         # It should be assigned in the actual model implementation, if any.
         self.text_transformer_args = None
 
+        self._register_load_state_dict_pre_hook(self._load_model_state_dict)
+
     def build_model(self) -> nn.Module:
         """
         Builds a model based on the provided configuration.
@@ -468,6 +470,13 @@ def _replace_known_params(self, params):
                 params[key] = patterns[value]
         return params
 
+    def _load_model_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        # 修改 state dict 中的键值
+        for key in list(state_dict.keys()):
+            new_key = 'model.' + key
+            state_dict[new_key] = state_dict.pop(key)
+        return state_dict
+    
     @abstractmethod
     def forward(self, *args, **kwargs):
         raise NotImplementedError("forward method is not implemented")