change the way of load checkpoint

Author: Cui-yshoho

mindone/diffusers/models/model_loading_utils.py

@@ -97,7 +97,7 @@ def load_state_dict(
             if disable_mmap:
                 return safe_load(open(checkpoint_file, "rb").read())
             else:
-                return safe_load_file(checkpoint_file)
+                return ms.load_checkpoint(checkpoint_file, format="safetensors")
         else:
             raise NotImplementedError(
                 f"Only supports deserialization of weights file in safetensors format, but got {checkpoint_file}"
@@ -140,11 +140,11 @@ def _load_state_dict_into_model(
                     and any(module_to_keep_in_fp32 in k.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules)
                     and dtype == ms.float16
                 ):
-                    v.set_dtype(ms.float32)
+                    state_dict[k] = ms.Parameter(v.to(ms.float32), name=k)
                 else:
-                    v.set_dtype(local_state[k].dtype)
+                    state_dict[k] = ms.Parameter(v.to(local_state[k].dtype), name=k)
             else:
-                v.set_dtype(local_state[k].dtype)
+                state_dict[k] = ms.Parameter(v.to(local_state[k].dtype), name=k)
         else:
             pass  # unexpect key keeps origin dtype
     cm = silence_mindspore_logger() if is_sharded else nullcontext()

mindone/transformers/modeling_utils.py

@@ -60,6 +60,7 @@
 import mindspore as ms
 from mindspore import Parameter, Tensor, mint, nn, ops
 from mindspore.nn import CrossEntropyLoss, Identity
+from mindspore.nn.utils import no_init_parameters
 
 from .activations import get_activation
 from .generation.utils import GenerationMixin
@@ -349,7 +350,7 @@ def _load_state_dict_into_model(model_to_load, state_dict, start_prefix, is_shar
     local_state = {v.name: v for k, v in model_to_load.parameters_and_names()}
     for k, v in state_dict.items():
         if k in local_state:
-            v.set_dtype(local_state[k].dtype)
+            state_dict[k] = ms.Parameter(v.to(local_state[k].dtype), name=k)
         else:
             pass  # unexpect key keeps origin dtype
     cm = silence_mindspore_logger() if is_sharded else nullcontext()
@@ -977,8 +978,8 @@ def _from_config(cls, config, **kwargs):
                 use_flash_attention_2=use_flash_attention_2,
                 mindspore_dtype=mindspore_dtype,
             )
-
-        model = cls(config, **kwargs)
+        with no_init_parameters():
+            model = cls(config, **kwargs)
 
         # We cannot set default mindspore dtype. So we need to cast model weights after creating.
         if mindspore_dtype is not None:
@@ -2348,7 +2349,8 @@ def from_pretrained(
             config, use_flash_attention_2=use_flash_attention_2, mindspore_dtype=mindspore_dtype
         )
 
-        model = cls(config, *model_args, **model_kwargs)
+        with no_init_parameters():
+            model = cls(config, *model_args, **model_kwargs)
 
         # Make sure to tie the weights correctly
         model.tie_weights()