mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/train/serialization.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2021 Huawei Technologies Co., Ltd
+# Copyright 2020-2024 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,19 +17,23 @@
 from __future__ import absolute_import
 from __future__ import division
 
+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
 
 import math
 import sys
 import time
+import google
 import numpy as np
 
 from mindspore.train.checkpoint_pb2 import Checkpoint
@@ -50,32 +54,41 @@ from mindspore.common.api import _generate_branch_control_input
 from mindspore.common.initializer import initializer, One
 from mindspore.common.parameter import Parameter, _offload_if_config
 from mindspore.common.tensor import Tensor
+from mindspore._c_expression import Tensor as Tensor_
 from mindspore.common._utils import is_shape_unknown
+from mindspore.common.file_system import FileSystem, _register_basic_file_system, _register_mindio_file_system
 from mindspore.communication.management import get_rank, get_group_size
 from mindspore.experimental import MapParameter
-from mindspore.parallel._cell_wrapper import get_allgather_cell
+from mindspore.ops import Cast
+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, _is_parallel_mode
+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.train._utils import read_proto
+from mindspore.parallel.checkpoint_transform import sync_pipeline_shared_parameters
+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 safetensors.numpy import save_file
+from safetensors import safe_open
 from ..ops.operations._opaque_predicate_registry import add_opaque_predicate, clean_funcs
-from ..ops.operations import Cast
 
 tensor_to_ms_type = {"Int8": mstype.int8, "UInt8": mstype.uint8, "Int16": mstype.int16, "UInt16": mstype.uint16,
                      "Int32": mstype.int32, "UInt32": mstype.uint32, "Int64": mstype.int64, "UInt64": mstype.uint64,
                      "Float16": mstype.float16, "Float32": mstype.float32, "Float64": mstype.float64,
-                     "Bool": mstype.bool_, "str": mstype.string, "BFloat16": mstype.bfloat16}
+                     "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",
-                     "BFloat16": np.float32}
+                     "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}
 
@@ -95,6 +108,92 @@ INT_64_MAX = 9223372036854775807
 
 cpu_cast = Cast().set_device("CPU")
 
+_ckpt_fs = FileSystem()
+
+
+def init_ckpt_file_system(fs: FileSystem):
+    """Initialize checkpoint file system"""
+    if _register_mindio_file_system(fs):
+        return
+    _register_basic_file_system(fs)
+
+
+# Initialize checkpoint file system
+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
+        self.param_dict_future = param_dict_future
+
+    def result(self):
+        param_dict = self.param_dict_future.result()
+        self.executor.shutdown()
+        return param_dict
+
 
 def _special_process_par(par, new_par):
     """
@@ -221,53 +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):
+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 open(ckpt_file_name, "ab") as f:
-                plain_data = None
-                if enc_key is not None:
-                    plain_data = BytesIO()
-
-                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:]
-                        if value[3].dtype == mstype.bfloat16:
-                            new_value[2] = cpu_cast(value[3], mstype.float32).asnumpy().reshape(-1)
-                        else:
-                            new_value[2] = value[3].asnumpy().reshape(-1)
-                        _write_parameter_data(name, new_value, f, enc_key, plain_data)
-                        _offload_if_config(value[3])
-                        continue
-                    if value[0] == "BFloat16_tensor":
-                        _write_bfloat16_data(name, value, f, enc_key, plain_data)
-                        continue
-                    if isinstance(value[2], Tensor):
-                        _write_hugeparameter(name, value, f)
-                        continue
-
-                    _write_parameter_data(name, value, f, enc_key, plain_data)
-
-                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)
-
-                os.chmod(ckpt_file_name, stat.S_IRUSR)
-
+                        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)
@@ -286,22 +404,7 @@ def _write_random_seed(name, value, f):
     f.write(checkpoint_list.SerializeToString())
 
 
-def _write_bfloat16_data(name, value, f, enc_key, plain_data):
-    """Write bfloat16 data 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(value[1])
-    param_tensor.tensor_type = value[2]
-    param_tensor.tensor_content = value[3].get_bytes()
-    if enc_key is None:
-        f.write(checkpoint_list.SerializeToString())
-    else:
-        plain_data.write(checkpoint_list.SerializeToString())
-
-
-def _write_parameter_data(name, value, f, enc_key, plain_data):
+def _write_parameter_data(name, value, f, enc_key, plain_data, crc_num=0, crc_check=False):
     """Write parameter data into protobuf file."""
     data_size = value[2].nbytes / 1024
     if data_size > SLICE_SIZE:
@@ -320,10 +423,40 @@ def _write_parameter_data(name, value, f, enc_key, plain_data):
         param_tensor.tensor_content = param_slice.tobytes()
 
         if enc_key is None:
-            f.write(checkpoint_list.SerializeToString())
+            output_data = checkpoint_list.SerializeToString()
+            if crc_check:
+                crc_num = binascii.crc32(output_data, crc_num)
+            f.write(output_data)
+        else:
+            plain_data.write(checkpoint_list.SerializeToString())
+
+    return crc_num
+
+
+def _write_parameter_bytes_data(name, value, f, enc_key, plain_data, crc_num=0, crc_check=False):
+    """Write parameter bytes data into protobuf file."""
+    bytes_value = value[2].get_bytes()
+    chunk_size = 1024 * SLICE_SIZE
+
+    for i in range(0, len(bytes_value), chunk_size):
+        checkpoint_list = Checkpoint()
+        param_value = checkpoint_list.value.add()
+        param_value.tag = name
+        param_tensor = param_value.tensor
+        param_tensor.dims.extend(value[0])
+        param_tensor.tensor_type = value[1]
+        param_tensor.tensor_content = bytes_value[i:i + chunk_size]
+
+        if enc_key is None:
+            output_data = checkpoint_list.SerializeToString()
+            if crc_check:
+                crc_num = binascii.crc32(output_data, crc_num)
+            f.write(output_data)
         else:
             plain_data.write(checkpoint_list.SerializeToString())
 
+    return crc_num
+
 
 def _write_mapparameter(name, value, f, map_param_inc=False):
     """Write map parameter into protobuf file."""
@@ -365,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)))
@@ -374,20 +510,32 @@ 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, **kwargs):
+                    async_save=False, append_dict=None, enc_key=None, enc_mode="AES-GCM", choice_func=None,
+                    crc_check=False, format="ckpt", **kwargs):
     r"""
     Save checkpoint to a specified file.
 
+    Note:
+        The `enc_mode` and `crc_check` parameters are mutually exclusive and cannot be configured simultaneously.
+
     Args:
         save_obj (Union[Cell, list, dict]): The object to be saved. The data type can be :class:`mindspore.nn.Cell`,
             list, or dict. If a list, it can be the returned value of `Cell.trainable_params()`, or a list of dict
@@ -409,6 +557,9 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                                  If returns ``True`` , the Parameter that matching the custom condition will be saved.
                                  If returns ``False`` , the Parameter that not matching the custom condition will not
                                  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:
@@ -420,7 +571,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
         >>> import mindspore as ms
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> ms.save_checkpoint(net, "./lenet.ckpt",
         ...                    choice_func=lambda x: x.startswith("conv") and not x.startswith("conv1"))
@@ -440,35 +591,57 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
 
     Tutorial Examples:
         - `Saving and Loading the Model - Saving and Loading the Model Weight
-          <https://mindspore.cn/tutorials/en/r2.2/beginner/save_load.html#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)
     enc_key = Validator.check_isinstance('enc_key', enc_key, (type(None), bytes))
     enc_mode = Validator.check_isinstance('enc_mode', enc_mode, str)
+    crc_check = Validator.check_isinstance('crc_check', crc_check, bool)
     map_param_inc = kwargs.get('incremental', False)
     logger.info("Execute the process of saving checkpoint files.")
-
-    save_obj = _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choice_func)
+    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)
+
+    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 not isinstance(value, str):
+            if isinstance(value, Generator):
+                value = value.get_state()
+            elif not isinstance(value, str):
                 value = Tensor(value)
             append_info_list.append({"name": k_name, "data": value})
         save_obj.extend(append_info_list)
 
     data_list = OrderedDict()
+    data_list_np = OrderedDict()
     with _ckpt_mutex:
         for param in save_obj:
             if param["name"] == "random_op":
-                data_list["random_op"] = param["data"]
+                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] = []
             if isinstance(param["data"], MapParameter):
                 data_list[param["name"]].append("mapparameter")
                 data_list[param["name"]].append(param["data"])
@@ -479,49 +652,48 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                 elif param["data"][0] == "offload_parameter":
                     data_list[key].append("offload_parameter")
                     _save_param_list_data(data_list, key, param)
-                elif param["data"][0] == "BFloat16_tensor":
-                    data_list[key].append("BFloat16_tensor")
-                    _save_param_list_data(data_list, key, param)
-                    continue
 
             if isinstance(param["data"], str):
-                data_list[key].append([0])
-                data_list[key].append('str')
-                data = np.array(param["data"])
-                data_list[key].append(data)
+                if os.getenv("AITURBO") == "1":
+                    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')
+                    data = np.array(param["data"])
+                    data_list[key].append(data)
             else:
                 if isinstance(param["data"], Parameter):
                     param["data"].init_data()
-                if isinstance(param["data"], Tensor) and param["data"].dtype == mstype.bfloat16:
-                    data_list[key].append("BFloat16_tensor")
-                    dims = []
-                    for dim in param["data"].shape:
-                        dims.append(dim)
-                    data_list[key].append(dims)
-                    data_list[key].append("BFloat16")
-                    data_list[key].append(cpu_cast(param["data"], mstype.float32))
-                    continue
-                dims = []
-                if param['data'].shape == ():
-                    dims.append(0)
+                if os.getenv("AITURBO") == "1":
+                    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:
                         dims.append(dim)
-                data_list[key].append(dims)
-                tensor_type = str(param["data"].dtype)
-                data_list[key].append(tensor_type)
-                if param["data"].dtype == mstype.bfloat16:
-                    data = cpu_cast(param["data"], mstype.float32).asnumpy().reshape(-1)
-                else:
-                    data = param["data"].asnumpy().reshape(-1)
-                data_list[key].append(data)
-
-    if async_save:
+                    data_list[key].append(dims)
+                    tensor_type = str(param["data"].dtype)
+                    data_list[key].append(tensor_type)
+                    data = param["data"]
+                    data_list[key].append(data)
+
+    if os.getenv("AITURBO") == "1":
+        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, 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), name="asyn_save_ckpt")
+        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)
+        _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc, crc_check, format)
 
     logger.info("Saving checkpoint process is finished.")
 
@@ -532,7 +704,21 @@ def _convert_list_to_param_list(save_obj, choice_func):
     if not save_obj:
         return param_list
     if isinstance(save_obj[0], dict):
-        param_list = [param for param in save_obj if choice_func is None or choice_func(param["name"])]
+        for param in save_obj:
+            if isinstance(param, dict) and "name" in param and "data" in param:
+                if not isinstance(param["name"], str):
+                    raise TypeError(f"For save_checkpoint, when save_obj is a list of dict items, the name in dict "
+                                    f"should be string, but got {type(param['name'])}.")
+                if not isinstance(param["data"], Tensor):
+                    raise TypeError(f"For save_checkpoint, when save_obj is a list of dict items, the data in dict "
+                                    f"should be parameter, but got {type(param['data'])}.")
+                if choice_func is not None and not choice_func(param["name"]):
+                    continue
+                each_param = {"name": param["name"], "data": param["data"]}
+                param_list.append(each_param)
+            else:
+                raise TypeError(f"For save_checkpoint, save_obj should be a list of dict items, and the dict should "
+                                f"have key values 'name' and 'value', but got {type(param)} and {param}.")
     else:
         for param in save_obj:
             if isinstance(param, Parameter):
@@ -585,6 +771,7 @@ def _convert_cell_param_and_names_to_dict(save_obj, choice_func):
 
 def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_func):
     """Convert nn.Cell to param_list."""
+    sync_pipeline_shared_parameters(save_obj)
     param_list = []
     parameter_layout_dict = save_obj.parameter_layout_dict
     if _is_in_auto_parallel_mode() and not parameter_layout_dict:
@@ -597,7 +784,7 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
         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")
+            append_dict.pop("random_op")
     for (key, value) in param_dict.items():
         each_param = {"name": key}
         if isinstance(value, MapParameter):
@@ -619,18 +806,16 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
             param_data.append(param_tensor.shape)
             param_data.append(str(param_tensor.dtype))
             param_data.append(value.key)
-        elif value.data.dtype == mstype.bfloat16:
-            param_data = ["BFloat16_tensor"]
-            param_data.append(cpu_cast(value.data, mstype.float32))
-            param_data.append(value.data.shape)
-            param_data.append("BFloat16")
-            param_data.append(value.key)
         else:
-            param_data = Tensor(value.data.asnumpy())
+            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:
+                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)
 
@@ -670,9 +855,9 @@ def _check_append_dict(append_dict):
         raise TypeError("For 'save_checkpoint', the argument 'append_dict' must be dict, but got "
                         "{}.".format(type(append_dict)))
     for key, value in append_dict.items():
-        if not isinstance(key, str) or not isinstance(value, (int, float, bool, str, Parameter, Tensor)):
+        if not isinstance(key, str) or not isinstance(value, (int, float, bool, str, Parameter, Tensor, Generator)):
             raise TypeError(f"For 'save_checkpoint', the type of dict 'append_info' must be key: string, "
-                            f"value: int, float or bool, but got key: {type(key)}, value: {type(value)}")
+                            f"value: int, float, bool or Generator, but got key: {type(key)}, value: {type(value)}")
     return append_dict
 
 
@@ -699,13 +884,13 @@ def load(file_name, **kwargs):
             - dec_key (bytes): Byte-type key used for decryption. The valid length is 16, 24, or 32.
             - dec_mode (Union[str, function]): Specifies the decryption mode, to take effect when dec_key is set.
 
-              - Option: 'AES-GCM', 'AES-CBC', 'SM4-CBC' or customized decryption. Default: 'AES-GCM'.
+              - Option: 'AES-GCM', 'AES-CBC', 'SM4-CBC' or customized decryption. Default: ``'AES-GCM'``.
               - For details of using the customized decryption, please check the `tutorial
-                <https://mindspore.cn/mindarmour/docs/en/r2.0/model_encrypt_protection.html>`_.
+                <https://mindspore.cn/mindarmour/docs/en/master/model_encrypt_protection.html>`_.
 
             - obf_func (function): A python function used for loading obfuscated MindIR model, which can refer to
               `obfuscate_model()
-              <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore/mindspore.obfuscate_model.html>`_.
+              <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.obfuscate_model.html>`_.
 
     Returns:
         GraphCell, a compiled graph that can executed by `GraphCell`.
@@ -735,7 +920,7 @@ def load(file_name, **kwargs):
 
     Tutorial Examples:
         - `Saving and Loading the Model - Saving and Loading MindIR
-          <https://mindspore.cn/tutorials/en/r2.2/beginner/save_load.html#saving-and-loading-mindir>`_
+          <https://mindspore.cn/tutorials/en/master/beginner/save_load.html#saving-and-loading-mindir>`_
     """
     if not isinstance(file_name, str):
         raise ValueError("For 'load', the argument 'file_name' must be string, but "
@@ -746,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)
@@ -776,7 +961,7 @@ def load(file_name, **kwargs):
     return graph
 
 
-def export_split_mindir(file_name, device_num=8, rank_id=0, dynamic=True, sapp=False):
+def export_split_mindir(file_name, device_num=8, rank_id=0, dynamic=True, sapp=True):
     """
     Auto Split MindIR.
 
@@ -784,10 +969,10 @@ def export_split_mindir(file_name, device_num=8, rank_id=0, dynamic=True, sapp=F
 
     Args:
         file_name (str): MindIR file name.
-        device_num (int): device number.
-        rank_id (int): rank id.
-        dynamic (bool): Indicates whether the model is a dynamic shape mindir model.
-        sapp (bool): Indicates whether to automatically generate split strategy through SAPP.
+        device_num (int): device number. Default: '8'.
+        rank_id (int): rank id. Default: '0'.
+        dynamic (bool): Indicates whether the model is a dynamic shape mindir model. Default: 'True'.
+        sapp (bool): Indicates whether to automatically generate split strategy through SAPP. Default: 'True'.
 
     Raises:
         ValueError: MindIR file does not exist or `file_name` is not a string.
@@ -809,7 +994,7 @@ def export_split_mindir(file_name, device_num=8, rank_id=0, dynamic=True, sapp=F
     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
@@ -909,13 +1094,14 @@ def obfuscate_model(obf_config, **kwargs):
             - customized_func (function): A python function used for customized function mode, which used for control
               the switch branch of obfuscation structure. The outputs of customized_func should be boolean and const (
               Reference to 'my_func()' in
-              `tutorials <https://www.mindspore.cn/mindarmour/docs/en/r2.0/dynamic_obfuscation_protection.html>`_).
+              `tutorials <https://www.mindspore.cn/mindarmour/docs/en/master/dynamic_obfuscation_protection.html>`_).
               This function needs to ensure that its result is constant for any input. Users can refer to opaque
               predicates. If customized_func is set, then it should be passed to :func:`mindspore.load` interface
               when loading obfuscated model.
             - obf_random_seed (int): Obfuscation random seed, which should be in (0, 9223372036854775807]. The
               structure of obfuscated models corresponding to different random seeds is different. If
-              `obf_random_seed` is set, then it should be passed to :class:`nn.GraphCell()` interface when loading
+              `obf_random_seed` is set, then it should be passed to :class:`mindspore.nn.GraphCell`
+              interface when loading
               obfuscated model. It should be noted that at least one of `customized_func` or `obf_random_seed` should
               be set, and the latter mode would be applied if both of them are set.
 
@@ -923,7 +1109,7 @@ def obfuscate_model(obf_config, **kwargs):
 
             - enc_key (bytes): Byte type key used for encryption. The valid length is 16, 24, or 32.
             - enc_mode (str): Specifies the encryption mode, to take effect when dec_key is set.
-              Option: 'AES-GCM' | 'AES-CBC' | 'SM4-CBC'. Default: 'AES-GCM'.
+              Options: ``'AES-GCM'`` | ``'AES-CBC'`` | ``'SM4-CBC'``. Default: ``'AES-GCM'``.
 
     Raises:
         TypeError: If `obf_config` is not a dict.
@@ -934,11 +1120,15 @@ def obfuscate_model(obf_config, **kwargs):
         ValueError: If `obf_ratio` is not provided in `obf_config`.
         ValueError: If both `customized_func` and `obf_random_seed` are not provided in `obf_config`.
         ValueError: If `obf_random_seed` is not in (0, 9223372036854775807].
-        ValueError: If `original_model_path` is not exist or `original_model_path` is not end with '.mindir'.
+        ValueError: If `original_model_path` does not exist or `original_model_path` does not end with '.mindir'.
 
     Examples:
         >>> import mindspore as ms
         >>> import mindspore.nn as nn
+        >>> import numpy as np
+        >>> # Download ori_net.mindir
+        >>> # https://gitee.com/mindspore/mindspore/blob/master/tests/ut/python/mindir/ori_net.mindir
+        >>> input1 = ms.Tensor(np.ones((1, 1, 32, 32)).astype(np.float32))
         >>> obf_config = {'original_model_path': "./net.mindir",
         ...          'save_model_path': "./obf_net",
         ...          'model_inputs': [input1, ],
@@ -998,12 +1188,81 @@ def obfuscate_model(obf_config, **kwargs):
     obf_net = nn.GraphCell(obf_graph)
     if obf_random_seed != 0:
         append_y_tensor = Tensor(np.ones((1, 1)).astype(np.int32))
-        model_inputs += [append_y_tensor,]
+        model_inputs += [append_y_tensor]
     export(obf_net, *model_inputs, file_name=saved_path, file_format="MINDIR", **kwargs)
 
 
+def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter_prefix, choice_func, dec_key,
+                          dec_mode, crc_check, format):
+    """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:
+            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 = []
+        map_data_list = [[], [], []]
+        map_shape_list = [0, 0, 0]
+        if specify_prefix:
+            logger.warning("For load_checkpoint, this parameter `specity_prefix` will be deprecated, "
+                           "please use `choice_func` instead.")
+        if filter_prefix:
+            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 \
+                    choice_func is not None and not choice_func(element.tag):
+                continue
+            if element.tensor.ByteSize() == 0:
+                _load_map_parameter(checkpoint_list, element, element_id, map_data_list, map_shape_list,
+                                    parameter_dict)
+                if element.tag in parameter_dict:
+                    map_data_list = [[], [], []]
+                    map_shape_list = [0, 0, 0]
+                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):
+                new_data = b"".join(param_data_list)
+                param_data_list.clear()
+                dims = element.tensor.dims
+                if data_type == 'str':
+                    str_value = np.frombuffer(new_data, np_type)
+                    parameter_dict[element.tag] = str(str_value[0])
+                else:
+                    if dims == [0]:
+                        dims = []
+                    param_data = Tensor_.convert_bytes_to_tensor(new_data, tuple(dims), ms_type)
+                    parameter = Parameter(param_data, name=element.tag)
+                    parameter_dict[element.tag] = parameter
+                    _offload_if_config(parameter)
+
+        logger.info("Loading checkpoint files process is finished.")
+
+    except BaseException as e:
+        logger.critical("Failed to load the checkpoint file '%s'.", ckpt_file_name)
+        raise ValueError(e.__str__() + "\nFor 'load_checkpoint', "
+                                       "failed to load the checkpoint file {}.".format(ckpt_file_name)) from e
+
+
 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):
+                    dec_key=None, dec_mode="AES-GCM", specify_prefix=None, choice_func=None,
+                    crc_check=False, remove_redundancy=False, format="ckpt"):
     """
     Load checkpoint info from a specified file.
 
@@ -1013,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.
@@ -1034,6 +1294,11 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
             and the return value is a bool. If returns ``True`` , the Parameter
             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
@@ -1076,83 +1341,42 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
 
     Tutorial Examples:
         - `Saving and Loading the Model - Saving and Loading the Model Weight
-          <https://mindspore.cn/tutorials/en/r2.2/beginner/save_load.html#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_ckpt_file_name(ckpt_file_name)
     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_format_and_other_params(format, dec_key, dec_mode, crc_check)
     logger.info("Execute the process of loading checkpoint files.")
-    checkpoint_list = _parse_ckpt_proto(ckpt_file_name, dec_key, dec_mode)
 
     parameter_dict = {}
-    try:
-        param_data_list = []
-        map_data_list = [[], [], []]
-        map_shape_list = [0, 0, 0]
-        if specify_prefix:
-            logger.warning("For load_checkpoint, this parameter `specity_prefix` will be deprecated, "
-                           "please use `choice_func` instead.")
-        if filter_prefix:
-            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 \
-                    choice_func is not None and not choice_func(element.tag):
-                continue
-            if element.tensor.ByteSize() == 0:
-                _load_map_parameter(checkpoint_list, element, element_id, map_data_list, map_shape_list, parameter_dict)
-                if element.tag in parameter_dict:
-                    map_data_list = [[], [], []]
-                    map_shape_list = [0, 0, 0]
-                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)
-            if data_type == "BFloat16":
-                dims = element.tensor.dims
-                param_data = np.frombuffer(data, np_type)
-                param_data = param_data.reshape(list(dims))
-                parameter = Parameter(Tensor(param_data, ms_type), name=element.tag)
-                parameter_dict[element.tag] = parameter
-                continue
-            element_data = np.frombuffer(data, np_type)
-            param_data_list.append(element_data)
-            if (element_id == len(checkpoint_list.value) - 1) or \
-                    (element.tag != checkpoint_list.value[element_id + 1].tag):
-                new_data = b"".join(param_data_list)
-                param_data = np.frombuffer(new_data, np_type)
-                param_data_list.clear()
-                dims = element.tensor.dims
-                if dims == [0] and data_type == 'str':
-                    parameter_dict[element.tag] = str(element_data[0])
-                else:
-                    if dims == [0] and 'Float' in data_type:
-                        param_data = float(param_data[0])
-                    if dims == [0] and 'Int' in data_type:
-                        param_data = int(param_data[0])
-                    if dims not in ([0], [1]):
-                        param_data = param_data.reshape(list(dims))
-                    parameter = Parameter(Tensor(param_data, ms_type), name=element.tag)
-                    parameter_dict[element.tag] = parameter
-                    _offload_if_config(parameter)
-
-        logger.info("Loading checkpoint files process is finished.")
 
-    except BaseException as e:
-        logger.critical("Failed to load the checkpoint file '%s'.", ckpt_file_name)
-        raise ValueError(e.__str__() + "\nFor 'load_checkpoint', "
-                                       "failed to load the checkpoint file {}.".format(ckpt_file_name)) from e
+    if os.getenv("AITURBO") == "1":
+        rank_id = get_rank()
+        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, 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):
+                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:
+                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, format)
 
     if not parameter_dict:
         raise ValueError(f"The loaded parameter dict is empty after filter or specify, please check whether "
@@ -1161,13 +1385,93 @@ 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)
 
     return parameter_dict
 
 
+def load_checkpoint_async(ckpt_file_name, net=None, strict_load=False, filter_prefix=None, dec_key=None,
+                          dec_mode="AES-GCM", specify_prefix=None, choice_func=None):
+    """
+    Load checkpoint info from a specified file asyncly.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Note:
+        - `specify_prefix` and `filter_prefix` do not affect each other.
+        - If none of the parameters are loaded from checkpoint file, it will throw ValueError.
+        - `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.
+
+    Args:
+        ckpt_file_name (str): Checkpoint file name.
+        net (Cell, optional): The network where the parameters will be loaded. Default: ``None`` .
+        strict_load (bool, optional): Whether to strict load the parameter into net. If ``False`` , it will load
+                                      parameter 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`` .
+        filter_prefix (Union[str, list[str], tuple[str]], optional): Deprecated(see `choice_func`). Parameters
+            starting with the `filter_prefix` will not be loaded. Default: ``None`` .
+        dec_key (Union[None, bytes], optional): Byte type key used for decryption. If the value is ``None`` ,
+                                                the decryption is not required. Default: ``None`` .
+        dec_mode (str, optional): This parameter is valid only when dec_key is not set to ``None`` . Specifies
+                                  the decryption mode, currently supports ``"AES-GCM"`` and ``"AES-CBC"``
+                                  and ``"SM4-CBC"`` . Default: ``"AES-GCM"`` .
+        specify_prefix (Union[str, list[str], tuple[str]], optional): Deprecated(see `choice_func`). Parameters
+            starting with the specify_prefix will be loaded. Default: ``None`` .
+        choice_func (Union[None, function], optional): Input value of the function is a Parameter name of type
+            string, and the return value is a bool. If returns ``True`` , the Parameter
+            that matches the custom condition will be loaded. If returns ``False`` , the Parameter that
+            matches the custom condition will be removed. Default: ``None`` .
+
+    Returns:
+        A custom inner class, calling its `result` method yields the :func:`mindspore.load_checkpoint` result.
+
+    Raises:
+        ValueError: Checkpoint file's format is incorrect.
+        ValueError: Parameter's dict is None after load checkpoint file.
+        TypeError: The type of `specify_prefix` or `filter_prefix` is incorrect.
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import nn
+        >>> from mindspore.train import Model
+        >>> from mindspore.amp import FixedLossScaleManager
+        >>> from mindspore import context
+        >>> from mindspore import load_checkpoint_async
+        >>> from mindspore import load_param_into_net
+        >>> context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
+        >>> # Create the dataset taking MNIST as an example. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
+        >>> dataset = create_dataset()
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> ckpt_file = "./checkpoint/LeNet5-1_32.ckpt"
+        >>> net = LeNet5()
+        >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
+        >>> loss_scale_manager = FixedLossScaleManager()
+        >>> optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
+        >>> model = Model(net, loss_fn=loss, optimizer=optim, metrics=None,
+        ...               loss_scale_manager=loss_scale_manager)
+        >>> pd_future = load_checkpoint_async(ckpt_file)
+        >>> model.build(train_dataset=dataset, epoch=2)
+        >>> param_dict = pd_future.result()
+        >>> load_param_into_net(net, param_dict)
+        >>> model.train(2, dataset)
+        >>> print("param dict len: ", len(param_dict), flush=True)
+    """
+    from concurrent.futures import ThreadPoolExecutor
+    executor = ThreadPoolExecutor(max_workers=2)
+    param_dict_future = executor.submit(load_checkpoint, ckpt_file_name, net, strict_load, filter_prefix,
+                                        dec_key, dec_mode, specify_prefix, choice_func)
+    return ParamDictFuture(executor, param_dict_future)
+
+
 def _load_map_parameter(checkpoint_list, element, element_id, map_data_list,
                         map_shape_list, parameter_dict):
     """load map parameter."""
@@ -1198,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))
@@ -1239,17 +1546,28 @@ def _check_prefix(prefix):
     return prefix
 
 
-def _parse_ckpt_proto(ckpt_file_name, dec_key, dec_mode):
+def _parse_ckpt_proto(ckpt_file_name, dec_key, dec_mode, crc_check):
     """Parse checkpoint protobuf."""
     checkpoint_list = Checkpoint()
     try:
         if dec_key is None:
-            with open(ckpt_file_name, "rb") as f:
+            with _ckpt_fs.open(ckpt_file_name, *_ckpt_fs.open_args) as f:
                 pb_content = f.read()
         else:
             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":
+            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:]
+            pb_content = pb_content[:-17]
+            if crc_check:
+                crc_num = int.from_bytes(crc_num_bytes, byteorder='big')
+                cal_crc_num = binascii.crc32(pb_content, 0)
+                if cal_crc_num != crc_num:
+                    raise ValueError("For 'load_checkpoint', the crc check is failed, "
+                                     "please check whether the ckpt file is damaged.")
         checkpoint_list.ParseFromString(pb_content)
     except BaseException as e:
         if _is_cipher_file(ckpt_file_name):
@@ -1282,17 +1600,40 @@ def _whether_load_param(specify_prefix, filter_prefix, param_name):
 
 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 load_param_into_net(net, parameter_dict, strict_load=False):
+def _check_load_param_into_net(net, parameter_dict):
+    """check load_param_into_net"""
+    if not isinstance(net, nn.Cell):
+        logger.critical("Failed to combine the net and the parameters.")
+        msg = ("For 'load_param_into_net', the argument 'net' should be a Cell, but got {}.".format(type(net)))
+        raise TypeError(msg)
+    if not isinstance(parameter_dict, dict):
+        logger.critical("Failed to combine the net and the parameters.")
+        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")
+
+
+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,
@@ -1301,10 +1642,13 @@ 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.
-        ckpt_not_load (List), the parameter name in checkpoint file which are not loaded into the network.
+        - param_not_load (List), the parameter name in model which are not loaded into the network.
+        - ckpt_not_load (List), the parameter name in checkpoint file which are not loaded into the network.
 
     Raises:
         TypeError: Argument is not a Cell, or parameter_dict is not a Parameter dictionary.
@@ -1313,7 +1657,7 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
         >>> import mindspore as ms
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> ckpt_file_name = "./checkpoint/LeNet5-1_32.ckpt"
         >>> param_dict = ms.load_checkpoint(ckpt_file_name, filter_prefix="conv1")
@@ -1323,20 +1667,9 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
 
     Tutorial Examples:
         - `Saving and Loading the Model - Saving and Loading the Model Weight
-          <https://mindspore.cn/tutorials/en/r2.2/beginner/save_load.html#saving-and-loading-the-model-weight>`_
+          <https://mindspore.cn/tutorials/en/master/beginner/save_load.html#saving-and-loading-the-model-weight>`_
     """
-    if not isinstance(net, nn.Cell):
-        logger.critical("Failed to combine the net and the parameters.")
-        msg = ("For 'load_param_into_net', the argument 'net' should be a Cell, but got {}.".format(type(net)))
-        raise TypeError(msg)
-    if not isinstance(parameter_dict, dict):
-        logger.critical("Failed to combine the net and the parameters.")
-        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")
+    _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.")
@@ -1345,8 +1678,11 @@ 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.")
-    if not _is_in_auto_parallel_mode():
+    for _, param in net.parameters_and_names():
+        param.from_ckpt = True
+    if not (_is_in_auto_parallel_mode() or _is_parallel_mode()):
         net.init_parameters_data()
     else:
         _init_parameter_data_in_parallel_mode(net, parameter_dict)
@@ -1360,7 +1696,7 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
             # Add has attr protection when load server checkpoint file on worker.
             if not hasattr(parameter_dict[param.name], "data"):
                 continue
-            new_param = copy.deepcopy(parameter_dict[param.name])
+            new_param = parameter_dict[param.name]
             _update_param(param, new_param, strict_load)
             ckpt_not_load.remove(param.name)
         else:
@@ -1369,18 +1705,31 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
     if param_not_load and not strict_load:
         _load_dismatch_prefix_params(net, parameter_dict, param_not_load, strict_load)
 
-    logger.debug("Params not matched(in net but not in parameter_dict):")
-    for param_name in param_not_load:
-        logger.debug("%s", param_name)
-
     logger.info("Loading parameters into net is finished.")
     if param_not_load:
         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)))
-        for param_name in param_not_load:
-            logger.warning("{} is not loaded.".format(param_name))
+                       "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 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
+        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
 
 
@@ -1486,7 +1835,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:
@@ -1494,6 +1843,23 @@ def _save_graph(network, file_name):
             f.write(graph_pb)
 
 
+def _reshape_tensor(tensor, dst_shape):
+    """reshape tensor to dst shape"""
+    np_tensor = tensor.asnumpy()
+    np_tensor = np_tensor.reshape(dst_shape)
+    return Tensor(np_tensor, tensor.dtype)
+
+
+def _check_param_for_integrate_save(pipeline_stages, uniform_split):
+    """check whether current settings and parameters are supported in integrated save checkpoint mode"""
+    if pipeline_stages > 1:
+        raise RuntimeError("Pipeline Parallel don't support Integrated save checkpoint now.")
+    if uniform_split == 0:
+        raise RuntimeError("For 'save_checkpoint' and in automatic model parallel scene, when set "
+                           "'integrated_save' to True, the checkpoint will be integrated save, it "
+                           "is only supports uniform split tensor now.")
+
+
 def _get_merged_param_data(net, parameter_layout_dict, param_name, param_data, integrated_save):
     """
     Gets the merged data(tensor) from tensor slice, by device arrangement and tensor map.
@@ -1507,7 +1873,7 @@ def _get_merged_param_data(net, parameter_layout_dict, param_name, param_data, i
         Tensor, the combined tensor which with the whole data value.
     """
     layout = parameter_layout_dict[param_name]
-    if len(layout) < 6:
+    if len(layout) < 8:
         logger.info("The layout dict does not contain the key %s", param_name)
         return param_data
 
@@ -1515,6 +1881,13 @@ def _get_merged_param_data(net, parameter_layout_dict, param_name, param_data, i
     tensor_map = layout[1]
     uniform_split = layout[4]
     opt_shard_group = layout[5]
+    before_reshape_slice_shape = layout[2]
+    before_reshape_full_shape = layout[6]
+    after_reshape_slice_shape = layout[7]
+    do_reshape = False
+    if before_reshape_full_shape and after_reshape_slice_shape \
+            and after_reshape_slice_shape != before_reshape_slice_shape:
+        do_reshape = True
 
     allgather_net = None
     mp_weight = False
@@ -1527,26 +1900,26 @@ def _get_merged_param_data(net, parameter_layout_dict, param_name, param_data, i
     else:
         logger.info("Need to create allgather net for %s", param_name)
         if integrated_save:
-            if context.get_auto_parallel_context("pipeline_stages") > 1:
-                raise RuntimeError("Pipeline Parallel don't support Integrated save checkpoint now.")
-            if uniform_split == 0:
-                raise RuntimeError("For 'save_checkpoint' and in automatic model parallel scene, when set "
-                                   "'integrated_save' to True, the checkpoint will be integrated save, it "
-                                   "is only supports uniform split tensor now.")
+            _check_param_for_integrate_save(context.get_auto_parallel_context("pipeline_stages"), uniform_split)
             # while any dim is not equal to -1, means param is split and needs to be merged
             # pipeline parallel need to be supported here later
             if mp_weight:
-                allgather_net = get_allgather_cell(opt_shard_group, bool(opt_shard_group))
+                allgather_net = get_allgather_cell(opt_shard_group, bool(opt_shard_group), do_reshape,
+                                                   tuple(after_reshape_slice_shape))
                 object.__setattr__(allgather_net, "keep_input_unchanged", True)
             elif opt_shard_group:
-                allgather_net = get_allgather_cell(opt_shard_group, False)
+                allgather_net = get_allgather_cell(opt_shard_group, False, do_reshape,
+                                                   tuple(after_reshape_slice_shape))
         elif opt_shard_group and context.get_auto_parallel_context("optimizer_weight_shard_aggregated_save"):
-            allgather_net = get_allgather_cell(opt_shard_group, False)
+            allgather_net = get_allgather_cell(opt_shard_group, False, do_reshape,
+                                               tuple(after_reshape_slice_shape))
         net.parallel_parameter_merge_net_dict[param_name] = allgather_net
     if allgather_net:
         param_data = allgather_net(param_data)
     if mp_weight and integrated_save:
         param_data = _reshape_param_data(param_data, dev_mat, tensor_map)
+        if do_reshape:
+            param_data = _reshape_tensor(param_data, before_reshape_full_shape)
     return param_data
 
 
@@ -1556,7 +1929,8 @@ def export(net, *inputs, file_name, file_format, **kwargs):
 
     Note:
         1. When exporting AIR, ONNX format, the size of a single tensor can not exceed 2GB.
-        2. When file_name does not have a suffix, the system will automatically add one according to the file_format.
+        2. When `file_name` does not have a suffix, the system will automatically add one
+           according to the `file_format`.
         3. Exporting functions decorated with :func:`mindspore.jit` to mindir format is supported.
         4. When exporting a function decorated with :func:`mindspore.jit`, the function should not involve
            class properties in calculations.
@@ -1576,7 +1950,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.
 
@@ -1586,9 +1960,9 @@ def export(net, *inputs, file_name, file_format, **kwargs):
               - For 'AIR' and 'ONNX' models, only customized encryption is supported.
               - For 'MINDIR', all options are supported. Option: 'AES-GCM', 'AES-CBC', 'SM4-CBC'
                 or Customized encryption.
-                Default: 'AES-GCM'.
+                Default: ``'AES-GCM'``.
               - For details of using the customized encryption, please check the `tutorial
-                <https://mindspore.cn/mindarmour/docs/en/r2.0/model_encrypt_protection.html>`_.
+                <https://mindspore.cn/mindarmour/docs/en/master/model_encrypt_protection.html>`_.
 
             - dataset (Dataset): Specifies the preprocessing method of the dataset, which is used to import the
               preprocessing of the dataset into MindIR.
@@ -1602,32 +1976,49 @@ def export(net, *inputs, file_name, file_format, **kwargs):
               - customized_func (function): A python function used for customized function mode, which used for control
                 the switch branch of obfuscation structure. The outputs of customized_func should be boolean and const (
                 Reference to 'my_func()' in
-                `tutorials <https://www.mindspore.cn/mindarmour/docs/en/r2.0/dynamic_obfuscation_protection.html>`_).
+                `tutorials <https://www.mindspore.cn/mindarmour/docs/en/master/dynamic_obfuscation_protection.html>`_).
                 This function needs to ensure that its result is constant for any input. Users can refer to opaque
                 predicates. If customized_func is set, then it should be passed to `load()` interface when loading
                 obfuscated model.
               - obf_random_seed (int): Obfuscation random seed, which should be in (0, 9223372036854775807]. The
                 structure of obfuscated models corresponding to different random seeds is different. If
-                `obf_random_seed` is set, then it should be passed to :class:`nn.GraphCell()` interface when loading
+                `obf_random_seed` is set, then it should be passed
+                to :class:`mindspore.nn.GraphCell` interface when loading
                 obfuscated model. It should be noted that at least one of `customized_func` or `obf_random_seed` should
                 be set, and the latter mode would be applied if both of them are set.
 
             - incremental (bool): export MindIR incrementally.
 
+            - custom_func (function): Functions for custom defined export policies. This function will be used to
+              customize the model during network export. Currently only support for files with mindir format. The
+              function only accepts one input representing the proto object of the mindir file. When modifying a model,
+              it is necessary to ensure the correctness of the `custom_func` , otherwise it may lead to model loading
+              failure or functional errors. Default: ``None`` .
+
     Examples:
         >>> import mindspore as ms
         >>> import numpy as np
         >>> from mindspore import Tensor
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> input_tensor = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32))
         >>> ms.export(net, input_tensor, file_name='lenet', file_format='MINDIR')
+        >>>
+        >>> # Export model in MindIR format and modified the model info using custom_func
+        >>> # The custom_func only support one input representing the Proto object of the model
+        >>> # And custom_func does not support return value
+        >>> def _custom_func(mindir_model):
+        ...     mindir_model.producer_name = "test11111"
+        ...     mindir_model.producer_version = "11.0"
+        ...     mindir_model.user_info["version"] = "11.0"
+        >>> ms.export(net, input_tensor, file_name="lenet", file_format='MINDIR', custom_func=_custom_func)
+
 
     Tutorial Examples:
         - `Saving and Loading the Model - Saving and Loading MindIR
-          <https://mindspore.cn/tutorials/en/r2.2/beginner/save_load.html#saving-and-loading-mindir>`_
+          <https://mindspore.cn/tutorials/en/master/beginner/save_load.html#saving-and-loading-mindir>`_
     """
     old_ms_jit_value = context.get_context("jit_syntax_level")
     context.set_context(jit_syntax_level=mindspore.STRICT)
@@ -1658,7 +2049,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)
@@ -1690,7 +2081,7 @@ def _get_funcgraph(net, *inputs):
         >>> from mindspore import Tensor
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> input_tensor = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32))
         >>> ms.get_funcgraph(net, input_tensor)
@@ -1712,6 +2103,8 @@ def _export(net, file_name, file_format, *inputs, **kwargs):
     logger.info("exporting model file:%s format:%s.", file_name, file_format)
     if "obf_config" in kwargs and file_format != "MINDIR":
         raise ValueError(f"Dynamic obfuscation only support for MindIR format, but got {file_format} format.")
+    if "custom_func" in kwargs and file_format != "MINDIR":
+        raise ValueError(f"Currently only support custom_func for MindIR format, but got {file_format} format.")
     if file_format == 'AIR':
         _save_air(net, file_name, *inputs, **kwargs)
     elif file_format == 'ONNX':
@@ -1749,8 +2142,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:
@@ -1860,24 +2253,24 @@ 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))
     data_file_name = os.path.join(dirname, external_local)
     f, parameter_size, offset = _get_data_file(is_encrypt, kwargs, data_file_name)
     try:
-        round_ = 0
+        round = 0
         names = []
         for param_proto in model.graph.parameter:
             name = param_proto.name[param_proto.name.find(":") + 1:]
             names.append((name, param_proto))
-        names.sort(key=lambda x: x[0])
+            names.sort(key=lambda x: x[0])
         for pairs in names:
             name = pairs[0]
             param_proto = pairs[1]
@@ -1900,8 +2293,8 @@ def _split_save(net_dict, model, file_name, is_encrypt, **kwargs):
             offset += (data_length + append_size)
             write_data = _encrypt_data(is_encrypt, write_data, kwargs)
             f.write(write_data)
-            round_ += 1
-            logger.debug(f"writing {round_}th split data, name:{name}")
+            round += 1
+            logger.debug(f"writing {round}th split data, name:{name}")
 
         graph_file_name = os.path.join(dirname, file_prefix + "_graph.mindir")
         if os.path.exists(graph_file_name):
@@ -1998,6 +2391,10 @@ def _save_mindir(net, file_name, *inputs, **kwargs):
         dataset = kwargs.get('dataset')
         _save_dataset_to_mindir(model, dataset)
 
+    custom_func = kwargs.get('custom_func', None)
+    if custom_func is not None:
+        custom_func(model)
+
     save_together = _save_together(net_dict, model)
     is_encrypt = lambda: 'enc_key' in kwargs.keys() and 'enc_mode' in kwargs.keys()
     if save_together:
@@ -2030,9 +2427,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:
@@ -2084,6 +2481,45 @@ def _save_dataset_to_mindir(model, dataset):
             model.preprocessor.op[-1].offload = op['offload'] if 'offload' in op.keys() else False
 
 
+def check_checkpoint(ckpt_file_name):
+    """
+    Check whether the checkpoint is valid.
+
+    Args:
+        ckpt_file_name (str): Checkpoint file name.
+
+    Returns:
+        bool, whether the checkpoint is valid.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ckpt_file_name = "./checkpoint/LeNet5-1_32.ckpt"
+        >>> check_result = ms.check_checkpoint(ckpt_file_name)
+        >>> print(check_result)
+        True
+    """
+    if not ckpt_file_name.endswith('.ckpt'):
+        return False
+    checkpoint_list = Checkpoint()
+    with _ckpt_fs.open(ckpt_file_name, *_ckpt_fs.open_args) as f:
+        pb_content = f.read()
+        if pb_content[-17:-10] == b"crc_num":
+            crc_num_bytes = pb_content[-10:]
+            pb_content = pb_content[:-17]
+            crc_num = int.from_bytes(crc_num_bytes, byteorder='big')
+            cal_crc_num = binascii.crc32(pb_content, 0)
+            if cal_crc_num != crc_num:
+                logger.warning("For 'check_checkpoint', the ckpt crc check is failed.")
+                return False
+        try:
+            checkpoint_list.ParseFromString(pb_content)
+        except google.protobuf.message.DecodeError as e:
+            logger.warning("For 'check_checkpoint', the ckpt parse is failed.")
+            logger.warning(e)
+            return False
+        return True
+
+
 def parse_print(print_file_name):
     """
     Parse data file generated by :class:`mindspore.ops.Print`.
@@ -2122,7 +2558,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 "
@@ -2411,14 +2847,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.
@@ -2428,13 +2865,21 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
                                        in at least one of them. Default: ``None`` .
         strict_load (bool): Whether to strict load the parameter into net. If ``False`` , it will load parameter
                             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
+                            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`` .
         dec_key (Union[None, bytes]): Byte type key used for decryption. If the value is ``None`` , the decryption
                                       is not required. Default: ``None`` .
         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.
@@ -2449,14 +2894,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/r2.2/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/r2.2/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/r2.2/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/docs/en/master/model_train/parallel/dynamic_cluster.html>`_ .
 
         >>> import os
         >>> import numpy as np
@@ -2538,6 +2983,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)
@@ -2582,17 +3075,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
@@ -2722,11 +3222,10 @@ def _merge_and_split(sliced_params, train_strategy, predict_strategy):
     param_name = merged_param.name
     tensor_layout = predict_strategy[param_name]
     rank = get_rank()
-    split_tensor = _load_tensor(merged_param.data, tensor_layout[0], tensor_layout[1], rank)
+    split_tensor = _load_tensor(merged_param.data, tensor_layout[0], tensor_layout[1], rank_id=rank)
     requires_grad = merged_param.requires_grad
     layerwise_parallel = merged_param.layerwise_parallel
-    data_type = merged_param.data.dtype
-    if data_type == mstype.bfloat16:
+    if merged_param.data.dtype == mstype.bfloat16:
         split_param = Parameter(Tensor(split_tensor, mstype.bfloat16), param_name, requires_grad, layerwise_parallel)
     else:
         split_param = Parameter(split_tensor, param_name, requires_grad, layerwise_parallel)
@@ -2765,7 +3264,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 = []
 
@@ -2794,7 +3293,7 @@ def _get_mindir_inputs(file_name):
 
 def convert_model(mindir_file, convert_file, file_format):
     """
-    Convert mindir model to other format model. Current version only support convert to "ONNX" format.
+    Convert mindir model to other format model. The current version only supports conversion to ONNX models.
 
     .. warning::
         This is an experimental API that is subject to change or deletion.