mindspore 2.2.11__cp39-cp39-win_amd64.whl → 2.3.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,6 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 
+import binascii
 import copy
 import json
 import os
@@ -30,6 +31,7 @@ from io import BytesIO
 import math
 import sys
 import time
+import google
 import numpy as np
 
 from mindspore.train.checkpoint_pb2 import Checkpoint
@@ -50,9 +52,12 @@ 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.ops import Cast
 from mindspore.parallel._cell_wrapper import get_allgather_cell
 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
@@ -61,21 +66,23 @@ from mindspore.parallel._parallel_serialization import _convert_to_list, _conver
     _restore_group_info_list
 from mindspore.parallel._ps_context import _set_checkpoint_load_status, _store_warm_up_ptr_by_tensor, \
     _store_warm_up_ptr_by_tensor_list, _cache_enable
+from mindspore.parallel.checkpoint_transform import sync_pipeline_shared_parameters
 from mindspore.train._utils import read_proto
 from mindspore._c_expression import load_mindir, _encrypt, _decrypt, _is_cipher_file, dynamic_obfuscate_mindir, \
     split_mindir, split_dynamic_mindir
+from mindspore.common.generator import Generator
+from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
+from mindspore.parallel.parameter_broadcast import parameter_broadcast
 from ..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 +102,30 @@ 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)
+
+
+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):
     """
@@ -176,7 +207,7 @@ def _update_param(param, new_param, strict_load):
 
 def _type_convert(param, new_param, strict_load):
     """Whether to convert parameter's type during load checkpoint into network."""
-    float_type = (mstype.float16, mstype.float32, mstype.float64)
+    float_type = (mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16)
     int_type = (mstype.int8, mstype.int16, mstype.int32, mstype.int64)
     if not strict_load and ({param.data.dtype, new_param.data.dtype}.issubset(float_type) or
                             {param.data.dtype, new_param.data.dtype}.issubset(int_type)):
@@ -221,18 +252,19 @@ 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):
     """Execute the process of saving checkpoint into file."""
     try:
         with _ckpt_mutex:
             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:
+            with _ckpt_fs.create(ckpt_file_name, *_ckpt_fs.create_args) as f:
                 plain_data = None
                 if enc_key is not None:
                     plain_data = BytesIO()
 
+                crc_num = 0
                 for name, value in data_list.items():
                     if name == "random_op":
                         _write_random_seed(name, value, f)
@@ -242,21 +274,21 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
                         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)
+                        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[0] == "BFloat16_tensor":
-                        _write_bfloat16_data(name, value, f, enc_key, plain_data)
+                    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):
+                    if isinstance(value[2], Tensor) and hasattr(value[2], "slice_num") and value[2].slice_num > 1:
                         _write_hugeparameter(name, value, f)
                         continue
 
-                    _write_parameter_data(name, value, f, enc_key, plain_data)
+                    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)
@@ -266,7 +298,10 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
                         f.write(_encrypt(block_data, len(block_data), enc_key, len(enc_key), enc_mode))
                         block_data = plain_data.read(max_block_size)
 
-                os.chmod(ckpt_file_name, stat.S_IRUSR)
+                if crc_check:
+                    f.write('crc_num'.encode() + crc_num.to_bytes(10, byteorder='big'))
+
+            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, "
@@ -286,22 +321,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 +340,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."""
@@ -384,10 +434,14 @@ def _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name):
 
 
 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, **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 +463,8 @@ 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`` .
         kwargs (dict): Configuration options dictionary.
 
     Raises:
@@ -420,7 +476,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,7 +496,7 @@ 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)
     integrated_save = Validator.check_bool(integrated_save)
@@ -448,24 +504,32 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
     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.")
+    global_step_num = kwargs.get('global_step_num', None)
 
     save_obj = _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choice_func)
 
     if append_dict:
         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"] = param["data"]
+                else:
+                    data_list["random_op"] = param["data"]
                 continue
             key = param["name"]
             data_list[key] = []
@@ -479,49 +543,41 @@ 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] = np.array(param["data"])
+                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] = param["data"].asnumpy()
                 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":
+        import aiturbo
+        ckpt_name = os.path.basename(ckpt_file_name)
+        aiturbo.save_ckpt(ckpt_name, global_step_num, data_list_np)
+    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),
+                     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)
 
     logger.info("Saving checkpoint process is finished.")
 
@@ -532,7 +588,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 +655,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 +668,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 +690,13 @@ 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
 
             # in automatic model parallel scenario, some parameters were split to all the devices,
             # which should be combined before saving
             if key in parameter_layout_dict:
+                param_data = Tensor(value.data)
                 param_data = _get_merged_param_data(save_obj, parameter_layout_dict, key, param_data,
                                                     integrated_save)
 
@@ -670,9 +736,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 +765,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 +801,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 "
@@ -776,7 +842,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 +850,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.
@@ -909,13 +975,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 +990,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 +1001,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 +1069,76 @@ 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):
+    """load parameter into parameter_dict"""
+    ckpt_file_name = _check_ckpt_file_name(ckpt_file_name)
+    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):
     """
     Load checkpoint info from a specified file.
 
@@ -1034,6 +1169,7 @@ 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`` .
 
     Returns:
         Dict, key is parameter name, value is a Parameter or string. When the `append_dict` parameter of
@@ -1076,83 +1212,31 @@ 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)
     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()
+        import 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)
+        for key, value in np_dict.items():
+            if isinstance(value, str):
+                parameter_dict[key] = value
+            else:
+                parameter_dict[key] = Parameter(Tensor(value), name=key)
+    else:
+        _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter_prefix, choice_func, dec_key,
+                              dec_mode, crc_check)
 
     if not parameter_dict:
         raise ValueError(f"The loaded parameter dict is empty after filter or specify, please check whether "
@@ -1168,6 +1252,86 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
     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."""
@@ -1239,17 +1403,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,13 +1457,33 @@ 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 _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):
     """
     Load parameters into network, return parameter list that are not loaded in the network.
@@ -1303,8 +1498,8 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
                             on the parameters of the same type, such as float32 to float16. Default: ``False`` .
 
     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 +1508,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 +1518,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.")
@@ -1346,6 +1530,8 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
 
     strict_load = Validator.check_bool(strict_load)
     logger.info("Execute the process of loading parameters into net.")
+    for _, param in net.parameters_and_names():
+        param.from_ckpt = True
     if not _is_in_auto_parallel_mode():
         net.init_parameters_data()
     else:
@@ -1360,7 +1546,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 +1555,21 @@ 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))
+        logger.warning("{} are not loaded.".format(param_not_load))
+    if os.getenv("AITURBO") == "1" and net.parameter_layout_dict is not None:
+        param_layout = net.parameter_layout_dict
+        param_redundancy = get_parameter_redundancy(param_layout)
+        remove_param_redundancy_dict = remove_param_redundancy(param_redundancy)
+        target_parameter_name_set = set(parameter_dict.keys())
+        for rank_id, param_name_set in remove_param_redundancy_dict:
+            if param_name_set == target_parameter_name_set:
+                parameter_broadcast(net, param_layout, rank_id)
     return param_not_load, ckpt_not_load
 
 
@@ -1494,6 +1683,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 +1713,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 +1721,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 +1740,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,10 +1769,13 @@ 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.
+        5. AIR format is deprecated, and will be removed in a future version, please use other format or use
+           MindSpore Lite to do offline inference.
 
     Args:
         net (Union[Cell, function]): MindSpore network.
@@ -1584,9 +1800,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.
@@ -1600,32 +1816,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)
@@ -1633,6 +1866,9 @@ def export(net, *inputs, file_name, file_format, **kwargs):
     supported_formats = ['AIR', 'ONNX', 'MINDIR']
     if file_format not in supported_formats:
         raise ValueError(f"For 'export', 'file_format' must be one of {supported_formats}, but got {file_format}.")
+    if file_format == 'AIR':
+        logger.warning("AIR format is deprecated, and will be removed in a future version, please use other format or "
+                       "use MindSpore Lite to do offline inference")
     Validator.check_file_name_by_regular(file_name)
     logger.info("exporting model file:%s format:%s.", file_name, file_format)
 
@@ -1685,7 +1921,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)
@@ -1707,6 +1943,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':
@@ -1867,12 +2105,12 @@ def _split_save(net_dict, model, file_name, is_encrypt, **kwargs):
     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]
@@ -1895,8 +2133,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):
@@ -1993,6 +2231,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:
@@ -2079,6 +2321,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`.
@@ -2423,7 +2704,7 @@ 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`` .
@@ -2444,14 +2725,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/tutorials/experts/en/master/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/tutorials/experts/en/master/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/tutorials/experts/en/master/parallel/dynamic_cluster.html>`_ .
 
         >>> import os
         >>> import numpy as np
@@ -2717,11 +2998,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)
@@ -2789,7 +3069,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.