mindspore 2.1.0__cp37-cp37m-manylinux1_x86_64.whl → 2.2.11__cp37-cp37m-manylinux1_x86_64.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
 
@@ -59,19 +59,23 @@ 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, \
-    split_mindir
+    split_mindir, split_dynamic_mindir
 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}
+                     "Bool": mstype.bool_, "str": mstype.string, "BFloat16": mstype.bfloat16}
 
 tensor_to_np_type = {"Int8": np.int8, "UInt8": np.uint8, "Int16": np.int16, "UInt16": np.uint16,
                      "Int32": np.int32, "UInt32": np.uint32, "Int64": np.int64, "UInt64": np.uint64,
-                     "Float16": np.float16, "Float32": np.float32, "Float64": np.float64, "Bool": np.bool_, "str": "U"}
+                     "Float16": np.float16, "Float32": np.float32, "Float64": np.float64, "Bool": np.bool_, "str": "U",
+                     "BFloat16": np.float32}
 
 np_type_convert = {"int32": np.int32, "float32": np.float32, "float16": np.float16, "float64": np.float64}
 
@@ -79,7 +83,7 @@ mindir_to_tensor_type = {1: mstype.float32, 2: mstype.uint8, 3: mstype.int8, 4:
                          5: mstype.int16, 6: mstype.int32, 7: mstype.int64, 10: mstype.float16,
                          11: mstype.float64, 12: mstype.uint32, 13: mstype.uint64}
 
-_ckpt_mutex = Lock()
+_ckpt_mutex = RLock()
 
 # unit is KB
 SLICE_SIZE = 512 * 1024
@@ -89,6 +93,8 @@ PARAMETER_SPLIT_SIZE = 1024 * 1024 * 1024
 ENCRYPT_BLOCK_SIZE = 64 * 1024
 INT_64_MAX = 9223372036854775807
 
+cpu_cast = Cast().set_device("CPU")
+
 
 def _special_process_par(par, new_par):
     """
@@ -105,7 +111,11 @@ def _special_process_par(par, new_par):
         if new_par.data.shape[par_shape_len + i] != 1:
             return False
 
-    new_val = new_par.data.asnumpy()
+    if new_par.data.dtype == mstype.bfloat16:
+        new_val = cpu_cast(new_par.data, mstype.float32).asnumpy()
+    else:
+        new_val = new_par.data.asnumpy()
+
     new_val = new_val.reshape(par.data.shape)
     par.set_data(Tensor(new_val, par.data.dtype))
     return True
@@ -126,7 +136,10 @@ 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)
+                if new_param.data.dtype == mstype.bfloat16:
+                    new_tensor = cpu_cast(new_param.data, param.data.dtype)
+                else:
+                    new_tensor = Tensor(new_param.data.asnumpy(), param.data.dtype)
                 param.set_data(new_tensor, param.sliced)
                 return
 
@@ -229,10 +242,16 @@ 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:]
-                        new_value[2] = value[3].asnumpy().reshape(-1)
+                        if value[3].dtype == mstype.bfloat16:
+                            new_value[2] = cpu_cast(value[3], mstype.float32).asnumpy().reshape(-1)
+                        else:
+                            new_value[2] = value[3].asnumpy().reshape(-1)
                         _write_parameter_data(name, new_value, f, enc_key, plain_data)
                         _offload_if_config(value[3])
                         continue
+                    if value[0] == "BFloat16_tensor":
+                        _write_bfloat16_data(name, value, f, enc_key, plain_data)
+                        continue
                     if isinstance(value[2], Tensor):
                         _write_hugeparameter(name, value, f)
                         continue
@@ -267,6 +286,21 @@ 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):
     """Write parameter data into protobuf file."""
     data_size = value[2].nbytes / 1024
@@ -333,8 +367,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,"
@@ -351,14 +385,15 @@ 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):
-    """
+    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`` .
@@ -370,16 +405,14 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                         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 return value is a bool.
+                                 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.
 
-            - incremental (bool): Whether export checkpoint for MapParameter incrementally.
-
     Raises:
-        TypeError: If the parameter `save_obj` is not `nn.Cell` or list 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.
 
@@ -387,17 +420,27 @@ 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.1/docs/mindspore/code/lenet.py
+        >>> # 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_dict = ms.load_checkpoint("./lenet.ckpt")
-        >>> print(param_dict)
+        ...                    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.1/beginner/save_load.html#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)
@@ -408,70 +451,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
     map_param_inc = kwargs.get('incremental', False)
     logger.info("Execute the process of saving checkpoint files.")
 
-    if isinstance(save_obj, nn.Cell):
-        if save_obj.ge_init and not save_obj.ge_sync_data:
-            from mindspore.train.callback._callback import set_cur_net
-            set_cur_net(save_obj)
-            save_obj.exec_checkpoint_graph()
-        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 = 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.
-            if is_graph_mode and _is_in_auto_parallel_mode() and (not_sliced or param.has_init):
-                continue
-            if choice_func is not None and not choice_func(param.name):
-                continue
-            param_dict[param.name] = param
-        param_list = []
-        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"]
-                param_data.append(value.data)
-                param_data.append(value.param_info.origin_shape)
-                param_data.append(str(value.dtype))
-                param_data.append(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)
-        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 = []
@@ -479,7 +459,7 @@ 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:
@@ -499,6 +479,10 @@ 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])
@@ -508,6 +492,15 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
             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)
@@ -517,7 +510,10 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                 data_list[key].append(dims)
                 tensor_type = str(param["data"].dtype)
                 data_list[key].append(tensor_type)
-                data = param["data"].asnumpy().reshape(-1)
+                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:
@@ -530,6 +526,130 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
     logger.info("Saving checkpoint process is finished.")
 
 
+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 _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)
+        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())
+
+            # 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 = []
@@ -585,7 +705,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.1/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`.
@@ -615,7 +735,7 @@ def load(file_name, **kwargs):
 
     Tutorial Examples:
         - `Saving and Loading the Model - Saving and Loading MindIR
-          <https://mindspore.cn/tutorials/en/r2.1/beginner/save_load.html#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 "
@@ -656,7 +776,7 @@ def load(file_name, **kwargs):
     return graph
 
 
-def export_split_mindir(file_name):
+def export_split_mindir(file_name, device_num=8, rank_id=0, dynamic=True, sapp=False):
     """
     Auto Split MindIR.
 
@@ -664,6 +784,10 @@ def export_split_mindir(file_name):
 
     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.
@@ -671,11 +795,9 @@ def export_split_mindir(file_name):
 
     Examples:
         >>> import mindspore as ms
-        >>> from mindspore.communication import init
         >>> context.set_context(mode=context.GRAPH_MODE)
         >>>
-        >>> init(backend_name="hccl")
-        >>> ms.export_split_mindir("net.mindir")
+        >>> ms.export_split_mindir("net.mindir", device_num=8, rank_id=0)
 
     """
     if not isinstance(file_name, str):
@@ -690,8 +812,11 @@ def export_split_mindir(file_name):
     file_name = os.path.abspath(file_name)
 
     logger.info("Execute the process of export and split mindir.")
-
-    graph = split_mindir(file_name)
+    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):
@@ -779,17 +904,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.
 
@@ -928,27 +1056,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.1/beginner/save_load.html#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)
@@ -979,8 +1107,7 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
                     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)
+                _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]
@@ -992,6 +1119,13 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
             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 \
@@ -1024,8 +1158,12 @@ 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
 
@@ -1061,7 +1199,7 @@ def _load_map_parameter(checkpoint_list, element, element_id, map_data_list,
 
 
 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)))
@@ -1175,7 +1313,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.1/docs/mindspore/code/lenet.py
+        >>> # 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")
@@ -1185,7 +1323,7 @@ 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.1/beginner/save_load.html#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.")
@@ -1219,6 +1357,9 @@ def load_param_into_net(net, parameter_dict, strict_load=False):
             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)
@@ -1243,6 +1384,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 = ""
@@ -1350,9 +1557,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.
@@ -1388,17 +1595,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.
 
@@ -1408,14 +1618,14 @@ def export(net, *inputs, file_name, file_format, **kwargs):
         >>> from mindspore import Tensor
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # 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.1/beginner/save_load.html#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)
@@ -1475,7 +1685,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.1/docs/mindspore/code/lenet.py
+        >>> # 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)
@@ -1657,10 +1867,17 @@ 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
+        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])
+        for pairs in names:
+            name = pairs[0]
+            param_proto = pairs[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:
@@ -1678,6 +1895,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}")
 
         graph_file_name = os.path.join(dirname, file_prefix + "_graph.mindir")
         if os.path.exists(graph_file_name):
@@ -1787,7 +2006,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,"
@@ -1797,10 +2016,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))
@@ -1831,7 +2050,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))
@@ -1862,7 +2081,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.
@@ -2039,8 +2258,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.
@@ -2050,7 +2269,7 @@ 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:
         >>> import mindspore as ms
@@ -2072,9 +2291,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.1/parallel/save_load.html>`_.
 
     Args:
         strategy_filename (str): Name of strategy file.
@@ -2096,8 +2312,6 @@ 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.1/parallel/save_load.html>`_.
 
     Args:
         sliced_parameters (list[Parameter]): Parameter slices in order of rank id.
@@ -2171,7 +2385,12 @@ def merge_sliced_parameter(sliced_parameters, strategy=None):
 
     layerwise_parallel = sliced_parameters[0].layerwise_parallel
     requires_grad = sliced_parameters[0].requires_grad
-    sliced_data = [parameter.data.asnumpy() for parameter in sliced_parameters]
+    sliced_data = []
+    for parameter in sliced_parameters:
+        if parameter.data.dtype == mstype.bfloat16:
+            sliced_data.append(cpu_cast(parameter.data, mstype.float32).asnumpy())
+        else:
+            sliced_data.append(parameter.data.asnumpy())
 
     if not strategy:
         merged_tensor = Tensor(np.concatenate(sliced_data))
@@ -2191,9 +2410,6 @@ 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.1/parallel/distributed_inference.html>`_ .
 
     Args:
         network (Cell): Network for distributed predication.
@@ -2218,6 +2434,104 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
     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)
@@ -2282,7 +2596,11 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
                 param_index = list(set(param_index))
                 param_index.sort()
                 for rank_num in param_index:
-                    param_stride.append(param_total_dict[param.name][rank_num].data.asnumpy())
+                    if param_total_dict[param.name][rank_num].data.dtype == mstype.bfloat16:
+                        param_stride.append(
+                            cpu_cast(param_total_dict[param.name][rank_num].data, mstype.float32).asnumpy())
+                    else:
+                        param_stride.append(param_total_dict[param.name][rank_num].data.asnumpy())
 
                 sliced_param = Parameter(Tensor(np.concatenate(param_stride)), name=param.name)
             else:
@@ -2297,7 +2615,10 @@ def load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy=
             split_param = _merge_and_split(sliced_params, _param_unique_strategy, predict_strategy)
         opt_shard_group = predict_strategy[param.name][5] if predict_strategy else None
         if opt_shard_group:
-            data = split_param.data.asnumpy()
+            if split_param.data.dtype == mstype.bfloat16:
+                data = cpu_cast(split_param.data, mstype.float32).asnumpy()
+            else:
+                data = split_param.data.asnumpy()
             rank = get_rank(opt_shard_group)
             size = get_group_size(opt_shard_group)
             try:
@@ -2395,10 +2716,15 @@ 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)
+    data_type = merged_param.data.dtype
+    if data_type == 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)
     return split_param
 
 
@@ -2407,7 +2733,7 @@ def _calculation_net_size(net):
     data_total = 0
     net_dict = net.parameters_dict()
     for name in net_dict:
-        data_total += sys.getsizeof(net_dict[name].data.asnumpy().tobytes()) / 1024
+        data_total += sys.getsizeof(net_dict[name].data.get_bytes()) / 1024
 
     return data_total