mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.1__cp310-cp310-win_amd64.whl

mindspore/train/serialization.py

@@ -21,10 +21,12 @@ import binascii
 import copy
 import json
 import os
+import re
 import shutil
 import stat
 import threading
 from threading import Thread, RLock
+from multiprocessing import Process
 from collections import defaultdict, OrderedDict
 from io import BytesIO
 
@@ -58,21 +60,25 @@ from mindspore.common.file_system import FileSystem, _register_basic_file_system
 from mindspore.communication.management import get_rank, get_group_size
 from mindspore.experimental import MapParameter
 from mindspore.ops import Cast
-from mindspore.parallel._cell_wrapper import get_allgather_cell
+from mindspore.parallel._cell_wrapper import get_allgather_cell, _single_parameter_broadcast
 from mindspore.parallel._tensor import _load_tensor, _get_tensor_strategy, _get_tensor_slice_index
 from mindspore.parallel._tensor import _reshape_param_data, _reshape_param_data_with_weight
-from mindspore.parallel._utils import _infer_rank_list, _remove_repeated_slices, _is_in_auto_parallel_mode
+from mindspore.parallel._utils import _infer_rank_list, _remove_repeated_slices, _is_in_auto_parallel_mode, \
+    _get_device_num
+from mindspore.parallel._auto_parallel_context import _get_auto_parallel_context
 from mindspore.parallel._parallel_serialization import _convert_to_list, _convert_to_layout, _build_searched_strategy, \
-    _restore_group_info_list
+    _restore_group_info_list, _get_param_list_when_first_dim_sharded
 from mindspore.parallel._ps_context import _set_checkpoint_load_status, _store_warm_up_ptr_by_tensor, \
     _store_warm_up_ptr_by_tensor_list, _cache_enable
 from mindspore.parallel.checkpoint_transform import sync_pipeline_shared_parameters
-from mindspore.train._utils import read_proto
+from mindspore.parallel.transform_safetensors import _load_parallel_checkpoint, _get_device_num_from_strategy, \
+    _extract_pipeline_stage_num
+from mindspore.train._utils import read_proto, get_parameter_redundancy
 from mindspore._c_expression import load_mindir, _encrypt, _decrypt, _is_cipher_file, dynamic_obfuscate_mindir, \
     split_mindir, split_dynamic_mindir
 from mindspore.common.generator import Generator
-from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
-from mindspore.parallel.parameter_broadcast import parameter_broadcast
+from safetensors.numpy import save_file
+from safetensors import safe_open
 from ..ops.operations._opaque_predicate_registry import add_opaque_predicate, clean_funcs
 
 tensor_to_ms_type = {"Int8": mstype.int8, "UInt8": mstype.uint8, "Int16": mstype.int16, "UInt16": mstype.uint16,
@@ -116,6 +122,68 @@ def init_ckpt_file_system(fs: FileSystem):
 init_ckpt_file_system(_ckpt_fs)
 
 
+def _get_cur_rank_dp(parameter_layout_dict):
+    """ Get dp and tp from layout dict. """
+    pp_num = _get_auto_parallel_context("pipeline_stages")
+    dev_num = _get_device_num()
+    global_rank = get_rank()
+    pipe_size = dev_num // pp_num
+    initial_rank = (global_rank // pipe_size) * pipe_size
+    parameter_redundancy_dict = get_parameter_redundancy(
+        parameter_layout_dict, initial_rank)
+    value_len = sys.maxsize
+    min_value = ()
+    for key, value in parameter_redundancy_dict.items():
+        if "accu_grads" in key or "inputs" in key:
+            continue
+        for item in value:
+            if len(item) < value_len and global_rank in item:
+                value_len = len(item)
+                min_value = item
+    return min_value
+
+
+def get_ckpt_path_with_strategy(cur_ckpt_path, cur_strategy_path):
+    """
+    Find available checkpoint file path from all backup checkpoint files of current rank.
+    It suppose that checkpoint path contains substring 'rank_{rank_id}' which is used to
+    distinguish between different path.If cur_ckpt_path doesn't have 'rank_{rank_id}' substring, will return
+    cur_ckpt_path itself when cur_ckpt_path is exist, otherwise return None.
+
+    Note:
+       This API must be called after the communication is initialized because the cluster information
+       needs to be obtained internally.
+
+    Args:
+        cur_ckpt_path (str): the checkpoint file path which cur rank needs.
+        cur_strategy_path (str): strategy file path for current rank.
+
+    Returns:
+        - new_ckpt_file (String), if found available checkpoint file , return it.
+        - None, if not found available checkpoint, return None.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore.communication import init
+        >>> from mindspore import get_ckpt_path_with_strategy
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.DATA_PARALLEL, gradients_mean=True)
+        >>> init()
+        >>> ckpt_file= "./rank_5/iteration-1_40.ckpt"
+        >>> strategy_file = "./src_pipeline_strategys/src_strategy_5.ckpt"
+        >>> ckpt_file_new = get_ckpt_path_with_strategy(ckpt_file, strategy_file)
+        >>> print(ckpt_file_new)
+    """
+    dp = _get_cur_rank_dp(cur_strategy_path)
+    pattern = r'rank_\d+'
+    for i in dp:
+        new_ckpt_path = re.sub(pattern, f"rank_{str(i)}", cur_ckpt_path)
+        if not os.path.isfile(new_ckpt_path):
+            continue
+        return new_ckpt_path
+    return None
+
+
 class ParamDictFuture:
     def __init__(self, executor, param_dict_future):
         self.executor = executor
@@ -252,57 +320,72 @@ def _save_weight(checkpoint_dir, model_name, iteration, params):
         logger.warning(f"Checkpoint dir: '{checkpoint_dir}' is not existed.")
 
 
-def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_param_inc=False, crc_check=False):
+def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_param_inc=False, crc_check=False,
+               format="ckpt"):
     """Execute the process of saving checkpoint into file."""
     try:
         with _ckpt_mutex:
+            file_name_list = list(os.path.splitext(ckpt_file_name))
+            file_name_list[1] = file_name_list[1].replace(f".{format}", ".tmp")
+            tmp_name = ''.join(file_name_list)
             if os.path.exists(ckpt_file_name):
                 os.chmod(ckpt_file_name, stat.S_IWUSR)
                 os.remove(ckpt_file_name)
-            with _ckpt_fs.create(ckpt_file_name, *_ckpt_fs.create_args) as f:
-                plain_data = None
-                if enc_key is not None:
-                    plain_data = BytesIO()
-
-                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
-                    if value[0] == "offload_parameter":
-                        new_value = value[1:]
-                        new_value[2] = value[3]
-                        _write_parameter_bytes_data(name, new_value, f, enc_key, plain_data)
-                        _offload_if_config(value[3])
-                        continue
-                    if value[1] == "str":
-                        crc_num = _write_parameter_data(name, value, f, enc_key, plain_data, crc_num, crc_check)
-                        continue
-                    if isinstance(value[2], np.ndarray):
-                        crc_num = _write_parameter_data(name, value, f, enc_key, plain_data, crc_num, crc_check)
-                        continue
-                    if isinstance(value[2], Tensor) and hasattr(value[2], "slice_num") and value[2].slice_num > 1:
-                        _write_hugeparameter(name, value, f)
-                        continue
-
-                    crc_num = _write_parameter_bytes_data(name, value, f, enc_key, plain_data, crc_num, crc_check)
-
-                if enc_key is not None:
-                    plain_data.seek(0)
-                    max_block_size = ENCRYPT_BLOCK_SIZE * 1024
-                    block_data = plain_data.read(max_block_size)
-                    while block_data:
-                        f.write(_encrypt(block_data, len(block_data), enc_key, len(enc_key), enc_mode))
+            if os.path.exists(tmp_name):
+                os.chmod(tmp_name, stat.S_IWUSR)
+                os.remove(tmp_name)
+            if format == "ckpt":
+                with _ckpt_fs.create(tmp_name, *_ckpt_fs.create_args) as f:
+                    plain_data = None
+                    if enc_key is not None:
+                        plain_data = BytesIO()
+
+                    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
+                        if value[0] == "offload_parameter":
+                            new_value = value[1:]
+                            new_value[2] = value[3]
+                            _write_parameter_bytes_data(name, new_value, f, enc_key, plain_data)
+                            _offload_if_config(value[3])
+                            continue
+                        if value[1] == "str":
+                            crc_num = _write_parameter_data(name, value, f, enc_key, plain_data, crc_num, crc_check)
+                            continue
+                        if isinstance(value[2], np.ndarray):
+                            crc_num = _write_parameter_data(name, value, f, enc_key, plain_data, crc_num, crc_check)
+                            continue
+                        if isinstance(value[2], Tensor) and hasattr(value[2], "slice_num") and value[2].slice_num > 1:
+                            _write_hugeparameter(name, value, f)
+                            continue
+
+                        crc_num = _write_parameter_bytes_data(name, value, f, enc_key, plain_data, crc_num, crc_check)
+
+                    if enc_key is not None:
+                        plain_data.seek(0)
+                        max_block_size = ENCRYPT_BLOCK_SIZE * 1024
                         block_data = plain_data.read(max_block_size)
-
-                if crc_check:
-                    f.write('crc_num'.encode() + crc_num.to_bytes(10, byteorder='big'))
-
+                        while block_data:
+                            f.write(_encrypt(block_data, len(block_data), enc_key, len(enc_key), enc_mode))
+                            block_data = plain_data.read(max_block_size)
+                    if crc_check:
+                        f.write('crc_num'.encode() + crc_num.to_bytes(10, byteorder='big'))
+            elif format == "safetensors":
+                save_dict = {}
+                for name, value in data_list.items():
+                    save_dict[name] = value[2].asnumpy()
+                save_file(save_dict, tmp_name)
+            if not os.path.exists(tmp_name):
+                logger.warning(f"Rename failed, can't find {tmp_name}, it is possible that multiple processes have "
+                               f"simultaneously modified a file.")
+            else:
+                os.rename(tmp_name, ckpt_file_name)
             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)
@@ -415,8 +498,11 @@ def _write_hugeparameter(name, value, f):
         offset += numpy_data.shape[0]
 
 
-def _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name):
+def _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name, format):
     """Check save_obj and ckpt_file_name for save_checkpoint."""
+    if format not in ["safetensors", "ckpt"]:
+        raise ValueError(f"For 'save_checkpoint', the format must be "
+                         f"'safetensors' or 'ckpt', but got {format}.")
     if not isinstance(save_obj, (nn.Cell, list, dict)):
         raise TypeError("For 'save_checkpoint', the parameter 'save_obj' must be nn.Cell, list or dict, "
                         "but got {}.".format(type(save_obj)))
@@ -424,18 +510,26 @@ def _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name):
         raise TypeError("For 'save_checkpoint', the parameter {} for checkpoint file name is invalid,"
                         "'ckpt_file_name' must be "
                         "string, but got {}.".format(ckpt_file_name, type(ckpt_file_name)))
-    ckpt_file_name = os.path.abspath(ckpt_file_name)
+    ckpt_file_name = os.path.realpath(ckpt_file_name)
     if os.path.isdir(ckpt_file_name):
         raise IsADirectoryError("For 'save_checkpoint', the parameter `ckpt_file_name`: {} is a directory, "
                                 "it must be a file name.".format(ckpt_file_name))
-    if not ckpt_file_name.endswith('.ckpt'):
-        ckpt_file_name += ".ckpt"
+    if not ckpt_file_name.endswith(format):
+        ckpt_file_name += f".{format}"
     return ckpt_file_name
 
 
+def _check_format_and_other_params(format, enc_key, enc_mode, crc_check=False, async_save=False, map_param_inc=False,
+                                   global_step_num=None):
+    param_not_default = (enc_key is not None or enc_mode != "AES-GCM" or crc_check or async_save
+                         or map_param_inc or global_step_num is not None)
+    if format == "safetensors" and param_not_default:
+        raise ValueError("For 'save_checkpoint', when format is 'safetensors', other param must be default.")
+
+
 def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                     async_save=False, append_dict=None, enc_key=None, enc_mode="AES-GCM", choice_func=None,
-                    crc_check=False, **kwargs):
+                    crc_check=False, format="ckpt", **kwargs):
     r"""
     Save checkpoint to a specified file.
 
@@ -465,6 +559,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                                  be saved. Default: ``None`` .
         crc_check (bool) : Whether to perform crc32 calculation when saving checkpoint and save the calculation
             result to the file. Default: ``False`` .
+        format (str): Format of the output file, can be "ckpt" or "safetensors". Default: "ckpt".
         kwargs (dict): Configuration options dictionary.
 
     Raises:
@@ -498,7 +593,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
         - `Saving and Loading the Model - Saving and Loading the Model Weight
           <https://mindspore.cn/tutorials/en/master/beginner/save_load.html#saving-and-loading-the-model-weight>`_
     """
-    ckpt_file_name = _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name)
+    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 = Validator.check_bool(async_save)
     append_dict = _check_append_dict(append_dict)
@@ -508,10 +603,19 @@ 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_format_and_other_params(format, enc_key, enc_mode, crc_check, async_save, map_param_inc, global_step_num)
 
-    save_obj = _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choice_func)
+    if append_dict and "__exception_save__" in append_dict:
+        s1 = mindspore.hal.Stream()
+        with mindspore.hal.StreamCtx(s1):
+            save_obj = _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choice_func)
+        s1.synchronize()
+    else:
+        save_obj = _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choice_func)
 
     if append_dict:
+        if "__exception_save__" in append_dict:
+            del append_dict["__exception_save__"]
         append_info_list = []
         for k_name, value in append_dict.items():
             if isinstance(value, Generator):
@@ -527,12 +631,17 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
         for param in save_obj:
             if param["name"] == "random_op":
                 if os.getenv("AITURBO") == "1":
-                    data_list_np["random_op"] = param["data"]
+                    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] = []
             if isinstance(param["data"], MapParameter):
                 data_list[param["name"]].append("mapparameter")
                 data_list[param["name"]].append(param["data"])
@@ -546,7 +655,10 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
 
             if isinstance(param["data"], str):
                 if os.getenv("AITURBO") == "1":
-                    data_list_np[key] = np.array(param["data"])
+                    data_list_np[key].append(np.array(param["data"]))
+                    if crc_check:
+                        bytes_value = data_list_np[key][0].tobytes()
+                        data_list_np[key].append(binascii.crc32(bytes_value))
                 else:
                     data_list[key].append([0])
                     data_list[key].append('str')
@@ -556,7 +668,10 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                 if isinstance(param["data"], Parameter):
                     param["data"].init_data()
                 if os.getenv("AITURBO") == "1":
-                    data_list_np[key] = param["data"].asnumpy()
+                    data_list_np[key].append(param["data"].asnumpy())
+                    if crc_check:
+                        bytes_value = data_list_np[key][0].tobytes()
+                        data_list_np[key].append(binascii.crc32(bytes_value))
                 else:
                     dims = []
                     for dim in param['data'].shape:
@@ -568,16 +683,17 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                     data_list[key].append(data)
 
     if os.getenv("AITURBO") == "1":
-        import aiturbo
+        from aiturbo.checkpoint import aiturbo_mindspore as aiturbo
         ckpt_name = os.path.basename(ckpt_file_name)
-        aiturbo.save_ckpt(ckpt_name, global_step_num, data_list_np)
+        aiturbo.save_ckpt(ckpt_name, global_step_num, data_list_np, crc_check)
     elif async_save:
         data_copy = copy.deepcopy(data_list)
-        thr = Thread(target=_exec_save, args=(ckpt_file_name, data_copy, enc_key, enc_mode, map_param_inc, crc_check),
+        thr = Thread(target=_exec_save,
+                     args=(ckpt_file_name, data_copy, enc_key, enc_mode, map_param_inc, crc_check, format),
                      name="asyn_save_ckpt")
         thr.start()
     else:
-        _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check)
+        _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check, format)
 
     logger.info("Saving checkpoint process is finished.")
 
@@ -692,11 +808,14 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
             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))
 
             # 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:
-                param_data = Tensor(value.data)
+                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)
 
@@ -812,7 +931,7 @@ def load(file_name, **kwargs):
     if not os.path.exists(file_name):
         raise ValueError("For 'load', the argument 'file_name'(MindIR file) does not exist, "
                          "please check whether the 'file_name' is correct.")
-    file_name = os.path.abspath(file_name)
+    file_name = os.path.realpath(file_name)
 
     # set customized functions for dynamic obfuscation
     obfuscated = _check_load_obfuscate(**kwargs)
@@ -875,7 +994,7 @@ def export_split_mindir(file_name, device_num=8, rank_id=0, dynamic=True, sapp=T
     if not os.path.exists(file_name):
         raise ValueError("For 'Split MindIR', the argument 'file_name'(MindIR file) does not exist, "
                          "please check whether the 'file_name' is correct.")
-    file_name = os.path.abspath(file_name)
+    file_name = os.path.realpath(file_name)
 
     logger.info("Execute the process of export and split mindir.")
     dynamic = True
@@ -1074,9 +1193,14 @@ def obfuscate_model(obf_config, **kwargs):
 
 
 def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter_prefix, choice_func, dec_key,
-                          dec_mode, crc_check):
+                          dec_mode, crc_check, format):
     """load parameter into parameter_dict"""
-    ckpt_file_name = _check_ckpt_file_name(ckpt_file_name)
+    ckpt_file_name = _check_ckpt_file_name(ckpt_file_name, format)
+    if format == "safetensors":
+        with safe_open(ckpt_file_name, framework='np') as f:
+            for k in f.keys():
+                parameter_dict[k] = Parameter(f.get_tensor(k))
+        return
     checkpoint_list = _parse_ckpt_proto(ckpt_file_name, dec_key, dec_mode, crc_check)
     try:
         param_data_list = []
@@ -1138,7 +1262,7 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
 
 def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=None,
                     dec_key=None, dec_mode="AES-GCM", specify_prefix=None, choice_func=None,
-                    crc_check=False):
+                    crc_check=False, remove_redundancy=False, format="ckpt"):
     """
     Load checkpoint info from a specified file.
 
@@ -1148,6 +1272,7 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
         - `specify_prefix` and `filter_prefix` are in the process of being deprecated,
           `choice_func` is recommended instead.
           And using either of those two args will override `choice_func` at the same time.
+        - When loading a checkpoint that has removed redundancy, the network should be compiled.
 
     Args:
         ckpt_file_name (str): Checkpoint file name.
@@ -1170,6 +1295,10 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
             that matches the custom condition will be loaded. If returns ``False`` , the Parameter that
             matches the custom condition will be removed. Default: ``None`` .
         crc_check (bool) : Whether to perform crc32 validation when loading checkpoint. Default: ``False`` .
+        remove_redundancy (bool): Whether to enable loading of checkpoint saved with redundancy removal.
+            Redundancy removal refers to eliminating redundant data in data parallelism mode. Default: ``False`` , means
+            redundant-free loading is not enabled.
+        format (str): Format of the input file, can be "ckpt" or "safetensors". Default: "ckpt".
 
     Returns:
         Dict, key is parameter name, value is a Parameter or string. When the `append_dict` parameter of
@@ -1219,24 +1348,35 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
     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_format_and_other_params(format, dec_key, dec_mode, crc_check)
     logger.info("Execute the process of loading checkpoint files.")
 
     parameter_dict = {}
 
     if os.getenv("AITURBO") == "1":
         rank_id = get_rank()
-        import aiturbo
+        from aiturbo.checkpoint import aiturbo_mindspore as aiturbo
         ckpt_path = os.path.dirname(ckpt_file_name)
         ckpt_name = os.path.basename(ckpt_file_name)
-        np_dict = aiturbo.load_ckpt(ckpt_path, ckpt_name, rank_id)
+        np_dict = aiturbo.load_ckpt(ckpt_path, ckpt_name, rank_id, crc_check)
         for key, value in np_dict.items():
+            if crc_check and len(value) != 2:
+                raise ValueError(f"When loading a checkpoint from AITurbo, if CRC check is enabled, "
+                                 f"the length of the value must be 2, but got {len(value)}.")
             if isinstance(value, str):
-                parameter_dict[key] = value
+                if crc_check and value[1] != binascii.crc32(np.array(value[0]).tobytes()):
+                    raise ValueError(f"When loading a checkpoint from AITurbo, the value of the string has not "
+                                     f"passed the CRC check and has been corrupted.")
+                parameter_dict[key] = value[0]
             else:
-                parameter_dict[key] = Parameter(Tensor(value), name=key)
+                if crc_check and value[1] != binascii.crc32(value[0].tobytes()):
+                    raise ValueError(f"When loading a checkpoint from AITurbo, the value of the parameter has not "
+                                     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)
+                              dec_mode, crc_check, format)
 
     if not parameter_dict:
         raise ValueError(f"The loaded parameter dict is empty after filter or specify, please check whether "
@@ -1245,7 +1385,7 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
     if _warm_up_host_cache_enabled(parameter_dict):
         (is_worker, net_dict, warm_up_dict) = _warm_up_host_cache(parameter_dict, net)
     if net is not None:
-        load_param_into_net(net, parameter_dict, strict_load)
+        load_param_into_net(net, parameter_dict, strict_load, remove_redundancy)
     if _warm_up_host_cache_enabled(parameter_dict):
         _warm_up_host_cache_post_process(is_worker, net_dict, warm_up_dict)
 
@@ -1362,17 +1502,20 @@ def _load_map_parameter(checkpoint_list, element, element_id, map_data_list,
         parameter_dict[element.tag] = map_array
 
 
-def _check_ckpt_file_name(ckpt_file_name):
+def _check_ckpt_file_name(ckpt_file_name, format):
     """Check function load_checkpoint's ckpt_file_name."""
     if not isinstance(ckpt_file_name, str):
         raise TypeError("For 'load_checkpoint', the argument 'ckpt_file_name' must be string, "
                         "but got {}.".format(type(ckpt_file_name)))
 
-    if ckpt_file_name[-5:] != ".ckpt":
-        raise ValueError("For 'load_checkpoint', the checkpoint file should end with '.ckpt', please "
+    if format not in ['ckpt', 'safetensors']:
+        raise ValueError("For 'load_checkpoint', the checkpoint file should end with '.ckpt' or '.safetensors', please "
                          "input the correct 'ckpt_file_name'.")
+    if not ckpt_file_name.endswith(format):
+        raise ValueError(f"For 'load_checkpoint', the checkpoint file format must same with 'format', but got "
+                         f"file_name:'{ckpt_file_name}', format:'{format}'")
 
-    ckpt_file_name = os.path.abspath(ckpt_file_name)
+    ckpt_file_name = os.path.realpath(ckpt_file_name)
     if not os.path.exists(ckpt_file_name):
         raise ValueError("For 'load_checkpoint', the checkpoint file: {} does not exist, please check "
                          "whether the 'ckpt_file_name' is correct.".format(ckpt_file_name))
@@ -1414,7 +1557,7 @@ def _parse_ckpt_proto(ckpt_file_name, dec_key, dec_mode, crc_check):
             pb_content = _decrypt(ckpt_file_name, dec_key, len(dec_key), dec_mode)
             if pb_content is None:
                 raise ValueError("For 'load_checkpoint', failed to decrypt the checkpoint file.")
-        if crc_check and pb_content[-17:-10] == b"crc_num":
+        if crc_check and pb_content[-17:-10] != b"crc_num":
             logger.warning("For 'load_checkpoint', the ckpt file do not contain the crc code, please check the file.")
         if pb_content[-17:-10] == b"crc_num":
             crc_num_bytes = pb_content[-10:]
@@ -1426,6 +1569,9 @@ def _parse_ckpt_proto(ckpt_file_name, dec_key, dec_mode, crc_check):
                     raise ValueError("For 'load_checkpoint', the crc check is failed, "
                                      "please check whether the ckpt file is damaged.")
         checkpoint_list.ParseFromString(pb_content)
+    except google.protobuf.message.DecodeError as e:
+        raise ValueError(f"Failed to read the checkpoint file {ckpt_file_name}. "
+                         f"The file may be corrupted, and the content cannot be parsed.") from e
     except BaseException as e:
         if _is_cipher_file(ckpt_file_name):
             err_info = "Failed to read the checkpoint file {}. The file may be encrypted or tempered with, " \
@@ -1455,19 +1601,6 @@ def _whether_load_param(specify_prefix, filter_prefix, param_name):
     return whether_load
 
 
-def _init_parameter_data_in_parallel_mode(net, parameter_dict):
-    """In parallel mode, only init the paraemters in ckpt."""
-    is_train_phase = net.phase.startswith('train')
-    for _, param in net.parameters_and_names():
-        if param.name in parameter_dict and param.from_ckpt and not is_train_phase:
-            param.shape = tuple(parameter_dict[param.name].shape)
-            continue
-        if param.name in parameter_dict and param.has_init:
-            logger.warning("{} is not init while load ckpt.".format(param.name))
-            new_tensor = param.init_data()
-            param._update_tensor_data(new_tensor)
-
-
 def _check_load_param_into_net(net, parameter_dict):
     """check load_param_into_net"""
     if not isinstance(net, nn.Cell):
@@ -1484,10 +1617,13 @@ def _check_load_param_into_net(net, parameter_dict):
         parameter_dict.pop("random_op")
 
 
-def load_param_into_net(net, parameter_dict, strict_load=False):
+def load_param_into_net(net, parameter_dict, strict_load=False, remove_redundancy=False):
     """
     Load parameters into network, return parameter list that are not loaded in the network.
 
+    Note:
+        - When loading a parameter dict that has removed redundancy, the network should be compiled.
+
     Args:
         net (Cell): The network where the parameters will be loaded.
         parameter_dict (dict): The dictionary generated by load checkpoint file,
@@ -1496,6 +1632,9 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
                             into net when parameter name's suffix in checkpoint file is the same as the
                             parameter in the network. When the types are inconsistent perform type conversion
                             on the parameters of the same type, such as float32 to float16. Default: ``False`` .
+        remove_redundancy (bool): Whether to enable loading of checkpoint saved with redundancy removal.
+            Redundancy removal refers to eliminating redundant data in data parallelism mode. Default: ``False`` , means
+            redundant-free loading is not enabled.
 
     Returns:
         - param_not_load (List), the parameter name in model which are not loaded into the network.
@@ -1529,13 +1668,10 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
             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.")
     for _, param in net.parameters_and_names():
         param.from_ckpt = True
-    if not _is_in_auto_parallel_mode():
-        net.init_parameters_data()
-    else:
-        _init_parameter_data_in_parallel_mode(net, parameter_dict)
     param_not_load = []
     ckpt_not_load = list(parameter_dict.keys())
     for _, param in net.parameters_and_names():
@@ -1548,6 +1684,8 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
                 continue
             new_param = parameter_dict[param.name]
             _update_param(param, new_param, strict_load)
+            if hasattr(param, "init_param") and not param.init_param:
+                param.init_param = True
             ckpt_not_load.remove(param.name)
         else:
             param_not_load.append(param.name)
@@ -1560,16 +1698,26 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
         logger.warning("For 'load_param_into_net', "
                        "{} parameters in the 'net' are not loaded, because they are not in the "
                        "'parameter_dict', please check whether the network structure is consistent "
-                       "when training and loading checkpoint.".format(len(param_not_load)))
+                       "when training and loading checkpoint. Another possibility is that "
+                       "the redundant loading is not enabled, but the loaded checkpoint is saved with "
+                       "redundancy removed. ".format(len(param_not_load)))
         logger.warning("{} are not loaded.".format(param_not_load))
-    if os.getenv("AITURBO") == "1" and net.parameter_layout_dict is not None:
+    if remove_redundancy:
+        parallel_mode = context.get_auto_parallel_context("parallel_mode")
+        if parallel_mode == "stand_alone":
+            raise TypeError(f"The deduplication feature for loading checkpoint can only be used "
+                            f"in parallel scenarios, but got {parallel_mode}.")
+        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.")
         param_layout = net.parameter_layout_dict
-        param_redundancy = get_parameter_redundancy(param_layout)
-        remove_param_redundancy_dict = remove_param_redundancy(param_redundancy)
-        target_parameter_name_set = set(parameter_dict.keys())
-        for rank_id, param_name_set in remove_param_redundancy_dict:
-            if param_name_set == target_parameter_name_set:
-                parameter_broadcast(net, param_layout, rank_id)
+        rank_id = get_rank()
+        device_num = _get_device_num()
+        stage_num = _get_auto_parallel_context("pipeline_stages")
+        chunk_size = device_num // stage_num
+        initial_rank = (rank_id // chunk_size) * chunk_size
+        _single_parameter_broadcast(net, param_layout, rank_id, initial_rank)
+
     return param_not_load, ckpt_not_load
 
 
@@ -1662,6 +1810,8 @@ def _load_dismatch_prefix_params(net, parameter_dict, param_not_load, strict_loa
                 if param.name in param_not_load and new_param_name in parameter_dict:
                     new_param = parameter_dict[new_param_name]
                     _update_param(param, new_param, strict_load)
+                    if hasattr(param, "init_param") and not param.init_param:
+                        param.init_param = True
                     param_not_load.remove(param.name)
 
 
@@ -1675,7 +1825,7 @@ def _save_graph(network, file_name):
     """
     logger.info("Execute the process of saving graph.")
 
-    file_name = os.path.abspath(file_name)
+    file_name = os.path.realpath(file_name)
     graph_pb = network.get_func_graph_proto()
     if graph_pb:
         with open(file_name, "wb") as f:
@@ -1790,7 +1940,7 @@ def export(net, *inputs, file_name, file_format, **kwargs):
             - AIR: Ascend Intermediate Representation. An intermediate representation format of Ascend model.
             - ONNX: Open Neural Network eXchange. An open format built to represent machine learning models.
             - MINDIR: MindSpore Native Intermediate Representation for Anf. An intermediate representation format
-              for MindSpore models.
+              for MindSpore models. MINDIR does not support operators which have dictionary attribute.
 
         kwargs (dict): Configuration options dictionary.
 
@@ -1889,7 +2039,7 @@ def export(net, *inputs, file_name, file_format, **kwargs):
                                + str(columns))
         inputs = tuple(inputs_col)
 
-    file_name = os.path.abspath(file_name)
+    file_name = os.path.realpath(file_name)
     if 'enc_key' in kwargs.keys():
         kwargs['enc_key'], kwargs['enc_mode'] = _check_key_mode_type(file_format, **kwargs)
     _export(net, file_name, file_format, *inputs, **kwargs)
@@ -1982,8 +2132,8 @@ def _save_air(net, file_name, *inputs, **kwargs):
     if os.path.exists(file_name):
         os.chmod(file_name, stat.S_IWUSR)
     if "/" in file_name:
-        real_path = os.path.abspath(file_name[:file_name.rfind("/")])
-        os.makedirs(real_path, exist_ok=True)
+        real_path = os.path.realpath(file_name[:file_name.rfind("/")])
+        os.makedirs(real_path, mode=0o700, exist_ok=True)
     if 'enc_key' in kwargs.keys() and 'enc_mode' in kwargs.keys():
         _executor.export(file_name, graph_id, enc_key=kwargs.get('enc_key'), encrypt_func=kwargs.get('enc_mode'))
     else:
@@ -2093,12 +2243,12 @@ def _split_save(net_dict, model, file_name, is_encrypt, **kwargs):
     file_prefix = file_name.split("/")[-1]
     if file_prefix.endswith(".mindir"):
         file_prefix = file_prefix[:-7]
-    current_path = os.path.abspath(file_name)
+    current_path = os.path.realpath(file_name)
     dirname = os.path.dirname(current_path)
     data_path = os.path.join(dirname, file_prefix + "_variables")
     if os.path.exists(data_path):
         shutil.rmtree(data_path)
-    os.makedirs(data_path, exist_ok=True)
+    os.makedirs(data_path, mode=0o700, exist_ok=True)
     os.chmod(data_path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
     index = 0
     external_local = os.path.join(file_prefix + "_variables", "data_" + str(index))
@@ -2267,9 +2417,9 @@ def _save_mindir_together(net_dict, model, file_name, is_encrypt, **kwargs):
                              "the data of parameter cannot be exported.".format(map_param_proto.name))
     if not file_name.endswith('.mindir'):
         file_name += ".mindir"
-    current_path = os.path.abspath(file_name)
+    current_path = os.path.realpath(file_name)
     dirname = os.path.dirname(current_path)
-    os.makedirs(dirname, exist_ok=True)
+    os.makedirs(dirname, mode=0o700, exist_ok=True)
     if os.path.exists(file_name):
         os.chmod(file_name, stat.S_IWUSR)
     with open(file_name, 'wb') as f:
@@ -2398,7 +2548,7 @@ def parse_print(print_file_name):
         [[ 1.00000000e+00,  2.00000000e+00,  3.00000000e+00,  4.00000000e+00],
         [ 5.00000000e+00,  6.00000000e+00,  7.00000000e+00,  8.00000000e+00]])]
     """
-    print_file_path = os.path.abspath(print_file_name)
+    print_file_path = os.path.realpath(print_file_name)
 
     if os.path.getsize(print_file_path) == 0:
         raise ValueError("For 'parse_print', the print file may be empty, please make sure enter the correct "
@@ -2687,14 +2837,15 @@ def merge_sliced_parameter(sliced_parameters, strategy=None):
     return merged_parameter
 
 
-def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=None,
-                                train_strategy_filename=None, strict_load=False, dec_key=None, dec_mode='AES-GCM'):
+def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_strategy=None,
+                                train_strategy_filename=None, strict_load=False, dec_key=None, dec_mode='AES-GCM',
+                                format='ckpt', unified_safetensors_dir=None, dst_safetensors_dir=None, rank_id=None):
     """
     Load checkpoint into net for distributed predication. Used in the case of distributed inference.
 
     Args:
         network (Cell): Network for distributed predication.
-        checkpoint_filenames (list[str]): The name of Checkpoint files in order of rank id.
+        checkpoint_filenames (list[str]): The name of Checkpoint files in order of rank id. Default: ``None`` .
         predict_strategy (dict): Strategy of predication process. It means that using one device to predict
                                  when setting predict_strategy as None. Default: ``None`` .
         train_strategy_filename (str): The filename of training strategy protocol buffer file.
@@ -2711,6 +2862,14 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
         dec_mode (str): This parameter is valid only when dec_key is not set to ``None`` . Specifies the decryption
                         mode, currently supports ``'AES-GCM'`` , ``'AES-CBC'``  and ``'SM4-CBC'`` .
                         Default: ``'AES-GCM'`` .
+        format (str): Input weight format to be loaded into the network.
+                      It can be set to either "ckpt" or "safetensors". Default: "ckpt".
+        unified_safetensors_dir (str): Directory of input weight files to be loaded into the network.
+                                       Default: ``None`` .
+        dst_safetensors_dir (str): In the save mode scenario, the save directory for safetensors.
+        rank_id (int): The logical sequence number of the card. In non save mode, it is automatically obtained
+                       globally by initializing the network; In save mode, save the file according to the input
+                       sequence number. If it is not input, save the entire file.
 
     Raises:
         TypeError: The type of inputs do not match the requirements.
@@ -2725,14 +2884,14 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `rank table startup
-            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/rank_table.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/rank_table.html>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun startup
-            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/mpirun.html>`_ .
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/mpirun.html>`_ .
 
             For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/master/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/docs/en/master/model_train/parallel/dynamic_cluster.html>`_ .
 
         >>> import os
         >>> import numpy as np
@@ -2814,6 +2973,54 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
         ...
         [ 1.6067538  1.6244187  1.5384722 ...  1.5449994  1.6195512  1.6176052]]
     """
+    if format not in ['safetensors', 'ckpt']:
+        raise ValueError(
+            f"For 'load_distributed_checkpoint', 'format' must be 'ckpt' or 'safetensors', but got {format}.")
+
+    if format == 'safetensors':
+        if unified_safetensors_dir is None:
+            raise ValueError(f"For 'load_distributed_checkpoint', 'unified_safetensors_dir' can not be None "
+                             f"when format is 'safetensors'.")
+        unsupport_param = [checkpoint_filenames, train_strategy_filename, dec_key]
+        for param in unsupport_param:
+            if param is not None:
+                raise ValueError(f"For 'load_distributed_checkpoint', {param} must be None "
+                                 f"when format is 'safetensors'.")
+        if strict_load or dec_mode != 'AES-GCM':
+            raise ValueError(f"For 'load_distributed_checkpoint', strict_load and dec_mode must be default "
+                             f"when format is 'safetensors'.")
+        if network is not None:
+            rank_id = get_rank()
+            _load_parallel_checkpoint(unified_safetensors_dir, predict_strategy, network, rank_id=rank_id)
+        else:
+            if dst_safetensors_dir is None:
+                raise ValueError(f"For 'load_distributed_checkpoint', 'dst_safetensors_dir' can not be None "
+                                 f"when network is None.")
+            if rank_id is not None:
+                _load_parallel_checkpoint(unified_safetensors_dir, predict_strategy, network, dst_safetensors_dir,
+                                          rank_id)
+            else:
+                dst_strategy_dict = _build_searched_strategy(predict_strategy)
+                dst_stage_device_num = _get_device_num_from_strategy(dst_strategy_dict)
+                dst_stage_num = _extract_pipeline_stage_num(dst_strategy_dict)
+                dst_device_num = dst_stage_device_num * dst_stage_num
+                processes = []
+                activate_processes = 0
+                for rank in range(0, dst_device_num):
+                    p = Process(target=_load_parallel_checkpoint, args=(
+                        unified_safetensors_dir, predict_strategy, network, dst_safetensors_dir, rank))
+                    p.start()
+                    processes.append(p)
+                    activate_processes += 1
+                    max_processes = 64
+                    if activate_processes >= max_processes:
+                        p = processes.pop(0)
+                        p.join()
+                        activate_processes -= 1
+                for p in processes:
+                    p.join()
+        return
+
     network = Validator.check_isinstance("network", network, nn.Cell)
     _check_checkpoint_file(checkpoint_filenames)
     _check_predict_strategy(predict_strategy)
@@ -2858,17 +3065,24 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
         param_rank = rank_list.get(param.name)[0]
         skip_merge_split = rank_list.get(param.name)[1]
         shard_stride = train_strategy.get(param.name)[4]
+        tensor_map = train_strategy.get(param.name)[1]
+        first_dim_shard_idx = tensor_map[0] if tensor_map else -1
+        device_arrangement = train_strategy.get(param.name)[0]
+        first_dim_shard_size = 1
+        if first_dim_shard_idx >= 0:
+            first_dim_shard_size = device_arrangement[-1 - first_dim_shard_idx]
         if train_strategy.get(param.name)[5]:
-            shard_size = ckpt_file_len / shard_stride / train_strategy.get(param.name)[5]
+            shard_size = int(ckpt_file_len / shard_stride / train_strategy.get(param.name)[5] / first_dim_shard_size)
         else:
             shard_size = 0
         for rank in param_rank:
             param_total_list = list(range(0, ckpt_file_len))
+            if first_dim_shard_size != 1:
+                param_total_list = _get_param_list_when_first_dim_sharded(device_arrangement, first_dim_shard_idx, rank)
             if shard_size > 0:
-                shard_total_list = []
-                for i in range(0, ckpt_file_len, shard_size):
-                    shard_total_list.append(param_total_list[i:i + shard_size])
-                param_total_list = shard_total_list[rank // shard_size]
+                rank_index = param_total_list.index(rank)
+                start = rank_index // shard_size * shard_size
+                param_total_list = param_total_list[start:start + shard_size]
             if shard_stride > 0:
                 param_stride = []
                 # merge pre parameter
@@ -3040,7 +3254,7 @@ def _get_mindir_inputs(file_name):
         >>> input_tensor = get_mindir_inputs("lenet.mindir")
     """
     Validator.check_file_name_by_regular(file_name)
-    file_name = os.path.abspath(file_name)
+    file_name = os.path.realpath(file_name)
     model = read_proto(file_name)
     input_tensor = []