mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/train/serialization.py

@@ -23,7 +23,7 @@ import os
 import shutil
 import stat
 import threading
-from threading import Thread, Lock
+from threading import Thread, RLock
 from collections import defaultdict, OrderedDict
 from io import BytesIO
 
@@ -48,7 +48,7 @@ from mindspore.common.api import _MindsporeFunctionExecutor
 from mindspore.common.api import _get_parameter_layout
 from mindspore.common.api import _generate_branch_control_input
 from mindspore.common.initializer import initializer, One
-from mindspore.common.parameter import Parameter
+from mindspore.common.parameter import Parameter, _offload_if_config
 from mindspore.common.tensor import Tensor
 from mindspore.common._utils import is_shape_unknown
 from mindspore.communication.management import get_rank, get_group_size
@@ -59,8 +59,11 @@ from mindspore.parallel._tensor import _reshape_param_data, _reshape_param_data_
 from mindspore.parallel._utils import _infer_rank_list, _remove_repeated_slices, _is_in_auto_parallel_mode
 from mindspore.parallel._parallel_serialization import _convert_to_list, _convert_to_layout, _build_searched_strategy, \
     _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.train._utils import read_proto
-from mindspore._c_expression import load_mindir, _encrypt, _decrypt, _is_cipher_file, dynamic_obfuscate_mindir
+from mindspore._c_expression import load_mindir, _encrypt, _decrypt, _is_cipher_file, dynamic_obfuscate_mindir, \
+    split_mindir, split_dynamic_mindir
 from ..ops.operations._opaque_predicate_registry import add_opaque_predicate, clean_funcs
 
 tensor_to_ms_type = {"Int8": mstype.int8, "UInt8": mstype.uint8, "Int16": mstype.int16, "UInt16": mstype.uint16,
@@ -72,11 +75,13 @@ tensor_to_np_type = {"Int8": np.int8, "UInt8": np.uint8, "Int16": np.int16, "UIn
                      "Int32": np.int32, "UInt32": np.uint32, "Int64": np.int64, "UInt64": np.uint64,
                      "Float16": np.float16, "Float32": np.float32, "Float64": np.float64, "Bool": np.bool_, "str": "U"}
 
+np_type_convert = {"int32": np.int32, "float32": np.float32, "float16": np.float16, "float64": np.float64}
+
 mindir_to_tensor_type = {1: mstype.float32, 2: mstype.uint8, 3: mstype.int8, 4: mstype.uint16,
                          5: mstype.int16, 6: mstype.int32, 7: mstype.int64, 10: mstype.float16,
                          11: mstype.float64, 12: mstype.uint32, 13: mstype.uint64}
 
-_ckpt_mutex = Lock()
+_ckpt_mutex = RLock()
 
 # unit is KB
 SLICE_SIZE = 512 * 1024
@@ -124,7 +129,7 @@ def _update_param(param, new_param, strict_load):
         if param.data.dtype != new_param.data.dtype:
             if _type_convert(param, new_param, strict_load):
                 new_tensor = Tensor(new_param.data.asnumpy(), param.data.dtype)
-                param.set_data(new_tensor)
+                param.set_data(new_tensor, param.sliced)
                 return
 
             logger.critical("Failed to combine the net and the parameters for param %s.", param.name)
@@ -205,7 +210,7 @@ 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"):
+def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_param_inc=False):
     """Execute the process of saving checkpoint into file."""
     try:
         with _ckpt_mutex:
@@ -213,37 +218,28 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM"):
                 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)
+                        _write_mapparameter(name, value, f, map_param_inc)
+                        continue
+                    if value[0] == "offload_parameter":
+                        new_value = value[1:]
+                        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 isinstance(value[2], Tensor):
                         _write_hugeparameter(name, value, f)
                         continue
 
-                    data_size = value[2].nbytes / 1024
-                    if data_size > SLICE_SIZE:
-                        slice_count = math.ceil(data_size / SLICE_SIZE)
-                        param_slice_list = np.array_split(value[2], slice_count)
-                    else:
-                        param_slice_list = [value[2]]
-
-                    for param_slice in param_slice_list:
-                        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 = param_slice.tobytes()
-
-                        if enc_key is None:
-                            f.write(checkpoint_list.SerializeToString())
-                        else:
-                            plain_data.write(checkpoint_list.SerializeToString())
+                    _write_parameter_data(name, value, f, enc_key, plain_data)
 
                 if enc_key is not None:
                     plain_data.seek(0)
@@ -261,18 +257,59 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM"):
         raise e
 
 
-def _write_mapparameter(name, value, f):
-    """Write map parameter into protobuf file."""
+def _write_random_seed(name, value, f):
+    """Write random op into protobuf file."""
     checkpoint_list = Checkpoint()
     param_value = checkpoint_list.value.add()
     param_value.tag = name
-    map_tensor = param_value.maptensor
-    for v in value[1:]:
-        tensor = map_tensor.tensor.add()
-        tensor.dims.extend(v[0])
-        tensor.tensor_type = v[1]
-        tensor.tensor_content = v[2].tobytes()
+    param_tensor = param_value.tensor
+    param_tensor.dims.extend(0)
+    param_tensor.tensor_type = "random_op"
+    param_tensor.tensor_content = value
+    f.write(checkpoint_list.SerializeToString())
+
+
+def _write_parameter_data(name, value, f, enc_key, plain_data):
+    """Write parameter data into protobuf file."""
+    data_size = value[2].nbytes / 1024
+    if data_size > SLICE_SIZE:
+        slice_count = math.ceil(data_size / SLICE_SIZE)
+        param_slice_list = np.array_split(value[2], slice_count)
+    else:
+        param_slice_list = [value[2]]
+
+    for param_slice in param_slice_list:
+        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 = param_slice.tobytes()
+
+        if enc_key is None:
+            f.write(checkpoint_list.SerializeToString())
+        else:
+            plain_data.write(checkpoint_list.SerializeToString())
+
+
+def _write_mapparameter(name, value, f, map_param_inc=False):
+    """Write map parameter into protobuf file."""
+    while True:
+        logger.info("Checkpoint save map_parameter.")
+        data_map_slice = value[1].export_slice_data(map_param_inc)
+        checkpoint_list = Checkpoint()
+        param_value = checkpoint_list.value.add()
+        param_value.tag = name
+        map_tensor = param_value.maptensor
+        for numpy_data in data_map_slice[:3]:
+            tensor_pro = map_tensor.tensor.add()
+            tensor_pro.dims.extend(numpy_data.shape)
+            tensor_pro.tensor_type = str(numpy_data.dtype)
+            tensor_pro.tensor_content = numpy_data.reshape(-1).tobytes()
         f.write(checkpoint_list.SerializeToString())
+        if data_map_slice[3]:
+            break
 
 
 def _write_hugeparameter(name, value, f):
@@ -298,8 +335,8 @@ def _write_hugeparameter(name, value, f):
 
 def _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name):
     """Check save_obj and ckpt_file_name for save_checkpoint."""
-    if not isinstance(save_obj, nn.Cell) and not isinstance(save_obj, list):
-        raise TypeError("For 'save_checkpoint', the parameter 'save_obj' must be nn.Cell or list, "
+    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)))
     if not isinstance(ckpt_file_name, str):
         raise TypeError("For 'save_checkpoint', the parameter {} for checkpoint file name is invalid,"
@@ -315,34 +352,63 @@ 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"):
-    """
+                    async_save=False, append_dict=None, enc_key=None, enc_mode="AES-GCM", choice_func=None, **kwargs):
+    r"""
     Save checkpoint to a specified file.
 
     Args:
-        save_obj (Union[Cell, list]): The cell object or data list(each element is a dictionary, like
-                                      [{"name": param_name, "data": param_data},...], the type of
-                                      param_name would be string, and the type of param_data would
-                                      be parameter or Tensor).
+        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
+            elements(each element is a dictionary, like [{"name": param_name, "data": param_data},...], the type of
+            `param_name` must be string, and the type of `param_data` must be parameter or Tensor); If dict,
+            it can be the returned value of `mindspore.load_checkpoint()`.
         ckpt_file_name (str): Checkpoint file name. If the file name already exists, it will be overwritten.
-        integrated_save (bool): Whether to integrated save in automatic model parallel scene. Default: True
-        async_save (bool): Whether to open an independent thread to save the checkpoint file. Default: False
+        integrated_save (bool): Whether to integrated save in automatic model parallel scene. Default: ``True`` .
+        async_save (bool): Whether to open an independent thread to save the checkpoint file. Default: ``False`` .
         append_dict (dict): Additional information that needs to be saved. The key of dict must be str, the value
-                            of dict must be one of int, float, bool, string, Parameter or Tensor. Default: None.
-        enc_key (Union[None, bytes]): Byte type key used for encryption. If the value is None, the encryption
-                                      is not required. Default: None.
-        enc_mode (str): This parameter is valid only when enc_key is not set to None. Specifies the encryption
-                        mode, currently supports 'AES-GCM' and 'AES-CBC' and 'SM4-CBC'. Default: 'AES-GCM'.
+                            of dict must be one of int, float, bool, string, Parameter or Tensor. Default: ``None`` .
+        enc_key (Union[None, bytes]): Byte type key used for encryption. If the value is ``None`` , the encryption
+                                      is not required. Default: ``None`` .
+        enc_mode (str): This parameter is valid only when enc_key is not set to ``None`` . Specifies the encryption
+                        mode, currently supports ``"AES-GCM"`` and ``"AES-CBC"`` and ``"SM4-CBC"`` .
+                        Default: ``"AES-GCM"`` .
+        choice_func (function) : A function for saving custom selected parameters. The input value of `choice_func` is
+                                 a parameter name in string type, and the returned value is a bool.
+                                 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`` .
+        kwargs (dict): Configuration options dictionary.
 
     Raises:
-        TypeError: If the parameter save_obj is not `nn.Cell` or list type. And if the parameter `integrated_save`
-                   and `async_save` are not bool type. If the parameter ckpt_file_name is not string type.
+        TypeError: If the parameter `save_obj` is not :class:`mindspore.nn.Cell` , list or dict type.
+        TypeError: If the parameter `integrated_save` or `async_save` is not bool type.
+        TypeError: If the parameter `ckpt_file_name` is not string type.
 
     Examples:
         >>> import mindspore as ms
         >>>
-        >>> net = Net()
-        >>> ms.save_checkpoint(net, "lenet.ckpt")
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/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"))
+        >>> param_dict1 = ms.load_checkpoint("./lenet.ckpt")
+        >>> print(param_dict1)
+        {'conv2.weight': Parameter (name=conv2.weight, shape=(16, 6, 5, 5), dtype=Float32, requires_grad=True)}
+        >>> params_list = net.trainable_params()
+        >>> ms.save_checkpoint(params_list, "./lenet_list.ckpt",
+        ...                    choice_func=lambda x: x.startswith("conv") and not x.startswith("conv2"))
+        >>> param_dict2 = ms.load_checkpoint("./lenet_list.ckpt")
+        >>> print(param_dict2)
+        {'conv1.weight': Parameter (name=conv1.weight, shape=(6, 1, 5, 5), dtype=Float32, requires_grad=True)}
+        >>> ms.save_checkpoint(param_dict2, "./lenet_dict.ckpt")
+        >>> param_dict3 = ms.load_checkpoint("./lenet_dict.ckpt")
+        >>> print(param_dict3)
+        {'conv1.weight': Parameter (name=conv1.weight, shape=(6, 1, 5, 5), dtype=Float32, requires_grad=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>`_
     """
     ckpt_file_name = _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name)
     integrated_save = Validator.check_bool(integrated_save)
@@ -350,46 +416,10 @@ 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)
-
+    map_param_inc = kwargs.get('incremental', False)
     logger.info("Execute the process of saving checkpoint files.")
 
-    if isinstance(save_obj, nn.Cell):
-        parameter_layout_dict = save_obj.parameter_layout_dict
-        if _is_in_auto_parallel_mode() and not parameter_layout_dict:
-            parameter_layout_dict = _get_parameter_layout()
-        save_obj.init_parameters_data()
-        param_dict = OrderedDict()
-        for _, param in save_obj.parameters_and_names():
-            param_dict[param.name] = param
-        param_list = []
-        for (key, value) in param_dict.items():
-            each_param = {"name": key}
-            if isinstance(value, MapParameter):
-                param_data = []
-                for export_data in value.export_data():
-                    param_data.append(Tensor(export_data))
-                    each_param["data"] = param_data
-                    param_list.append(each_param)
-                continue
-
-            if value.data.is_persistent_data():
-                # list save persistent_data: [Tensor, shape, type, param.key]
-                param_data = ["persistent_data"]
-                param_data.append(value.data)
-                param_data.append(value.param_info.origin_shape)
-                param_data.append(str(value.dtype))
-                param_data.append(value.key)
-            else:
-                param_data = Tensor(value.data.asnumpy())
-
-            # 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 = _get_merged_param_data(save_obj, parameter_layout_dict, key, param_data, integrated_save)
-
-            each_param["data"] = param_data
-            param_list.append(each_param)
-        save_obj = param_list
+    save_obj = _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choice_func)
 
     if append_dict:
         append_info_list = []
@@ -397,19 +427,27 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
             if not isinstance(value, str):
                 value = Tensor(value)
             append_info_list.append({"name": k_name, "data": value})
-            save_obj.extend(append_info_list)
+        save_obj.extend(append_info_list)
 
     data_list = OrderedDict()
     with _ckpt_mutex:
         for param in save_obj:
+            if param["name"] == "random_op":
+                data_list["random_op"] = param["data"]
+                continue
             key = param["name"]
             data_list[key] = []
+            if isinstance(param["data"], MapParameter):
+                data_list[param["name"]].append("mapparameter")
+                data_list[param["name"]].append(param["data"])
+                continue
             if isinstance(param["data"], list):
                 if param["data"][0] == "persistent_data":
-                    _save_persistent_data(data_list, key, param)
-                else:
-                    _save_mapparameter(data_list, param)
-                continue
+                    _save_param_list_data(data_list, key, param)
+                elif param["data"][0] == "offload_parameter":
+                    data_list[key].append("offload_parameter")
+                    _save_param_list_data(data_list, key, param)
+
             if isinstance(param["data"], str):
                 data_list[key].append([0])
                 data_list[key].append('str')
@@ -435,28 +473,130 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
         thr = Thread(target=_exec_save, args=(ckpt_file_name, data_copy, enc_key, enc_mode), name="asyn_save_ckpt")
         thr.start()
     else:
-        _exec_save(ckpt_file_name, data_list, enc_key, enc_mode)
+        _exec_save(ckpt_file_name, data_list, enc_key, enc_mode, map_param_inc)
 
     logger.info("Saving checkpoint process is finished.")
 
 
-def _save_mapparameter(data_list, param):
-    """Save map parameter into save_obj."""
-    data_list[param["name"]].append("mapparameter")
-    for value in param["data"]:
-        dims = []
-        tmp_list = []
-        for dim in value.shape:
-            dims.append(dim)
-        tmp_list.append(dims)
-        tensor_type = str(value.dtype)
-        tmp_list.append(tensor_type)
-        data = value.asnumpy().reshape(-1)
-        tmp_list.append(data)
-        data_list[param["name"]].append(tmp_list)
+def _convert_list_to_param_list(save_obj, choice_func):
+    """Convert a list of Parameter to param_list."""
+    param_list = []
+    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"])]
+    else:
+        for param in save_obj:
+            if isinstance(param, Parameter):
+                if choice_func is not None and not choice_func(param.name):
+                    continue
+                each_param = {"name": param.name, "data": param}
+                param_list.append(each_param)
+            else:
+                raise TypeError(f"For save_checkpoint, when save_obj is made up by list of Parameter,"
+                                f"the param should be parameter, but got {type(param)}")
+    return param_list
+
+
+def _convert_dict_to_param_dict(save_obj, choice_func):
+    """Convert a dict of Parameter to param_list."""
+    param_list = []
+    for (key, value) in save_obj.items():
+        if isinstance(key, str) and isinstance(value, (Parameter, str)):
+            if choice_func is not None and not choice_func(key):
+                continue
+            each_param = {"name": key, "data": value}
+            param_list.append(each_param)
+        else:
+            raise TypeError(f"For save_checkpoint, when save_obj is made up by dict, the key should be str and"
+                            f"value should be Parameter, but got the type of key is {type(key)} and"
+                            f"the type of value is {type(value)}")
+    return param_list
+
+
+def _convert_cell_param_and_names_to_dict(save_obj, choice_func):
+    """Convert cell.parameters_and_names to OrderedDict."""
+    param_dict = OrderedDict()
+    for _, param in save_obj.parameters_and_names():
+        not_sliced = not param.sliced
+        is_graph_mode = context.get_context('mode') == context.GRAPH_MODE
+        # All parameters are initialized immediately under PyNative mode, skip this judgement.
+        judgment = not_sliced or param.has_init
+        if is_graph_mode and _is_in_auto_parallel_mode() and judgment:
+            continue
+        if choice_func is not None and not choice_func(param.name):
+            continue
+        # Add suffix for cache_enabled parameter, and then parameter can carry key info.
+        # Notice that suffix needs be removed when loading into net.
+        if param.cache_enable:
+            param_dict[param.name + ".__param_key__" + str(param.key)] = param
+        else:
+            param_dict[param.name] = param
+    return param_dict
 
 
-def _save_persistent_data(data_list, key, param):
+def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_func):
+    """Convert nn.Cell to param_list."""
+    param_list = []
+    parameter_layout_dict = save_obj.parameter_layout_dict
+    if _is_in_auto_parallel_mode() and not parameter_layout_dict:
+        parameter_layout_dict = _get_parameter_layout()
+    if not _is_in_auto_parallel_mode():
+        save_obj.init_parameters_data()
+    param_dict = _convert_cell_param_and_names_to_dict(save_obj, choice_func)
+    if append_dict and "random_op" in append_dict:
+        phase = 'train' + '.' + str(save_obj.create_time) + '.' + str(id(save_obj)) + '.' + save_obj.arguments_key
+        if phase in save_obj.compile_cache and _executor.has_compiled(phase):
+            random_byte = _executor._graph_executor.get_random_status(phase)
+            param_list.append({"name": "random_op", "data": random_byte})
+        append_dict.pop("random_op")
+    for (key, value) in param_dict.items():
+        each_param = {"name": key}
+        if isinstance(value, MapParameter):
+            each_param["data"] = value
+            param_list.append(each_param)
+            continue
+
+        if value.data.is_persistent_data():
+            # list save persistent_data: [Tensor, shape, type, param.key]
+            param_data = ["persistent_data", value.data, value.param_info.origin_shape, str(value.dtype), value.key]
+        elif value.data.offload_file_path() != "":
+            # list save offload data: [Param, shape, type, param.key]
+            param_data = ["offload_parameter"]
+            param_tensor = value.data
+            if key in parameter_layout_dict:
+                param_tensor = _get_merged_param_data(save_obj, parameter_layout_dict, key, param_tensor,
+                                                      integrated_save)
+            param_data.append(param_tensor)
+            param_data.append(param_tensor.shape)
+            param_data.append(str(param_tensor.dtype))
+            param_data.append(value.key)
+        else:
+            param_data = Tensor(value.data.asnumpy())
+
+            # 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 = _get_merged_param_data(save_obj, parameter_layout_dict, key, param_data,
+                                                    integrated_save)
+
+        each_param["data"] = param_data
+        param_list.append(each_param)
+    return param_list
+
+
+def _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choice_func):
+    """Convert a save_obj to param_list."""
+    if isinstance(save_obj, list):
+        return _convert_list_to_param_list(save_obj, choice_func)
+
+    if isinstance(save_obj, dict):
+        return _convert_dict_to_param_dict(save_obj, choice_func)
+
+    return _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_func)
+
+
+def _save_param_list_data(data_list, key, param):
     """Save persistent data into save_obj."""
     dims = []
     # persistent_data shape can not be ()
@@ -511,7 +651,7 @@ def load(file_name, **kwargs):
 
             - 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.0/api_python/mindspore/mindspore.obfuscate_model.html>` .
+              <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore/mindspore.obfuscate_model.html>`_.
 
     Returns:
         GraphCell, a compiled graph that can executed by `GraphCell`.
@@ -538,6 +678,10 @@ def load(file_name, **kwargs):
         [[[[4. 6. 4.]
            [6. 9. 6.]
            [4. 6. 4.]]]]
+
+    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>`_
     """
     if not isinstance(file_name, str):
         raise ValueError("For 'load', the argument 'file_name' must be string, but "
@@ -578,6 +722,57 @@ def load(file_name, **kwargs):
     return graph
 
 
+def export_split_mindir(file_name, device_num=8, rank_id=0, dynamic=True, sapp=False):
+    """
+    Auto Split MindIR.
+
+    The returned object can be executed by a `GraphCell`, see class :class:`mindspore.nn.GraphCell` for more details.
+
+    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.
+
+    Raises:
+        ValueError: MindIR file does not exist or `file_name` is not a string.
+        RuntimeError: Failed to split MindIR file.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> context.set_context(mode=context.GRAPH_MODE)
+        >>>
+        >>> ms.export_split_mindir("net.mindir", device_num=8, rank_id=0)
+
+    """
+    if not isinstance(file_name, str):
+        raise ValueError("For 'Split MindIR', the argument 'file_name' must be string, but "
+                         "got {}.".format(type(file_name)))
+    if not file_name.endswith(".mindir"):
+        raise ValueError("For 'Split MindIR', the argument 'file_name'(MindIR file) should end with '.mindir', "
+                         "please input the correct 'file_name'.")
+    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)
+
+    logger.info("Execute the process of export and split mindir.")
+    dynamic = True
+    if dynamic:
+        graph = split_dynamic_mindir(file_name, device_num, rank_id, sapp)
+    else:
+        graph = split_mindir(file_name)
+
+    if graph is None:
+        if _is_cipher_file(file_name):
+            raise RuntimeError("Export and split MindIR failed. The file may be encrypted and decrypt failed, you "
+                               "can check whether the values of the arguments 'dec_key' and 'dec_mode'"
+                               " are the same as when exported MindIR file, or check the file integrity.")
+        raise RuntimeError("Export and split MindIR failed.")
+    return graph
+
+
 def _check_param_type(param_config, key, target_type, requested):
     """check type of parameters"""
     if key in param_config:
@@ -655,17 +850,20 @@ def obfuscate_model(obf_config, **kwargs):
             - model_inputs (list(Tensor)): The inputs of the original model, the values of Tensor can be random, which
               is the same as using :func:`mindspore.export`.
             - obf_ratio (Union(float, str)): The ratio of nodes in original model that would be obfuscated. `obf_ratio`
-              should be in range of (0, 1] or in ["small", "medium", "large"].
+              should be in range of (0, 1] or in ["small", "medium", "large"]. "small", "medium" and "large" are
+              correspond to 0.1, 0.3, and 0.6 respectively.
             - 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. This
-              function needs to ensure that its result is constant for any input. Users can refer to opaque
+              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>`_).
+              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): The random seed used for determine the distribution of confusion branches and the
-              weight confusion coefficient, which should be in (0, 9223372036854775807]. If `obf_random_seed` is set,
-              then it should be passed to :class:`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.
+            - 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
+              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.
 
         kwargs (dict): Configuration options dictionary.
 
@@ -685,12 +883,14 @@ def obfuscate_model(obf_config, **kwargs):
         ValueError: If `original_model_path` is not exist or `original_model_path` is not end with '.mindir'.
 
     Examples:
+        >>> import mindspore as ms
+        >>> import mindspore.nn as nn
         >>> obf_config = {'original_model_path': "./net.mindir",
         ...          'save_model_path': "./obf_net",
         ...          'model_inputs': [input1, ],
         ...          'obf_ratio': 0.1, 'obf_random_seed': 173262358423}
-        >>> obfuscate_model(obf_config)
-        >>> obf_func = load("obf_net.mindir")
+        >>> ms.obfuscate_model(obf_config)
+        >>> obf_func = ms.load("obf_net.mindir")
         >>> obf_net = nn.GraphCell(obf_func, obf_random_seed=173262358423)
         >>> print(obf_net(input1).asnumpy())
     """
@@ -762,23 +962,24 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
 
     Args:
         ckpt_file_name (str): Checkpoint file name.
-        net (Cell): The network where the parameters will be loaded. Default: None
-        strict_load (bool): Whether to strict load the parameter into net. If False, it will load parameter
+        net (Cell): The network where the parameters will be loaded. 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
-                            on the parameters of the same type, such as float32 to float16. Default: False.
+                            on the parameters of the same type, such as float32 to float16. Default: ``False`` .
         filter_prefix (Union[str, list[str], tuple[str]]): Deprecated(see `choice_func`). Parameters starting with the
-            filter_prefix will not be loaded. Default: None.
-        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' and 'AES-CBC' and 'SM4-CBC'. Default: 'AES-GCM'.
+            filter_prefix will not be loaded. Default: ``None`` .
+        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"`` and ``"AES-CBC"`` and ``"SM4-CBC"`` .
+                        Default: ``"AES-GCM"`` .
         specify_prefix (Union[str, list[str], tuple[str]]): Deprecated(see `choice_func`). Parameters starting with the
-            specify_prefix will be loaded. Default: None.
+            specify_prefix will be loaded. Default: ``None`` .
         choice_func (Union[None, function]) : 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.
+            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:
         Dict, key is parameter name, value is a Parameter or string. When the `append_dict` parameter of
@@ -801,23 +1002,27 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
         >>> print(param_dict["conv2.weight"])
         Parameter (name=conv2.weight, shape=(16, 6, 5, 5), dtype=Float32, requires_grad=True)
         >>> def func(param_name):
-        >>>     whether_load = False
-        >>>     if param_name.startswith("conv"):
-        >>>         whether_load = True
-        >>>     if param_name.startswith("conv1"):
-        >>>         whether_load = False
-        >>>     return whether_load
+        ...     whether_load = False
+        ...     if param_name.startswith("conv"):
+        ...         whether_load = True
+        ...     if param_name.startswith("conv1"):
+        ...         whether_load = False
+        ...     return whether_load
         >>> param_dict1 = ms.load_checkpoint(ckpt_file_name, choice_func=func)
         >>> print(param_dict1["conv2.weight"])
         Parameter (name=conv2.weight, shape=(16, 6, 5, 5), dtype=Float32, requires_grad=True)
         >>> def func(param_name):
-        >>>     whether_load = False
-        >>>     if param_name.startswith("conv1"):
-        >>>         whether_load = True
-        >>>     return whether_load
+        ...     whether_load = False
+        ...     if param_name.startswith("conv1"):
+        ...         whether_load = True
+        ...     return whether_load
         >>> param_dict2 = ms.load_checkpoint(ckpt_file_name, choice_func=func)
         >>> print(param_dict2)
         {'conv1.weight': Parameter (name=conv1.weight, shape=(6, 1, 5, 5), dtype=Float32, requires_grad=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>`_
     """
     ckpt_file_name = _check_ckpt_file_name(ckpt_file_name)
     specify_prefix = _check_prefix(specify_prefix)
@@ -830,6 +1035,8 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
     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.")
@@ -837,13 +1044,19 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
             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_mapparameter(element, parameter_dict)
+                _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
@@ -856,7 +1069,8 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
             param_data_list.append(element_data)
             if (element_id == len(checkpoint_list.value) - 1) or \
                     (element.tag != checkpoint_list.value[element_id + 1].tag):
-                param_data = np.concatenate((param_data_list), axis=0)
+                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':
@@ -868,7 +1082,9 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
                         param_data = int(param_data[0])
                     if dims not in ([0], [1]):
                         param_data = param_data.reshape(list(dims))
-                    parameter_dict[element.tag] = Parameter(Tensor(param_data, ms_type), name=element.tag)
+                    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.")
 
@@ -881,14 +1097,48 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
         raise ValueError(f"The loaded parameter dict is empty after filter or specify, please check whether "
                          f"'filter_prefix' or 'specify_prefix' are set correctly.")
 
+    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)
+    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_map_parameter(checkpoint_list, element, element_id, map_data_list,
+                        map_shape_list, parameter_dict):
+    """load map parameter."""
+    logger.info("Checkpoint load map_parameter.")
+    if (element_id != len(checkpoint_list.value) - 1) and \
+            element.tag == checkpoint_list.value[element_id + 1].tag:
+        for index, tensor in enumerate(element.maptensor.tensor):
+            data = tensor.tensor_content
+            data_type = tensor.tensor_type
+            np_type = np_type_convert.get(data_type)
+            element_data = np.frombuffer(data, np_type)
+            map_data_list[index].append(element_data)
+            map_shape_list[index] += tensor.dims[0]
+    else:
+        map_array = []
+        for index, tensor in enumerate(element.maptensor.tensor):
+            data = tensor.tensor_content
+            data_type = tensor.tensor_type
+            np_type = np_type_convert.get(data_type)
+            element_data = np.frombuffer(data, np_type)
+            map_data_list[index].append(element_data)
+            new_data = b"".join(map_data_list[index])
+            param_data = np.frombuffer(new_data, np_type)
+            dims = tensor.dims
+            dims[0] += map_shape_list[index]
+            param_data = param_data.reshape(list(dims))
+            map_array.append(param_data)
+        parameter_dict[element.tag] = map_array
+
+
 def _check_ckpt_file_name(ckpt_file_name):
-    """Check function load_checkpoint's cket_file_name."""
+    """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)))
@@ -969,18 +1219,13 @@ def _whether_load_param(specify_prefix, filter_prefix, param_name):
     return whether_load
 
 
-def _load_mapparameter(element, parameter_dict):
-    """Load map parameter from ckpt file."""
-    map_array = []
-    for tensor in element.maptensor.tensor:
-        data = tensor.tensor_content
-        data_type = tensor.tensor_type
-        np_type = tensor_to_np_type.get(data_type)
-        element_data = np.frombuffer(data, np_type)
-        dims = tensor.dims
-        param_data = element_data.reshape(list(dims))
-        map_array.append(param_data)
-    parameter_dict[element.tag] = map_array
+def _init_parameter_data_in_parallel_mode(net, parameter_dict):
+    """In parallel mode, only init the paraemters in ckpt."""
+    for _, param in net.parameters_and_names():
+        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):
@@ -991,10 +1236,10 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
         net (Cell): The network where the parameters will be loaded.
         parameter_dict (dict): The dictionary generated by load checkpoint file,
                                it is a dictionary consisting of key: parameters's name, value: parameter.
-        strict_load (bool): Whether to strict load the parameter into net. If False, it will load parameter
+        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
-                            on the parameters of the same type, such as float32 to float16. Default: 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.
@@ -1006,25 +1251,33 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
     Examples:
         >>> import mindspore as ms
         >>>
-        >>> net = Net()
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/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")
         >>> param_not_load, _ = ms.load_param_into_net(net, param_dict)
         >>> print(param_not_load)
         ['conv1.weight']
+
+    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>`_
     """
     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")
     for key, value in parameter_dict.items():
-        if not isinstance(key, str) or not isinstance(value, (Parameter, str)):
+        if not isinstance(key, str) or not isinstance(value, (Parameter, str, list)):
             logger.critical("Load parameters into net failed.")
             msg = ("For 'parameter_dict', the element in the argument 'parameter_dict' should be a "
                    "'str' and 'Parameter' , but got {} and {}.".format(type(key), type(value)))
@@ -1032,11 +1285,20 @@ 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.")
-    net.init_parameters_data()
+    if not _is_in_auto_parallel_mode():
+        net.init_parameters_data()
+    else:
+        _init_parameter_data_in_parallel_mode(net, parameter_dict)
     param_not_load = []
     ckpt_not_load = list(parameter_dict.keys())
     for _, param in net.parameters_and_names():
         if param.name in parameter_dict:
+            if isinstance(param, MapParameter):
+                param.import_data(parameter_dict[param.name])
+                continue
+            # 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])
             _update_param(param, new_param, strict_load)
             ckpt_not_load.remove(param.name)
@@ -1061,6 +1323,72 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
     return param_not_load, ckpt_not_load
 
 
+def _warm_up_host_cache_enabled(parameter_dict):
+    """Warm up host cache enabled."""
+    if _cache_enable():
+        return True
+    for key in parameter_dict.keys():
+        if key.find(".__param_key__") != -1:
+            return True
+    return False
+
+
+def _warm_up_host_cache(parameter_dict, net):
+    """Warm up host cache."""
+    ms_role = os.getenv("MS_ROLE")
+    is_worker = ms_role == "MS_WORKER"
+    param_key_dict = {}
+    # Traverse key, value in parameter_dict, warm up param key and record param key into param_key_dict.
+    if is_worker:
+        net.init_parameters_data()
+        net_dict = {}
+        for name, value in net.parameters_and_names():
+            net_dict[name] = value
+        for param_name, value in parameter_dict.items():
+            pos = param_name.find(".__param_key__")
+            if pos != -1:
+                net_param_name = param_name[:pos]
+                param_key_dict[param_name] = net_param_name
+                net_value = None
+                if net_param_name not in net_dict:
+                    logger.warning("net param name : %s is not in net", net_param_name)
+                else:
+                    net_value = net_dict.get(net_param_name, None)
+                pos += len(".__param_key__")
+                param_key = int(param_name[pos:])
+                value_is_map_parameter = isinstance(value, list) and len(value) == 3
+                if value_is_map_parameter and (net_value is None or isinstance(net_value, Parameter)):
+                    key_tensor = Tensor.from_numpy(value[0])
+                    value_tensor = Tensor.from_numpy(value[1])
+                    status_tensor = Tensor.from_numpy(value[2])
+                    _store_warm_up_ptr_by_tensor_list(param_key, key_tensor, value_tensor, status_tensor)
+                elif not isinstance(value, list) and isinstance(net_value, Parameter):
+                    _store_warm_up_ptr_by_tensor(param_key, value)
+                else:
+                    logger.warning("Unknown matches parameter type %s and net_value %s", type(value), type(net_value))
+    else:
+        for param_name, value in parameter_dict.items():
+            pos = param_name.find(".__param_key__")
+            if pos != -1:
+                net_param_name = param_name[:pos]
+                param_key_dict[param_name] = net_param_name
+    # Split param key from parameter_dict since worker cannot load param key.
+    warm_up_dict = {}
+    for key, value in param_key_dict.items():
+        if is_worker:
+            warm_up_dict[value] = parameter_dict.pop(key)
+        else:
+            parameter_dict[value] = parameter_dict.pop(key)
+    return (is_worker, parameter_dict, warm_up_dict)
+
+
+def _warm_up_host_cache_post_process(is_worker, net_dict, warm_up_dict):
+    """Warm up host cache post process."""
+    if is_worker:
+        net_dict.update(warm_up_dict)
+    _set_checkpoint_load_status(True)
+
+
 def _load_dismatch_prefix_params(net, parameter_dict, param_not_load, strict_load):
     """When some net parameter did not load, try to continue loading."""
     prefix_name = ""
@@ -1161,31 +1489,6 @@ def _get_merged_param_data(net, parameter_layout_dict, param_name, param_data, i
     return param_data
 
 
-def _fill_param_into_net(net, parameter_list):
-    """
-    Fills parameter_list into net.
-
-    Args:
-        net (Cell): train network.
-        parameter_list (list): parameters list from ge callback.
-    """
-    parameter_dict = {}
-    for each_param in parameter_list:
-        param_name = each_param["name"]
-        if isinstance(each_param["data"], Parameter):
-            each_param["data"].init_data()
-        np_val = each_param["data"].asnumpy()
-        if np_val.shape == (1,):
-            parameter_dict[param_name] = Parameter(np_val, name=param_name)
-        elif np_val.shape == ():
-            parameter_dict[param_name] = Parameter(Tensor(np_val.tolist(), mstype.pytype_to_dtype(np_val.dtype)),
-                                                   name=param_name)
-        else:
-            parameter_dict[param_name] = Parameter(Tensor(np_val), name=param_name)
-
-    load_param_into_net(net, parameter_dict, strict_load=True)
-
-
 def export(net, *inputs, file_name, file_format, **kwargs):
     """
     Export the MindSpore network into an offline model in the specified format.
@@ -1193,9 +1496,9 @@ 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.
-        3. Exporting functions decorated with 'jit' to mindir format is supported.
-        4. When exporting a function decorated with 'jit', the function should not involve class properties in
-           calculations.
+        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.
 
     Args:
         net (Union[Cell, function]): MindSpore network.
@@ -1231,17 +1534,20 @@ def export(net, *inputs, file_name, file_format, **kwargs):
 
               - type (str): The type of obfuscation, only 'dynamic' is supported until now.
               - obf_ratio (float, str): The ratio of nodes in original model that would be obfuscated. `obf_ratio`
-                should be in range of (0, 1] or in ["small", "medium", "large"].
+                should be in range of (0, 1] or in ["small", "medium", "large"]. "small", "medium" and "large" are
+                correspond to 0.1, 0.3, and 0.6 respectively.
               - 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. This
-                function needs to ensure that its result is constant for any input. Users can refer to opaque
+                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>`_).
+                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): The random seed used for determine the distribution of confusion branches and the
-                weight confusion coefficient, which should be in (0, 9223372036854775807]. If `obf_random_seed` is set,
-                then it should be passed to :class:`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.
+              - 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
+                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.
 
@@ -1250,10 +1556,19 @@ def export(net, *inputs, file_name, file_format, **kwargs):
         >>> import numpy as np
         >>> from mindspore import Tensor
         >>>
-        >>> net = LeNet()
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/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')
+
+    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>`_
     """
+    old_ms_jit_value = context.get_context("jit_syntax_level")
+    context.set_context(jit_syntax_level=mindspore.STRICT)
+
     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}.")
@@ -1282,6 +1597,47 @@ def export(net, *inputs, file_name, file_format, **kwargs):
         kwargs['enc_key'], kwargs['enc_mode'] = _check_key_mode_type(file_format, **kwargs)
     _export(net, file_name, file_format, *inputs, **kwargs)
 
+    context.set_context(jit_syntax_level=old_ms_jit_value)
+
+
+def _get_funcgraph(net, *inputs):
+    """
+    Compile the MindSpore network and get FuncGraph.
+
+    Arg:
+        net (Union[Cell, function]): MindSpore network.
+        inputs (Union[Tensor, Dataset, List, Tuple, Number, Bool]): It represents the inputs
+             of the `net`, if the network has multiple inputs, set them together. While its type is Dataset,
+             it represents the preprocess behavior of the `net`, data preprocess operations will be serialized.
+             In second situation, you should adjust batch size of dataset script manually which will impact on
+             the batch size of 'net' input. Only supports parse "image" column from dataset currently.
+
+    Returns:
+        FuncGraph, a mindspore._c_expression.FuncGraph obj.
+
+    Raises:
+        ValueError: input `net` is not a nn.Cell.
+
+    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
+        >>> net = LeNet5()
+        >>> input_tensor = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32))
+        >>> ms.get_funcgraph(net, input_tensor)
+
+    """
+    if not isinstance(net, nn.Cell):
+        raise ValueError(f"For get_funcgraph's parameter 'net', currently only support Cell right now.")
+    phase_name = "lite_infer_predict" if _is_in_auto_parallel_mode() else "lite_infer_get_func_graph"
+    graph_id, _ = _executor.compile(net, *inputs, phase=phase_name, do_convert=False)
+    # pylint: disable=protected-access
+    func_graph = _executor._get_func_graph(net, graph_id)
+    return func_graph
+
 
 def _export(net, file_name, file_format, *inputs, **kwargs):
     """
@@ -1290,7 +1646,6 @@ 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 file_format == 'AIR':
         _save_air(net, file_name, *inputs, **kwargs)
     elif file_format == 'ONNX':
@@ -1454,7 +1809,7 @@ def _split_save(net_dict, model, file_name, is_encrypt, **kwargs):
         for param_proto in model.graph.parameter:
             name = param_proto.name[param_proto.name.find(":") + 1:]
             param = net_dict[name]
-            raw_data = param.data.asnumpy().tobytes()
+            raw_data = param.data.get_bytes()
             data_length = len(raw_data)
             append_size = 0
             if data_length % 64 != 0:
@@ -1508,7 +1863,7 @@ def _msfunc_info(net, *inputs):
 
 def _cell_info(net, incremental, *inputs):
     """Get mindir stream and net dict of cell"""
-    phase_name = "predict" if _is_in_auto_parallel_mode() else "export.mindir"
+    phase_name = "export.mindir"
     graph_id, _ = _executor.compile(net, *inputs, phase=phase_name, do_convert=False)
     # pylint: disable=protected-access
     mindir_stream = _executor._get_func_graph_proto(net, graph_id, 'mind_ir', incremental=incremental)
@@ -1581,7 +1936,7 @@ def _save_mindir_together(net_dict, model, file_name, is_encrypt, **kwargs):
     for param_proto in model.graph.parameter:
         param_name = param_proto.name[param_proto.name.find(":") + 1:]
         if param_name in net_dict.keys():
-            param_data = net_dict[param_name].data.asnumpy().tobytes()
+            param_data = net_dict[param_name].data.get_bytes()
             param_proto.raw_data = param_data
         else:
             raise ValueError("The parameter '{}' is not belongs to any cell,"
@@ -1591,10 +1946,10 @@ def _save_mindir_together(net_dict, model, file_name, is_encrypt, **kwargs):
         map_param_name = map_param_proto.name[map_param_proto.name.find(":") + 1:]
         if map_param_name in net_dict.keys():
             map_parameter = net_dict[map_param_name]
-            key_nparr, value_nparr, status_nparr = map_parameter.export_data(incremental)
-            map_param_proto.key_tensor.raw_data = key_nparr.tobytes()
-            map_param_proto.value_tensor.raw_data = value_nparr.tobytes()
-            map_param_proto.status_tensor.raw_data = status_nparr.tobytes()
+            key_bytes, value_bytes, status_bytes = map_parameter.export_bytes(incremental)
+            map_param_proto.key_tensor.raw_data = key_bytes
+            map_param_proto.value_tensor.raw_data = value_bytes
+            map_param_proto.status_tensor.raw_data = status_bytes
         else:
             raise ValueError("The map_parameter '{}' is not belongs to any cell,"
                              "the data of parameter cannot be exported.".format(map_param_proto.name))
@@ -1625,7 +1980,7 @@ def _save_together(net_dict, model):
     for param_proto in model.graph.parameter:
         name = param_proto.name[param_proto.name.find(":") + 1:]
         if name in net_dict.keys():
-            data_total += sys.getsizeof(net_dict[name].data.asnumpy().tobytes()) / 1024
+            data_total += sys.getsizeof(net_dict[name].data.get_bytes()) / 1024
         else:
             raise ValueError("The parameter '{}' is not belongs to any cell,"
                              "the data of parameter cannot be exported.".format(param_proto.name))
@@ -1656,7 +2011,7 @@ def _save_dataset_to_mindir(model, dataset):
 
 def parse_print(print_file_name):
     """
-    Parse data file generated by mindspore.ops.Print.
+    Parse data file generated by :class:`mindspore.ops.Print`.
 
     Args:
         print_file_name (str): The file name needs to be parsed.
@@ -1671,9 +2026,7 @@ def parse_print(print_file_name):
     Examples:
         >>> import numpy as np
         >>> import mindspore as ms
-        >>> import mindspore.ops as ops
-        >>> from mindspore import nn
-        >>> from mindspore import Tensor
+        >>> from mindspore import nn, Tensor, ops
         >>> ms.set_context(mode=ms.GRAPH_MODE, print_file_path='log.data')
         >>> class PrintInputTensor(nn.Cell):
         ...         def __init__(self):
@@ -1688,8 +2041,7 @@ def parse_print(print_file_name):
         >>> net = PrintInputTensor()
         >>> net(input_pra)
         >>>
-        >>> import mindspore
-        >>> data = mindspore.parse_print('./log.data')
+        >>> data = ms.parse_print('./log.data')
         >>> print(data)
         ['print:', Tensor(shape=[2, 4], dtype=Float32, value=
         [[ 1.00000000e+00,  2.00000000e+00,  3.00000000e+00,  4.00000000e+00],
@@ -1836,8 +2188,8 @@ def _merge_param_with_strategy(sliced_data, parameter_name, strategy, is_even):
 def restore_group_info_list(group_info_file_name):
     """
     Build rank list, the checkpoint of ranks in the rank list has the same contents with the local rank
-    who saves the group_info_file_name. To save the group info file, please export GROUP_INFO_FILE environment variables
-    like "export GROUP_INFO_FILE=/data/group_info.pb".
+    who saves the `group_info_file_name`. To save the group info file, please export GROUP_INFO_FIL
+    environment variables like "export GROUP_INFO_FILE=/data/group_info.pb".
 
     Args:
         group_info_file_name (str): Name of group information file.
@@ -1847,10 +2199,11 @@ def restore_group_info_list(group_info_file_name):
 
     Raises:
         ValueError: group information file is incorrect.
-        TypeError: group_info_file_name is not str.
+        TypeError: `group_info_file_name` is not str.
 
     Examples:
-        >>> restore_list = restore_group_info_list("./group_info.pb")
+        >>> import mindspore as ms
+        >>> ms.restore_list = restore_group_info_list("./group_info.pb")
     """
     if not isinstance(group_info_file_name, str):
         raise TypeError(f"For 'restore_group_info_list', the argument 'group_info_file_name' should be str, "
@@ -1868,9 +2221,6 @@ def restore_group_info_list(group_info_file_name):
 def build_searched_strategy(strategy_filename):
     """
     Build strategy of every parameter in network. Used in the case of distributed inference.
-    For details of it, please check:
-    `Saving and Loading Models in Hybrid Parallel Mode
-    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/save_load.html>`_.
 
     Args:
         strategy_filename (str): Name of strategy file.
@@ -1880,10 +2230,11 @@ def build_searched_strategy(strategy_filename):
 
     Raises:
         ValueError: Strategy file is incorrect.
-        TypeError: strategy_filename is not a string.
+        TypeError: `strategy_filename` is not a string.
 
     Examples:
-        >>> strategy = build_searched_strategy("./strategy_train.ckpt")
+        >>> import mindspore as ms
+        >>> strategy = ms.build_searched_strategy("./strategy_train.ckpt")
     """
     return _build_searched_strategy(strategy_filename)
 
@@ -1891,14 +2242,12 @@ def build_searched_strategy(strategy_filename):
 def merge_sliced_parameter(sliced_parameters, strategy=None):
     """
     Merge parameter slices into one parameter. Used in the case of distributed inference.
-    For details of it, please check:
-    `<https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/save_load.html>`_.
 
     Args:
         sliced_parameters (list[Parameter]): Parameter slices in order of rank id.
         strategy (Optional[dict]): Parameter slice strategy, whose key is parameter name and
             value is slice strategy of this parameter. If strategy is None, just merge
-            parameter slices in 0 axis order. Default: None.
+            parameter slices in 0 axis order. Default: ``None``.
 
     Returns:
         Parameter, the merged parameter which has the whole data.
@@ -1986,32 +2335,128 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
                                 train_strategy_filename=None, strict_load=False, dec_key=None, dec_mode='AES-GCM'):
     """
     Load checkpoint into net for distributed predication. Used in the case of distributed inference.
-    For details of distributed inference, please check:
-    `Distributed Inference
-    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/distributed_inference.html>`_ .
 
     Args:
         network (Cell): Network for distributed predication.
         checkpoint_filenames (list[str]): The name of Checkpoint files in order of rank id.
         predict_strategy (dict): Strategy of predication process. It means that using one device to predict
-                                 when setting predict_strategy as None. Default: None.
+                                 when setting predict_strategy as None. Default: ``None`` .
         train_strategy_filename (str): The filename of training strategy protocol buffer file.
                                        When train_strategy_filename is None, the training strategy file will be
                                        obtained from context.get_auto_parallel_context("strategy_ckpt_load_file").
                                        Therefore, the training strategy file needs to be specified
-                                       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
+                                       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
-                            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'.
+                            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'`` .
 
     Raises:
         TypeError: The type of inputs do not match the requirements.
         ValueError: Failed to load checkpoint into net.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            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>`_
+            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>`_ .
+
+            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>`_ .
+
+        >>> import os
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> import mindspore.dataset as ds
+        >>> from mindspore import nn, ops, train
+        >>> from mindspore.communication import init
+        >>>
+        >>> step_per_epoch = 4
+        >>> device_num = 8
+        >>>
+        >>> # Define the network structure.
+        >>> class Net(nn.Cell):
+        ...     def __init__(self, matmul_size, strategy=None):
+        ...         super().__init__()
+        ...         matmul_np = np.full(matmul_size, 0.5, dtype=np.float32)
+        ...         self.matmul_weight = ms.Parameter(ms.Tensor(matmul_np))
+        ...         self.matmul = ops.MatMul()
+        ...         self.neg = ops.Neg()
+        ...         if strategy is not None:
+        ...             self.matmul.shard(strategy)
+        ...
+        ...     def construct(self, inputs):
+        ...         x = self.matmul(inputs, self.matmul_weight)
+        ...         x = self.neg(x)
+        ...         return x
+        >>>
+        >>> # Create dataset.
+        >>> def get_dataset(*inputs):
+        ...     def generate():
+        ...         for _ in range(step_per_epoch):
+        ...             yield inputs
+        ...     return generate
+        >>>
+        >>> # Train network and save distributed checkpoint.
+        >>> def train_net():
+        ...     ms.set_context(mode=ms.GRAPH_MODE)
+        ...     init()
+        ...     np.random.seed(1)
+        ...     input_data = np.random.rand(16, 96).astype(np.float32)
+        ...     label_data = np.random.rand(16, 16).astype(np.float32)
+        ...     fake_dataset = get_dataset(input_data, label_data)
+        ...     dataset = ds.GeneratorDataset(fake_dataset, ["input", "label"])
+        ...
+        ...     # Set parallel strategy.
+        ...     strategy = ((1, 4), (4, 1))
+        ...     ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL, device_num=device_num,
+        ...                                  strategy_ckpt_save_file="./train_strategy.ckpt")
+        ...     network = Net(matmul_size=(96, 16), strategy=strategy)
+        ...     net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
+        ...     net_loss = nn.SoftmaxCrossEntropyWithLogits(reduction="mean")
+        ...     model = ms.Model(network=network, loss_fn=net_loss, optimizer=net_opt)
+        ...     ckpt_config = train.CheckpointConfig(keep_checkpoint_max=1, integrated_save=False)
+        ...     global_rank_id = int(os.getenv("RANK_ID"))
+        ...     ckpt_path = "./rank_{}_ckpt".format(global_rank_id)
+        ...     ckpt_callback = train.ModelCheckpoint(prefix="parallel", directory=ckpt_path, config=ckpt_config)
+        ...     model.train(epoch=2, train_dataset=dataset, callbacks=[ckpt_callback], dataset_sink_mode=False)
+        ...     ms.reset_auto_parallel_context()
+        >>>
+        >>> # Load distributed checkpoint and test.
+        >>> def load_model():
+        ...     ms.set_context(mode=ms.GRAPH_MODE)
+        ...     init()
+        ...     ms.set_auto_parallel_context(full_batch=True, parallel_mode="semi_auto_parallel",
+        ...                                  strategy_ckpt_load_file="./train_strategy.ckpt", device_num=device_num)
+        ...     predict_data = ms.Tensor(np.random.randn(128, 96).astype(np.float32))
+        ...     network = Net(matmul_size=(96, 16))
+        ...     model = ms.Model(network)
+        ...     predict_layout = model.infer_predict_layout(ms.Tensor(predict_data))
+        ...     ckpt_file_list = ["./rank_{}_ckpt/parallel-2_4.ckpt".format(i) for i in range(0, device_num)]
+        ...     ms.load_distributed_checkpoint(network, ckpt_file_list, predict_layout)
+        ...     predict_result = model.predict(predict_data)
+        ...     print(predict_result)
+        >>>
+        >>> train_net()
+        >>> load_model()
+        [[-7.3259363 -7.497216  -7.398196  ... -7.374962  -7.204874  -7.234935 ]
+        [ 3.362938   3.3535435  3.3832688 ...  3.4263954  3.279045   3.3202887]
+        ...
+        [ 1.6067538  1.6244187  1.5384722 ...  1.5449994  1.6195512  1.6176052]]
     """
     network = Validator.check_isinstance("network", network, nn.Cell)
     _check_checkpoint_file(checkpoint_filenames)
@@ -2127,6 +2572,11 @@ def async_ckpt_thread_status():
     Returns:
         bool, True, Asynchronous save checkpoint thread is running.
         False, Asynchronous save checkpoint thread is not executing.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.async_ckpt_thread_status()
+        False
     """
     thr_list = threading.enumerate()
     return True in [ele.getName() == "asyn_save_ckpt" for ele in thr_list]
@@ -2184,7 +2634,8 @@ def _merge_and_split(sliced_params, train_strategy, predict_strategy):
         return merged_param
     param_name = merged_param.name
     tensor_layout = predict_strategy[param_name]
-    split_tensor = _load_tensor(merged_param.data, tensor_layout[0], tensor_layout[1])
+    rank = get_rank()
+    split_tensor = _load_tensor(merged_param.data, tensor_layout[0], tensor_layout[1], rank)
     requires_grad = merged_param.requires_grad
     layerwise_parallel = merged_param.layerwise_parallel
     split_param = Parameter(split_tensor, param_name, requires_grad, layerwise_parallel)
@@ -2268,7 +2719,8 @@ def convert_model(mindir_file, convert_file, file_format):
         ValueError: If the parameter `file_format` is not "ONNX".
 
     Examples:
-        >>> convert_model("lenet.mindir", "lenet.onnx", "ONNX")
+        >>> import mindspore as ms
+        >>> ms.convert_model("lenet.mindir", "lenet.onnx", "ONNX")
     """
     Validator.check_file_name_by_regular(mindir_file)
     Validator.check_file_name_by_regular(convert_file)