a new method

Author: Cui-yshoho

mindone/diffusers/models/model_loading_utils.py

@@ -35,7 +35,6 @@
 
 import mindspore as ms
 from mindspore import nn, ops
-from mindspore.ops import Cast
 
 from ...safetensors.mindspore import load as safe_load
 from ..utils import (
@@ -51,7 +50,8 @@
 )
 
 logger = logging.get_logger(__name__)
-cpu_cast = Cast().set_device("CPU")
+ms.Parameter._data = ms.Tensor.data
+ms.Parameter.data_ptr = ms.Tensor.data_ptr
 
 _CLASS_REMAPPING_DICT = {
     "Transformer2DModel": {
@@ -146,11 +146,11 @@ def _load_state_dict_into_model(
                 if keep_in_fp32_modules is not None and any(
                     module_to_keep_in_fp32 in k.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules
                 ):
-                    state_dict[k] = ms.Parameter(cpu_cast(v.data, ms.float32), name=k)
+                    v._data = v.to(device="CPU", dtype=ms.float32)
                 else:
-                    state_dict[k] = ms.Parameter(cpu_cast(v.data, local_state[k].dtype), name=k)
+                    v._data = v.to(device="CPU", dtype=local_state[k].dtype)
             else:
-                state_dict[k] = ms.Parameter(cpu_cast(v.data, local_state[k].dtype), name=k)
+                v._data = v.to(device="CPU", dtype=local_state[k].dtype)
         else:
             pass  # unexpect key keeps origin dtype
     cm = silence_mindspore_logger() if is_sharded else nullcontext()

mindone/diffusers/models/modeling_utils.py

@@ -37,6 +37,7 @@
 from mindspore.nn.utils import no_init_parameters
 
 from mindone.safetensors.mindspore import save_file as safe_save_file
+from mindone.utils.modeling_patch import patch_nn_default_dtype, unpatch_nn_default_dtype
 
 from .. import __version__
 from ..utils import (
@@ -61,7 +62,9 @@
     load_state_dict,
     split_torch_state_dict_into_shards,
 )
-from .modeling_patch import patch_nn_default_dtype, restore_nn_default_dtype
+
+ms.Parameter._data = ms.Tensor.data
+ms.Parameter.data_ptr = ms.Tensor.data_ptr
 
 
 class ContextManagers:
@@ -853,12 +856,12 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                     f"{mindspore_dtype} needs to be of type `mindspore.Type`, e.g. `mindspore.float16`, but is {type(mindspore_dtype)}."
                 )
 
+        if mindspore_dtype is not None:
+            patch_nn_default_dtype(dtype=mindspore_dtype, force=True)
         with no_init_parameters():
-            if mindspore_dtype is not None:
-                patch_nn_default_dtype(dtype=mindspore_dtype, force=True)
             model = cls.from_config(config, **unused_kwargs)
-            if mindspore_dtype is not None:
-                restore_nn_default_dtype()
+        if mindspore_dtype is not None:
+            unpatch_nn_default_dtype()
 
         state_dict = None
         if not is_sharded:
@@ -915,17 +918,17 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
     def to(self, dtype: Optional[ms.Type] = None):
         for p in self.get_parameters():
             if p.dtype != dtype:
-                p.set_dtype(dtype)
+                p._data = p.to(device="CPU", dtype=dtype)
         return self
 
     def half(self):
         for p in self.get_parameters():
-            p.set_dtype(ms.float16)
+            p._data = p.to(device="CPU", dtype=ms.float16)
         return self
 
     def float(self):
         for p in self.get_parameters():
-            p.set_dtype(ms.float32)
+            p._data = p.to(device="CPU", dtype=ms.float32)
         return self
 
     def compile_repeated_blocks(self, *args, **kwargs):

mindone/transformers/modeling_patch.py

@@ -62,7 +62,8 @@
 from mindspore import Parameter, Tensor, mint, nn, ops
 from mindspore.nn import CrossEntropyLoss, Identity
 from mindspore.nn.utils import no_init_parameters
-from mindspore.ops import Cast
+
+from mindone.utils.modeling_patch import patch_nn_default_dtype, unpatch_nn_default_dtype
 
 from .activations import get_activation
 from .generation.utils import GenerationMixin
@@ -81,7 +82,6 @@
     prune_linear_layer,
 )
 from .modeling_attn_mask_utils import dtype_to_min
-from .modeling_patch import patch_nn_default_dtype, restore_nn_default_dtype
 from .utils.generic import _CAN_RECORD_REGISTRY, OutputRecorder
 from .utils.import_utils import is_sdpa_available
 
@@ -113,7 +113,8 @@
 ]
 
 logger = logging.get_logger(__name__)
-cpu_cast = Cast().set_device("CPU")
+ms.Parameter._data = ms.Tensor.data
+ms.Parameter.data_ptr = ms.Tensor.data_ptr
 
 _init_weights = True
 
@@ -377,7 +378,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:
-            state_dict[k] = ms.Parameter(cpu_cast(v.data, local_state[k].dtype), name=k)
+            v._data = v.to(device="CPU", dtype=local_state[k].dtype)
         else:
             pass  # unexpect key keeps origin dtype
     cm = silence_mindspore_logger() if is_sharded else nullcontext()
@@ -514,17 +515,17 @@ def _get_name(self):
     def to(self, dtype: Optional[ms.Type] = None):
         for p in self.get_parameters():
             if p.dtype != dtype:
-                p.set_dtype(dtype)
+                p._data = p.to(device="CPU", dtype=dtype)
         return self
 
     def float(self):
         for p in self.get_parameters():
-            p.set_dtype(ms.float32)
+            p._data = p.to(device="CPU", dtype=ms.float32)
         return self
 
     def half(self):
         for p in self.get_parameters():
-            p.set_dtype(ms.float16)
+            p._data = p.to(device="CPU", dtype=ms.float16)
         return self
 
     @property
@@ -1162,12 +1163,12 @@ def _from_config(cls, config, **kwargs):
         if "attn_implementation" in kwargs:
             config._attn_implementation = kwargs.pop("attn_implementation")
 
+        if mindspore_dtype is not None:
+            patch_nn_default_dtype(dtype=mindspore_dtype, force=True)
         with no_init_parameters():
-            if mindspore_dtype is not None:
-                patch_nn_default_dtype(dtype=mindspore_dtype, force=True)
             model = cls(config, **kwargs)
-            if mindspore_dtype is not None:
-                restore_nn_default_dtype()
+        if mindspore_dtype is not None:
+            unpatch_nn_default_dtype()
 
         # We cannot set default mindspore dtype. So we need to cast model weights after creating.
         if mindspore_dtype is not None:
@@ -2763,12 +2764,12 @@ def from_pretrained(
 
         config = copy.deepcopy(config)  # We do not want to modify the config inplace in from_pretrained.
 
+        if mindspore_dtype is not None:
+            patch_nn_default_dtype(dtype=mindspore_dtype, force=True)
         with no_init_parameters():
-            if mindspore_dtype is not None:
-                patch_nn_default_dtype(dtype=mindspore_dtype, force=True)
             model = cls(config, *model_args, **model_kwargs)
-            if mindspore_dtype is not None:
-                restore_nn_default_dtype()
+        if mindspore_dtype is not None:
+            unpatch_nn_default_dtype()
 
         # Make sure to tie the weights correctly
         model.tie_weights()

mindone/transformers/modeling_utils.py

@@ -1,4 +1,5 @@
 from .env import init_env
 from .logger import set_logger
+from .modeling_patch import patch_nn_default_dtype, unpatch_nn_default_dtype
 from .params import count_params
 from .weight_norm import WeightNorm

mindone/utils/__init__.py

@@ -40,7 +40,7 @@ def _new_init(self, *args, _orig_init=_orig_init, **kwargs):
                     setattr(attr, "__init__", _new_init)
 
 
-def restore_nn_default_dtype():
+def unpatch_nn_default_dtype():
     """
     Manually restore the original __init__ of all patched nn / mint.nn Cells.
     """