mindspore 2.6.0rc1__cp311-cp311-win_amd64.whl → 2.7.0__cp311-cp311-win_amd64.whl

mindspore/train/serialization.py

@@ -31,15 +31,14 @@ from multiprocessing import active_children
 import multiprocessing as mp
 from collections import OrderedDict
 from io import BytesIO
+from functools import partial
 
 import math
 import sys
 import time
-import google
 import numpy as np
-
-from safetensors.numpy import save_file, load_file
-from safetensors import safe_open
+from safetensors.numpy import save_file
+import google
 
 from mindspore.train.checkpoint_pb2 import Checkpoint
 from mindspore.train.mind_ir_pb2 import ModelProto as mindir_model
@@ -53,7 +52,6 @@ from mindspore.log import vlog_print
 from mindspore._checkparam import check_input_data, check_input_dataset
 from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
-from mindspore.common import np_dtype
 from mindspore.common.api import _cell_graph_executor as _executor
 from mindspore.common.api import _JitExecutor
 from mindspore.common.api import _get_parameter_layout
@@ -76,6 +74,7 @@ from mindspore.parallel.checkpoint_transform import restore_group_info_list as n
 from mindspore.parallel.checkpoint_transform import load_distributed_checkpoint as new_load_distributed_checkpoint
 from mindspore.parallel.checkpoint_transform import merge_sliced_parameter as new_merge_sliced_parameter
 from mindspore.parallel.checkpoint_transform import build_searched_strategy as new_build_searched_strategy
+from mindspore.parallel.transform_safetensors import _fast_safe_open
 from mindspore.train._utils import read_proto, get_parameter_redundancy, _progress_bar, _load_and_transform
 from mindspore._c_expression import load_mindir, _encrypt, _decrypt, _is_cipher_file, \
     split_mindir, split_dynamic_mindir
@@ -86,12 +85,9 @@ tensor_to_ms_type = {"Int8": mstype.int8, "UInt8": mstype.uint8, "Int16": mstype
                      "Float16": mstype.float16, "Float32": mstype.float32, "Float64": mstype.float64,
                      "Bool": mstype.bool_, "str": mstype.string, "BFloat16": mstype.bfloat16, "Int4": mstype.qint4x2}
 
-tensor_to_np_type = {"Int8": np.int8, "UInt8": np.uint8, "Int16": np.int16, "UInt16": np.uint16,
-                     "Int32": np.int32, "UInt32": np.uint32, "Int64": np.int64, "UInt64": np.uint64,
-                     "Float16": np.float16, "Float32": np.float32, "Float64": np.float64, "Bool": np.bool_, "str": "U"}
-
-if hasattr(np_dtype, "bfloat16"):
-    tensor_to_np_type["BFloat16"] = np_dtype.bfloat16
+_tensor_to_np_type = {"Int8": np.int8, "UInt8": np.uint8, "Int16": np.int16, "UInt16": np.uint16,
+                      "Int32": np.int32, "UInt32": np.uint32, "Int64": np.int64, "UInt64": np.uint64,
+                      "Float16": np.float16, "Float32": np.float32, "Float64": np.float64, "Bool": np.bool_, "str": "U"}
 
 np_type_convert = {"int32": np.int32, "float32": np.float32, "float16": np.float16, "float64": np.float64}
 
@@ -99,6 +95,8 @@ mindir_to_tensor_type = {1: mstype.float32, 2: mstype.uint8, 3: mstype.int8, 4:
                          5: mstype.int16, 6: mstype.int32, 7: mstype.int64, 10: mstype.float16,
                          11: mstype.float64, 12: mstype.uint32, 13: mstype.uint64}
 
+safetensors_to_mstype = {'Int4': mstype.qint4x2}
+
 _ckpt_mutex = RLock()
 
 # unit is KB
@@ -112,6 +110,21 @@ INT_64_MAX = 9223372036854775807
 cpu_cast = Cast().set_device("CPU")
 
 _ckpt_fs = FileSystem()
+_ckpt_fs_initialized = False
+
+
+def tensor_to_np_type(tensor_type_str):
+    """tensor to numpy type"""
+    if tensor_type_str == "BFloat16":
+        from mindspore.common import np_dtype
+        if not np_dtype.np_dtype_valid(True):
+            raise TypeError(
+                "The Numpy bfloat16 data type is not supported now, please ensure that the current "
+                "Numpy version is not less than the version when the mindspore is compiled, "
+                "and the major versions are same."
+            )
+        return np_dtype.bfloat16
+    return _tensor_to_np_type.get(tensor_type_str)
 
 
 def init_ckpt_file_system(fs: FileSystem):
@@ -121,8 +134,12 @@ def init_ckpt_file_system(fs: FileSystem):
     _register_basic_file_system(fs)
 
 
-# Initialize checkpoint file system
-init_ckpt_file_system(_ckpt_fs)
+def _ensure_ckpt_fs_initialized():
+    """Ensure checkpoint file system is initialized"""
+    global _ckpt_fs_initialized
+    if not _ckpt_fs_initialized:
+        init_ckpt_file_system(_ckpt_fs)
+        _ckpt_fs_initialized = True
 
 
 def _wait_async_process_save_ckpt():
@@ -272,10 +289,7 @@ def _update_param(param, new_param, strict_load):
 
         if param.data.dtype != new_param.data.dtype:
             if _type_convert(param, new_param, strict_load):
-                if new_param.data.dtype == mstype.bfloat16:
-                    new_tensor = cpu_cast(new_param.data, param.data.dtype)
-                else:
-                    new_tensor = Tensor(new_param.data.asnumpy(), param.data.dtype)
+                new_tensor = Tensor(new_param.data.asnumpy(), param.data.dtype)
                 param.set_data(new_tensor, param.sliced)
                 return
 
@@ -313,7 +327,7 @@ def _update_param(param, new_param, strict_load):
 def _type_convert(param, new_param, strict_load):
     """Whether to convert parameter's type during load checkpoint into network."""
     float_type = (mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16)
-    int_type = (mstype.int8, mstype.int16, mstype.int32, mstype.int64)
+    int_type = (mstype.int8, mstype.int16, mstype.int32, mstype.int64, mstype.qint4x2)
     if not strict_load and ({param.data.dtype, new_param.data.dtype}.issubset(float_type) or
                             {param.data.dtype, new_param.data.dtype}.issubset(int_type)):
         logger.warning(f"The type of {new_param.name}:{new_param.data.dtype} in 'parameter_dict' is different from "
@@ -359,7 +373,7 @@ def _save_weight(checkpoint_dir, model_name, iteration, params):
 
 
 def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_param_inc=False, crc_check=False,
-               format="ckpt"):
+               format="ckpt", remove_redundancy=None):
     """Execute the process of saving checkpoint into file."""
     try:
         with _ckpt_mutex:
@@ -383,9 +397,6 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
 
                     crc_num = 0
                     for name, value in data_list.items():
-                        if name == "random_op":
-                            _write_random_seed(name, value, f)
-                            continue
                         if value[0] == "mapparameter":
                             _write_mapparameter(name, value, f, map_param_inc)
                             continue
@@ -428,16 +439,19 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
             elif format == "safetensors":
                 save_dict = {}
                 crc_num = 0
+                meta_data = {"format": "ms"}
+                if remove_redundancy is not None and isinstance(remove_redundancy, bool):
+                    meta_data["remove_redundancy"] = str(remove_redundancy)
                 for name in sorted(data_list.keys()):
                     value = data_list[name]
                     if isinstance(value[2], np.ndarray):
+                        if value[1] == str(mstype.qint4x2):
+                            meta_data[name] = str(mstype.qint4x2)
                         save_dict[name] = value[2]
                     else:
-                        bytes_data = value[2].get_bytes()
-                        np_type = tensor_to_np_type.get(value[1])
-                        np_array = np.frombuffer(bytes_data, np_type)
-                        new_np_array = np_array.reshape(value[0])
-                        save_dict[name] = new_np_array
+                        if value[2].dtype == mstype.qint4x2:
+                            meta_data[name] = str(mstype.qint4x2)
+                        save_dict[name] = value[2].asnumpy()
 
                     if crc_check:
                         crc_num = binascii.crc32(bytes(name, encoding='utf-8'), crc_num)
@@ -445,10 +459,12 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
                             bytes(save_dict[name]), crc_num)
                 safetensors_save_time_start = time.time()
                 if crc_check:
-                    save_file(save_dict, tmp_name, metadata={
-                        "crc_num": str(crc_num)})
+                    meta_data.update({"crc_num": str(crc_num)})
+                if save_dict:
+                    save_file(save_dict, tmp_name, metadata=meta_data)
                 else:
                     save_file(save_dict, tmp_name)
+
                 safetensors_save_time_end = time.time()
                 cost_time = safetensors_save_time_end - safetensors_save_time_start
                 vlog_print("1", "ME", __file__, sys._getframe().f_lineno, f"Save safetensors io cost time:{cost_time}.")
@@ -457,25 +473,13 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
                                f"simultaneously modified a file.")
             elif _ckpt_fs.backend != "mindio":
                 os.rename(tmp_name, ckpt_file_name)
-            os.chmod(ckpt_file_name, stat.S_IRUSR)
+                os.chmod(ckpt_file_name, stat.S_IRUSR)
     except BaseException as e:
         logger.critical("Failed to save the checkpoint file %s. Maybe don't have the permission to write files, "
                         "or the disk space is insufficient and so on.", ckpt_file_name)
         raise e
 
 
-def _write_random_seed(name, value, f):
-    """Write random op into protobuf file."""
-    checkpoint_list = Checkpoint()
-    param_value = checkpoint_list.value.add()
-    param_value.tag = name
-    param_tensor = param_value.tensor
-    param_tensor.dims.extend(0)
-    param_tensor.tensor_type = "random_op"
-    param_tensor.tensor_content = value
-    f.write(checkpoint_list.SerializeToString())
-
-
 def _write_parameter_data(name, value, f, enc_key, plain_data, crc_num=0, crc_check=False, ckpt_total_io_time=0):
     """Write parameter data into protobuf file."""
     data_size = value[2].nbytes / 1024
@@ -599,7 +603,7 @@ def _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name, format):
     return ckpt_file_name
 
 
-def _check_load_checkpoint_upsupported_param(format, dec_key, dec_mode):
+def _check_load_checkpoint_unsupported_param(format, dec_key, dec_mode):
     """check load checkpoint unsupported param"""
     if format != "safetensors":
         return
@@ -614,7 +618,7 @@ def _check_load_checkpoint_upsupported_param(format, dec_key, dec_mode):
                              f"be set to default value '{default_value}', but got '{current_value}'.")
 
 
-def _check_save_checkpoint_upsupported_param(format, enc_key, enc_mode, map_param_inc=False, global_step_num=None):
+def _check_save_checkpoint_unsupported_param(format, enc_key, enc_mode, map_param_inc=False, global_step_num=None):
     """check save checkpoint unsupported param"""
     if format != "safetensors":
         return
@@ -644,11 +648,11 @@ def _check_async_save(async_save):
 
 
 def _async_process_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_param_inc=False,
-                        crc_check=False, format="ckpt", cond=None):
+                        crc_check=False, format="ckpt", cond=None, remove_redundancy=None):
     """Check whether the process is pulled up successfully, execute the process of saving checkpoint into file."""
     with cond:
         cond.notify()
-    _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check, format)
+    _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check, format, remove_redundancy)
 
 
 def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
@@ -729,6 +733,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
           <https://mindspore.cn/tutorials/en/master/beginner/save_load.html#saving-and-loading-the-model-weight>`_
     """
     start_save_time = time.time()
+    _ensure_ckpt_fs_initialized()
     ckpt_file_name = _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name, format)
     integrated_save = Validator.check_bool(integrated_save)
     async_save = _check_async_save(async_save)
@@ -739,7 +744,9 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
     map_param_inc = kwargs.get('incremental', False)
     logger.info("Execute the process of saving checkpoint files.")
     global_step_num = kwargs.get('global_step_num', None)
-    _check_save_checkpoint_upsupported_param(format, enc_key, enc_mode, map_param_inc, global_step_num)
+    remove_redundancy = kwargs.get('remove_redundancy', None)
+    remove_redundancy = Validator.check_isinstance("remove_redundancy", remove_redundancy, (type(None), bool))
+    _check_save_checkpoint_unsupported_param(format, enc_key, enc_mode, map_param_inc, global_step_num)
 
     if append_dict and "__exception_save__" in append_dict:
         s1 = mindspore.hal.Stream()
@@ -768,16 +775,6 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
     data_list_np = OrderedDict()
     with _ckpt_mutex:
         for param in save_obj:
-            if param["name"] == "random_op":
-                if os.getenv("AITURBO") == "1":
-                    data_list_np["random_op"] = []
-                    data_list_np["random_op"].append(param["data"])
-                    if crc_check:
-                        bytes_value = bytes(data_list_np[key][0])
-                        data_list_np[key].append(binascii.crc32(bytes_value))
-                else:
-                    data_list["random_op"] = param["data"]
-                continue
             key = param["name"]
             data_list[key] = []
             data_list_np[key] = []
@@ -841,7 +838,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                 while process_flag:
                     process = ctx.Process(target=_async_process_save,
                                           args=(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check,
-                                                format, cond), daemon=True, name="asyn_save_ckpt")
+                                                format, cond, remove_redundancy), daemon=True, name="asyn_save_ckpt")
                     process.start()
                     with cond:
                         wait_flag = cond.wait(timeout=5)
@@ -854,11 +851,12 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
             data_copy = copy.deepcopy(data_list)
             _wait_async_thread_save_ckpt()
             thr = Thread(target=_exec_save,
-                         args=(ckpt_file_name, data_copy, enc_key, enc_mode, map_param_inc, crc_check, format),
+                         args=(ckpt_file_name, data_copy, enc_key, enc_mode, map_param_inc, crc_check, format,
+                               remove_redundancy),
                          name="asyn_save_ckpt")
             thr.start()
     else:
-        _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check, format)
+        _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check, format, remove_redundancy)
 
     mstx.range_end(range_id)
     logger.info("Saving checkpoint process is finished.")
@@ -926,10 +924,13 @@ def _convert_dict_to_param_dict(save_obj, choice_func):
     """Convert a dict of Parameter to param_list."""
     param_list = []
     for (key, value) in save_obj.items():
-        if isinstance(key, str) and (isinstance(value, (Parameter, str)) or _is_buffer_type(value)):
+        if isinstance(key, str):
             if choice_func is not None and not choice_func(key):
                 continue
-            each_param = {"name": key, "data": value}
+            if isinstance(value, np.ndarray):
+                each_param = {"name": key, "data": Parameter(Tensor.from_numpy(value))}
+            if isinstance(value, (Parameter, str)) or _is_buffer_type(value):
+                each_param = {"name": key, "data": value}
             param_list.append(each_param)
         else:
             raise TypeError(f"For save_checkpoint, when save_obj is made up by dict, the key should be str and"
@@ -941,16 +942,12 @@ def _convert_dict_to_param_dict(save_obj, choice_func):
 def _convert_cell_param_and_names_to_dict(save_obj, choice_func, is_parallel_mode):
     """Convert cell.parameters_and_names to OrderedDict."""
     param_dict = OrderedDict()
+    is_graph_mode = context.get_context('mode') == context.GRAPH_MODE
     for _, param in save_obj.parameters_and_names():
-        if param.name.startswith("accu_grads") or param.name.endswith("expert_load"):
-            continue
-        not_sliced = not param.sliced
-        is_graph_mode = context.get_context('mode') == context.GRAPH_MODE
         # All parameters are initialized immediately under PyNative mode, skip this judgement.
-        judgment = not_sliced or param.has_init
         if param.param_info.is_pipeline_shared_param:
             continue
-        if is_graph_mode and is_parallel_mode and judgment:
+        if is_parallel_mode and is_graph_mode and (not param.sliced or param.has_init):
             continue
         if choice_func is not None and not choice_func(param.name):
             continue
@@ -974,12 +971,6 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
     if not is_parallel_mode:
         save_obj.init_parameters_data()
     param_dict = _convert_cell_param_and_names_to_dict(save_obj, choice_func, is_parallel_mode)
-    if append_dict and "random_op" in append_dict:
-        phase = 'train' + '.' + str(save_obj.create_time) + '.' + str(id(save_obj)) + '.' + save_obj.arguments_key
-        if phase in save_obj.compile_cache and _executor.has_compiled(phase):
-            random_byte = _executor._graph_executor.get_random_status(phase)
-            param_list.append({"name": "random_op", "data": random_byte})
-            append_dict.pop("random_op")
     for (key, value) in param_dict.items():
         each_param = {"name": key}
         if isinstance(value, MapParameter):
@@ -1002,15 +993,14 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
             param_data.append(str(param_tensor.dtype))
             param_data.append(value.key)
         else:
-            param_data = value.data
             if append_dict and "__exception_save__" in append_dict:
                 param_data = Tensor(Tensor_.move_to(value, "CPU", False))
+            else:
+                param_data = Tensor(value.data)
 
             # in automatic model parallel scenario, some parameters were split to all the devices,
             # which should be combined before saving
             if key in parameter_layout_dict:
-                if not append_dict or "__exception_save__" not in append_dict:
-                    param_data = Tensor(value.data)
                 param_data = _get_merged_param_data(save_obj, parameter_layout_dict, key, param_data,
                                                     integrated_save)
 
@@ -1215,12 +1205,26 @@ def _check_param_type(param_config, key, target_type, requested):
     return None
 
 
+def _check_remove_redundancy(remove_redundancy, f):
+    """Check whether remove_redundancy is consistent with the safetensors file."""
+    if f.metadata() is not None and "remove_redundancy" in f.metadata().keys():
+        if f.metadata()["remove_redundancy"] == "True" and not remove_redundancy:
+            logger.warning("For 'load_checkpoint', the safetensors file is deduplicated, "
+                           "but remove_redundancy is set to False.")
+            return True
+        if f.metadata()["remove_redundancy"] == "False" and remove_redundancy:
+            logger.warning("For 'load_checkpoint', the safetensors file is non-deduplicated, "
+                           "but remove_redundancy is set to True.")
+            return False
+    return remove_redundancy
+
+
 def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter_prefix, choice_func, dec_key,
-                          dec_mode, crc_check, format):
+                          dec_mode, crc_check, format, remove_redundancy):
     """load parameter into parameter_dict"""
     ckpt_file_name = _check_ckpt_file_name(ckpt_file_name, format)
     if format == "safetensors":
-        with safe_open(ckpt_file_name, framework='np') as f:
+        with _fast_safe_open(ckpt_file_name, framework='np') as f:
             cal_crc_num = 0
             total_io_cost_time = 0
             for k in sorted(f.keys()):
@@ -1234,8 +1238,13 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                 io_end_time = time.time()
                 io_cost_time = io_end_time - io_start_time
                 total_io_cost_time += io_cost_time
-                parameter_dict[k] = Parameter(Tensor.from_numpy(value))
-
+                if f.metadata() is not None and k in f.metadata().keys():
+                    sf_dtype = f.metadata()[k]
+                    ms_dtype = safetensors_to_mstype[sf_dtype]
+                    parameter_dict[k] = Parameter(Tensor(value, dtype=ms_dtype))
+                else:
+                    parameter_dict[k] = Parameter(Tensor.from_numpy(value))
+            remove_redundancy = _check_remove_redundancy(remove_redundancy, f)
             vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
                        f"Load safetensors io cost time:{total_io_cost_time}.")
             if crc_check:
@@ -1248,7 +1257,7 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                     if cal_crc_num != crc_num:
                         raise ValueError("For 'load_checkpoint', the crc check has failed. "
                                          "Please check whether the ckpt file is damaged.")
-        return
+        return remove_redundancy
     checkpoint_list = _parse_ckpt_proto(ckpt_file_name, dec_key, dec_mode, crc_check)
     try:
         param_data_list = []
@@ -1261,9 +1270,6 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
             logger.warning("For load_checkpoint, this parameter `filter_prefix` will be deprecated, "
                            "please use `choice_func` instead.")
         for element_id, element in enumerate(checkpoint_list.value):
-            if element.tag == "random_op":
-                parameter_dict["random_op"] = element.tensor.tensor_content
-                continue
             if not _whether_load_param(specify_prefix, filter_prefix, element.tag):
                 continue
             if specify_prefix is None and filter_prefix is None and \
@@ -1278,11 +1284,7 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                 continue
             data = element.tensor.tensor_content
             data_type = element.tensor.tensor_type
-            np_type = tensor_to_np_type.get(data_type)
             ms_type = tensor_to_ms_type[data_type]
-            if data_type == 'str':
-                str_length = int(len(data) / 4)
-                np_type = np_type + str(str_length)
             param_data_list.append(data)
             if (element_id == len(checkpoint_list.value) - 1) or \
                     (element.tag != checkpoint_list.value[element_id + 1].tag):
@@ -1290,6 +1292,8 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                 param_data_list.clear()
                 dims = element.tensor.dims
                 if data_type == 'str':
+                    str_length = int(len(data) / 4)
+                    np_type = "U" + str(str_length)
                     str_value = np.frombuffer(new_data, np_type)
                     parameter_dict[element.tag] = str(str_value[0])
                 else:
@@ -1301,6 +1305,7 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                     _offload_if_config(parameter)
 
         logger.info("Loading checkpoint files process is finished.")
+        return remove_redundancy
 
     except BaseException as e:
         logger.critical("Failed to load the checkpoint file '%s'.", ckpt_file_name)
@@ -1320,6 +1325,9 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
           And using either of those two args will override `choice_func` at the same time.
         - If none of the parameters are loaded from checkpoint file, it will throw ValueError.
         - When loading a checkpoint that has removed redundancy, the network should be compiled.
+        - When `net` is not None, it will verify whether the `remove_redundancy` parameter matches the
+          deduplication flag in the loaded safetensors file. If they are different, load the file according to
+          the deduplication flag in the file.
 
     Args:
         ckpt_file_name (str): Checkpoint file name.
@@ -1392,13 +1400,14 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
     """
     start_load_time = time.time()
     vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin load checkpoint.")
+    _ensure_ckpt_fs_initialized()
     specify_prefix = _check_prefix(specify_prefix)
     filter_prefix = _check_prefix(filter_prefix)
     dec_key = Validator.check_isinstance('dec_key', dec_key, (type(None), bytes))
     dec_mode = Validator.check_isinstance('dec_mode', dec_mode, str)
     crc_check = Validator.check_isinstance('crc_check', crc_check, bool)
     remove_redundancy = Validator.check_isinstance('remove_redundancy', remove_redundancy, bool)
-    _check_load_checkpoint_upsupported_param(format, dec_key, dec_mode)
+    _check_load_checkpoint_unsupported_param(format, dec_key, dec_mode)
     logger.info("Execute the process of loading checkpoint files.")
 
     parameter_dict = {}
@@ -1424,8 +1433,8 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
                                      f"passed the CRC check and has been corrupted.")
                 parameter_dict[key] = Parameter(Tensor(value[0]), name=key)
     else:
-        _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter_prefix, choice_func, dec_key,
-                              dec_mode, crc_check, format)
+        remove_redundancy = _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter_prefix,
+                                                  choice_func, dec_key, dec_mode, crc_check, format, remove_redundancy)
 
     if not parameter_dict:
         raise ValueError(f"The loaded parameter dict is empty after filter or specify, please check whether "
@@ -1672,9 +1681,22 @@ def _check_load_param_into_net(net, parameter_dict):
         msg = ("For 'load_param_into_net', the argument 'parameter_dict' should be a dict, "
                "but got {}.".format(type(parameter_dict)))
         raise TypeError(msg)
-    if "random_op" in parameter_dict.keys():
-        net._add_attr("random_op_snapshot", parameter_dict["random_op"])
-        parameter_dict.pop("random_op")
+    for key, value in parameter_dict.items():
+        if not isinstance(key, str) or not isinstance(value, (Parameter, str, list)):
+            logger.critical("Load parameters into net failed.")
+            msg = ("For 'parameter_dict', the element in the argument 'parameter_dict' should be a "
+                   "'str' and 'Parameter' , but got {} and {}.".format(type(key), type(value)))
+            raise TypeError(msg)
+
+
+def _check_remove_redundancy_net(net):
+    """Check whether the network is compiled with the remove_redundancy feature."""
+    if get_group_size() == 1:
+        raise TypeError(f"The deduplication feature for loading checkpoint can only be used "
+                        f"in parallel scenarios, but got stand_alone.")
+    if not net.compile_cache and not net.parameter_layout_dict:
+        raise ValueError("When loading a parameter dict that has removed redundancy, "
+                         "the network should be compiled.")
 
 
 def load_param_into_net(net, parameter_dict, strict_load=False, remove_redundancy=False):
@@ -1721,18 +1743,14 @@ def load_param_into_net(net, parameter_dict, strict_load=False, remove_redundanc
           <https://mindspore.cn/tutorials/en/master/beginner/save_load.html#saving-and-loading-the-model-weight>`_
     """
     _check_load_param_into_net(net, parameter_dict)
-    for key, value in parameter_dict.items():
-        if not isinstance(key, str) or not isinstance(value, (Parameter, str, list)):
-            logger.critical("Load parameters into net failed.")
-            msg = ("For 'parameter_dict', the element in the argument 'parameter_dict' should be a "
-                   "'str' and 'Parameter' , but got {} and {}.".format(type(key), type(value)))
-            raise TypeError(msg)
 
     strict_load = Validator.check_bool(strict_load)
     remove_redundancy = Validator.check_isinstance('remove_redundancy', remove_redundancy, bool)
     logger.info("Execute the process of loading parameters into net.")
     param_not_load = []
+    param_loaded = set()
     ckpt_not_load = list(parameter_dict.keys())
+    is_parallel_mode = _is_auto_parallel_mode(net)
     for _, param in net.parameters_and_names():
         if param.param_info.is_pipeline_shared_param:
             continue
@@ -1748,22 +1766,23 @@ def load_param_into_net(net, parameter_dict, strict_load=False, remove_redundanc
             if hasattr(param, "init_param") and not param.init_param:
                 param.init_param = True
             ckpt_not_load.remove(param.name)
+            param_loaded.add(param.name)
         else:
+            if param.name.startswith("accu_grads"):
+                continue
+            if param.param_info.is_pipeline_shared_param:
+                continue
+            if is_parallel_mode and not param.sliced:
+                continue
             param_not_load.append(param.name)
 
     if param_not_load and not strict_load:
         _load_dismatch_prefix_params(net, parameter_dict, param_not_load, strict_load)
 
     if remove_redundancy:
-        if get_group_size() == 1:
-            raise TypeError(f"The deduplication feature for loading checkpoint can only be used "
-                            f"in parallel scenarios, but got stand_alone.")
-        if not net.compile_cache and not net.parameter_layout_dict:
-            raise ValueError("When loading a parameter dict that has removed redundancy, "
-                             "the network should be compiled.")
+        _check_remove_redundancy_net(net)
         param_layout = net.parameter_layout_dict
-        _single_parameter_broadcast(net, param_layout, param_not_load)
-        mindspore.hal.synchronize()
+        _single_parameter_broadcast(net, param_layout, param_not_load, param_loaded)
 
     logger.info("Loading parameters into net is finished.")
     if param_not_load:
@@ -1878,9 +1897,10 @@ def _save_graph(network, file_name):
         file_name (str): Graph file name into which the graph will be saved.
     """
     logger.info("Execute the process of saving graph.")
-
     file_name = os.path.realpath(file_name)
     graph_pb = network.get_func_graph_proto()
+    if os.path.isfile(file_name) and graph_pb:
+        os.remove(file_name)
     if graph_pb:
         with open(file_name, "wb") as f:
             os.chmod(file_name, stat.S_IRUSR | stat.S_IWUSR)
@@ -2193,6 +2213,11 @@ def _save_onnx(net, file_name, *inputs, **kwargs):
         file_name += ".onnx"
     if os.path.exists(file_name):
         os.chmod(file_name, stat.S_IWUSR)
+    else:
+        dir_path = os.path.dirname(file_name)
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path, mode=0o700, exist_ok=True)
+        os.chmod(dir_path, 0o700)
     with open(file_name, 'wb') as f:
         f.write(onnx_stream)
         os.chmod(file_name, stat.S_IRUSR)
@@ -2242,7 +2267,7 @@ def _get_data_file(is_encrypt, kwargs, data_file_name):
     if is_encrypt():
         place_holder_data = _encrypt(place_holder_data, len(place_holder_data), kwargs["enc_key"],
                                      len(kwargs["enc_key"]), kwargs["enc_mode"])
-    parameter_size = (offset / 1024)
+    parameter_size = offset / 1024
     try:
         f = open(data_file_name, "wb")
         f.write(place_holder_data)
@@ -2284,9 +2309,11 @@ def _split_save(net_dict, model, file_name, is_encrypt, **kwargs):
     external_local = os.path.join(file_prefix + "_variables", "data_" + str(index))
     data_file_name = os.path.join(dirname, external_local)
     f, parameter_size, offset = _get_data_file(is_encrypt, kwargs, data_file_name)
+
+    round = 0
+    names = []
+
     try:
-        round = 0
-        names = []
         for param_proto in model.graph.parameter:
             name = param_proto.name[param_proto.name.find(":") + 1:]
             names.append((name, param_proto))
@@ -2587,7 +2614,7 @@ def parse_print(print_file_name):
                 dims = print_.tensor.dims
                 data_type = print_.tensor.tensor_type
                 data = print_.tensor.tensor_content
-                np_type = tensor_to_np_type.get(data_type)
+                np_type = tensor_to_np_type(data_type)
                 param_data = np.fromstring(data, np_type)
                 ms_type = tensor_to_ms_type.get(data_type)
                 if dims and dims != [0]:
@@ -2730,28 +2757,35 @@ def convert_model(mindir_file, convert_file, file_format):
         export(net, *net_input, file_name=convert_file, file_format=file_format)
 
 
-def _transform_tensor_to_numpy(path, name_map=None):
-    return _load_and_transform(path, name_map, mindspore.load_checkpoint, lambda v, new_name: v.asnumpy())
+def _load_ckpt_to_new_name_map(path, name_map=None):
+    return _load_and_transform(path, name_map, mindspore.load_checkpoint, None)
 
 
-def _transform_numpy_to_tensor(path, name_map=None):
-    return _load_and_transform(path, name_map, load_file, lambda v, new_name: mindspore.Parameter(v, name=new_name))
+def _load_sf_to_new_name_map(path, name_map=None):
+    load_func = partial(mindspore.load_checkpoint, format="safetensors")
+    return _load_and_transform(path, name_map, load_func, None)
 
 
 def _process_file(file_info):
     cur_ckpt_path, name_map, save_path, file = file_info
-    param_dict_numpy = _transform_tensor_to_numpy(cur_ckpt_path, name_map)
+    if name_map is not None:
+        param_dict = _load_ckpt_to_new_name_map(cur_ckpt_path, name_map)
+    else:
+        param_dict = mindspore.load_checkpoint(cur_ckpt_path)
     safetensors_filename = file.replace(".ckpt", ".safetensors")
     dst_file = os.path.join(save_path, safetensors_filename)
-    save_file(param_dict_numpy, dst_file)
+    mindspore.save_checkpoint(param_dict, dst_file, format='safetensors')
 
 
 def _process_file_safetensors(file_info):
     cur_safe_path, name_map, save_path, file = file_info
-    param_dict_tensor = _transform_numpy_to_tensor(cur_safe_path, name_map)
+    if name_map is not None:
+        param_dict = _load_sf_to_new_name_map(cur_safe_path, name_map)
+    else:
+        param_dict = mindspore.load_checkpoint(cur_safe_path, format="safetensors")
     ckpt_filename = file.replace(".safetensors", ".ckpt")
     dst_file = os.path.join(save_path, ckpt_filename)
-    mindspore.save_checkpoint(param_dict_tensor, dst_file)
+    mindspore.save_checkpoint(param_dict, dst_file)
 
 
 def _gather_safetensors_tasks(file_path, save_path, file_name_regex, name_map):
@@ -2862,10 +2896,14 @@ def ckpt_to_safetensors(file_path, save_path=None, name_map=None, file_name_rege
         if save_path and not os.path.exists(save_path):
             os.makedirs(save_path, exist_ok=True)
 
-        param_dict_numpy = _transform_tensor_to_numpy(file_path, name_map)
+        if name_map is not None:
+            param_dict = _load_ckpt_to_new_name_map(file_path, name_map)
+        else:
+            param_dict = mindspore.load_checkpoint(file_path)
+
         safetensors_filename = os.path.basename(file_path).replace(".ckpt", ".safetensors")
         dst_file = os.path.join(save_path if save_path else os.path.dirname(file_path), safetensors_filename)
-        save_file(param_dict_numpy, dst_file)
+        mindspore.save_checkpoint(param_dict, dst_file, format='safetensors')
 
 
 def safetensors_to_ckpt(file_path, save_path=None, name_map=None, file_name_regex=None, processes_num=1):
@@ -2924,10 +2962,14 @@ def safetensors_to_ckpt(file_path, save_path=None, name_map=None, file_name_rege
         if save_path and not os.path.exists(save_path):
             os.makedirs(save_path, exist_ok=True)
 
-        param_dict_tensor = _transform_numpy_to_tensor(file_path, name_map)
+        if name_map is not None:
+            param_dict = _load_sf_to_new_name_map(file_path, name_map)
+        else:
+            param_dict = mindspore.load_checkpoint(file_path, format="safetensors")
+
         ckpt_filename = os.path.basename(file_path).replace(".safetensors", ".ckpt")
         dst_file = os.path.join(save_path if save_path else os.path.dirname(file_path), ckpt_filename)
-        mindspore.save_checkpoint(param_dict_tensor, dst_file)
+        mindspore.save_checkpoint(param_dict, dst_file)
 
 
 def restore_group_info_list(group_info_file_name):