mindspore 2.1.0__cp38-cp38-win_amd64.whl → 2.2.11__cp38-cp38-win_amd64.whl

mindspore/parallel/_offload_context.py

@@ -33,6 +33,7 @@ class _OffloadConfig:
     OFFLOAD_PARAM = "offload_param"
     OFFLOAD_PATH = "offload_path"
     OFFLOAD_CPU_SIZE = "offload_cpu_size"
+    OFFLOAD_CHECKPOINT = "offload_checkpoint"
     OFFLOAD_DISK_SIZE = "offload_disk_size"
     ENABLE_AIO = "enable_aio"
     AIO_BLOCK_SIZE = "aio_block_size"
@@ -84,6 +85,16 @@ class _OffloadContext:
         Validator.check_string(offload_param.lower(), ["cpu", "disk"], "offload_param", "set_offload_param")
         self._context_handle.set_offload_param(offload_param.lower())
 
+    def set_offload_checkpoint(self, offload_checkpoint):
+        """Set offload_checkpoint"""
+        if not isinstance(offload_checkpoint, str):
+            raise TypeError("For 'set_offload_checkpoint', "
+                            "the argument 'offload_checkpoint' must be str, but got the type : {}."
+                            .format(type(offload_checkpoint)))
+        Validator.check_string(offload_checkpoint.lower(), ["cpu", "disk"], "offload_checkpoint",
+                               "set_offload_checkpoint")
+        self._context_handle.set_offload_checkpoint(offload_checkpoint.lower())
+
     def set_offload_path(self, offload_path):
         """Set offload_path"""
         if not isinstance(offload_path, str):
@@ -194,7 +205,8 @@ class _OffloadContext:
                                    _OffloadConfig.HBM_RATIO, _OffloadConfig.OFFLOAD_CPU_SIZE,
                                    _OffloadConfig.OFFLOAD_DISK_SIZE, _OffloadConfig.ENABLE_AIO,
                                    _OffloadConfig.AIO_BLOCK_SIZE, _OffloadConfig.AIO_QUEUE_DEPTH,
-                                   _OffloadConfig.ENABLE_PINNED_MEM, _OffloadConfig.AUTO_OFFLOAD]:
+                                   _OffloadConfig.ENABLE_PINNED_MEM, _OffloadConfig.AUTO_OFFLOAD,
+                                   _OffloadConfig.OFFLOAD_CHECKPOINT]:
                 unknown_config.append(config_name)
 
             if unknown_config:
@@ -220,7 +232,8 @@ class _OffloadContext:
             _OffloadConfig.AUTO_OFFLOAD: self._context_handle.auto_offload(),
             _OffloadConfig.HOST_MEM_BLOCk_SIZE: self._context_handle.host_mem_block_size(),
             _OffloadConfig.CPU_RATIO: self._context_handle.cpu_ratio(),
-            _OffloadConfig.HBM_RATIO: self._context_handle.hbm_ratio()
+            _OffloadConfig.HBM_RATIO: self._context_handle.hbm_ratio(),
+            _OffloadConfig.OFFLOAD_CHECKPOINT: self._context_handle.offload_checkpoint()
         }
         return offload_config
 
@@ -257,5 +270,6 @@ _set_offload_context_func_map = {
     _OffloadConfig.AUTO_OFFLOAD: offload_context().set_auto_offload,
     _OffloadConfig.HOST_MEM_BLOCk_SIZE: offload_context().set_host_mem_block_size,
     _OffloadConfig.CPU_RATIO: offload_context().set_cpu_ratio,
-    _OffloadConfig.HBM_RATIO: offload_context().set_hbm_ratio
+    _OffloadConfig.HBM_RATIO: offload_context().set_hbm_ratio,
+    _OffloadConfig.OFFLOAD_CHECKPOINT: offload_context().set_offload_checkpoint
 }

mindspore/parallel/_parallel_serialization.py

@@ -330,12 +330,13 @@ def _rank_list_for_transform_parallel_checkpoint(rank_id, src_strategy_list, dst
         device_list = list(range(0, np.prod(from_tensor_layout[0])))
         param_rank_list = _get_needed_rank_list_by_layouts(from_tensor_layout, to_tensor_layout, device_list, rank_id)
         param_rank_list_new = [rank % from_device_num for rank in param_rank_list]
-        param_rank_list_new = set(param_rank_list_new)
-        result_list.update(param_rank_list_new)
+        param_rank_set_new = set(param_rank_list_new)
+        result_list.update(param_rank_set_new)
     return list(result_list)
 
 
-def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, src_strategy_list, dst_strategy_list):
+def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, src_strategy_list,
+                                   dst_strategy_list, param_type_dict):
     """
     Transform model parallel dimension for distributed checkpoint files.
     """
@@ -397,15 +398,21 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
         transform_tensor = ms.Tensor(param_total_dict[param_name][rank_id % device_num])
         requires_grad = param_attr_dict[param_name][rank_id % device_num][0]
         layerwise_parallel = param_attr_dict[param_name][rank_id % device_num][1]
-        transform_param_dict[param_name] = ms.Parameter(transform_tensor, param_name, requires_grad, layerwise_parallel)
+        transform_para = ms.Parameter(transform_tensor, param_name, requires_grad, layerwise_parallel)
+        if param_type_dict[param_name][rank_id % device_num] == "BFloat16":
+            transform_para.set_dtype(ms.bfloat16)
+        transform_param_dict[param_name] = transform_para
 
     # Handle those parameter like learning_rate, global_step which not in strategy_file.
     for param_name, _ in param_total_dict.items():
         if param_name not in transform_param_dict:
-            transform_param_dict[param_name] = ms.Parameter(
+            transform_para = ms.Parameter(
                 ms.Tensor(param_total_dict[param_name][rank_id % device_num]), param_name,
                 param_attr_dict[param_name][rank_id % device_num][0],
                 param_attr_dict[param_name][rank_id % device_num][1])
+            if param_type_dict[param_name][rank_id % device_num] == "BFloat16":
+                transform_para.set_dtype(ms.bfloat16)
+            transform_param_dict[param_name] = transform_para
 
     transform_param_list = [{"name": param_name, "data": param_data}
                             for param_name, param_data in transform_param_dict.items()]

mindspore/parallel/_ps_context.py

@@ -228,3 +228,15 @@ def _enable_distributed_mindrt():
     This method is used to distinguish from old distributed training mode.
     '''
     return ps_context().enable_distributed_mindrt()
+
+
+def _set_checkpoint_load_status(status):
+    return ps_context().set_checkpoint_load_status(status)
+
+
+def _store_warm_up_ptr_by_tensor(param_key, tensor):
+    return ps_context().store_warm_up_ptr_by_tensor(param_key, tensor)
+
+
+def _store_warm_up_ptr_by_tensor_list(param_key, key_tensor, value_tensor, status_tensor):
+    return ps_context().store_warm_up_ptr_by_tensor_list(param_key, key_tensor, value_tensor, status_tensor)

mindspore/parallel/_tensor.py

@@ -17,7 +17,7 @@ from __future__ import division
 from __future__ import absolute_import
 
 import numpy as np
-
+from mindspore.common import dtype as mstype
 from mindspore.common.tensor import Tensor
 from mindspore.communication.management import get_rank, get_group_size
 from mindspore._c_expression import TensorTransform
@@ -41,7 +41,7 @@ def _get_tensor_strategy(dev_mat, tensor_map):
         if dim == -1:
             tensor_strategy.append(1)
         else:
-            tensor_strategy.append(dev_mat[-dim-1])
+            tensor_strategy.append(dev_mat[-dim - 1])
     return tensor_strategy
 
 
@@ -198,7 +198,7 @@ def _get_slice_index(dev_mat, tensor_map, opt_shard_group):
     return tensor_slice_index
 
 
-def _load_tensor(tensor, dev_mat, tensor_map):
+def _load_tensor(tensor, dev_mat, tensor_map, rank_id=-1):
     """
     Get the tensor slice of the local device by the device matrix and the tensor map
 
@@ -216,16 +216,21 @@ def _load_tensor(tensor, dev_mat, tensor_map):
         >>> tensor_map = [1, -1]
         >>> tensor_slice = _load_tensor(tensor, dev_mat, tensor_map)
     """
-    rank = get_rank()
+    if rank_id == -1:
+        rank = get_rank()
+    else:
+        rank = rank_id
     tensor_strategy = _get_tensor_strategy(dev_mat, tensor_map)
     tensor_slice_index = _get_tensor_slice_index(dev_mat, tensor_strategy, tensor_map, rank)
+    if tensor.dtype == mstype.bfloat16:
+        tensor = tensor.float()
     np_tensor = tensor.asnumpy()
     np_tensor_list = _chunk_tensor_by_strategy(np_tensor, tensor_strategy)
     np_tensor_slice = np_tensor_list[int(tensor_slice_index)]
     return np_tensor_slice
 
 
-def _load_tensor_by_layout(tensor, layout):
+def _load_tensor_by_layout(tensor, layout, rank_id):
     """
     Load tensor by layout.
 
@@ -246,19 +251,19 @@ def _load_tensor_by_layout(tensor, layout):
         raise ValueError("The length of layout must be larger than 5! layout is {}".format(layout))
     dev_mat = layout[0]
     tensor_map = layout[1]
-    if len(tensor_map) == 0:
+    if not tensor_map:
         return tensor
     uniform_split = layout[4]
     group = layout[5]
     if uniform_split == 0:
         raise RuntimeError("The load tensor only support uniform split now")
-    tensor_slice = _load_tensor(tensor, dev_mat, tensor_map)
+    tensor_slice = _load_tensor(tensor, dev_mat, tensor_map, rank_id)
     if group:
         # get a totally shard tensor slice for parallel optimizer
         rank = get_rank(group)
         size = get_group_size(group)
         tensor_slice = np.split(tensor_slice, size)[rank]
-    return Tensor(tensor_slice)
+    return Tensor(tensor_slice, tensor.dtype)
 
 
 def _reshape_param_data(param_data, dev_mat, tensor_map):
@@ -315,7 +320,6 @@ def _reshape_param_data(param_data, dev_mat, tensor_map):
     return Tensor(tensor_slices_new[0])
 
 
-
 def _extract_layout_item(layout_item):
     dev_matrix = layout_item[0]
     tensor_map = layout_item[1]
@@ -541,6 +545,7 @@ def _check_operator(operator):
 
 def _apply_operator(operator_name):
     """apply transform operator"""
+
     def _apply_reshape_operator(numpy_data, reshape_op):
         """
         Apply reshape operator.
@@ -597,8 +602,8 @@ def _apply_operator(operator_name):
             raise ValueError("The slice operator information is wrong.")
         shape_size = len(slice_op[1]) // 3
         begin = slice_op[1][:shape_size]
-        end = slice_op[1][shape_size:shape_size*2]
-        stride = slice_op[1][shape_size*2:]
+        end = slice_op[1][shape_size:shape_size * 2]
+        stride = slice_op[1][shape_size * 2:]
         slice_index = []
         for begin_i, end_i, strides_i in zip(begin, end, stride):
             s = slice(begin_i, end_i, strides_i)
@@ -637,8 +642,8 @@ def _reshape_param_data_with_weight(param_data, dev_mat, field_size):
     for i in range(len(tensor_slices[0][0])):
         tensor_slices_new = np.array(tensor_slices[0][:, i]).reshape(field_size, -1)
         for j in range(1, device_count):
-            tensor_slices_new = np.concatenate((tensor_slices_new,\
-                                   np.array(tensor_slices[j][:, i]).reshape(field_size, -1)), axis=1)
+            tensor_slices_new = np.concatenate((tensor_slices_new, \
+                                                np.array(tensor_slices[j][:, i]).reshape(field_size, -1)), axis=1)
         tensor_slices_col.append(tensor_slices_new)
     new_tensor = np.array(tensor_slices_col[0]).reshape(-1, 1)
     for i in range(1, len(tensor_slices_col)):

mindspore/parallel/_transformer/layers.py

@@ -424,9 +424,11 @@ class _Linear(Cell):
         self.out_channels = out_channels
         if not (isinstance(activation, str) or activation is None or issubclass(activation, nn.Cell)):
             raise TypeError(f"For Linear cell, the activation should str type or nn.Cell type, but got {activation}.")
-        if isinstance(weight_init, Tensor) and (weight_init.ndim != 2 or weight_init.shape[0] != out_channels or
-                                                weight_init.shape[1] != in_channels):
-            raise ValueError("The shape of parameter 'weight_init' is error, please check shape of 'weight_init'.")
+
+        if isinstance(weight_init, Tensor):
+            if weight_init.ndim != 2 or weight_init.shape[0] != out_channels \
+                                or weight_init.shape[1] != in_channels:
+                raise ValueError("The shape of parameter 'weight_init' is error, please check shape of 'weight_init'.")
         weight_shape = [out_channels, in_channels] if transpose_b else [in_channels, out_channels]
         self.expert_num = expert_num
         self.outer_batch = outer_batch

mindspore/parallel/_transformer/loss.py

@@ -139,6 +139,7 @@ class _NLLLoss(Cell):
             self.add = P.Add().shard(((dp, mp), ()))
 
     def construct(self, softmax_result, one_hot_label):
+        """The forward of _NLLLoss"""
         log_softmax_result = self.log(self.add(softmax_result, self.eps_const))
         loss = self.mul(log_softmax_result, one_hot_label)
         loss_unsum = self.neg(loss)

mindspore/parallel/_transformer/moe.py

@@ -273,7 +273,7 @@ class MoE(Cell):
         if self.group_wise_a2a:
             # If capacity can't div by mp, pad for mp shard.
             if capacity % self.mp != 0:
-                pad_size = self.mp-(capacity % self.mp)
+                pad_size = self.mp - (capacity % self.mp)
             if pad_size != 0:
                 capacity += pad_size
                 pad_tensor = self.stride_slice_dp(expert_input, (0, 0, 0, 0),
@@ -330,7 +330,7 @@ class MoE(Cell):
         # Pad capacity for comp_comm_parallel_degree split.
         pad_size = 0
         if capacity % self.comp_comm_parallel_degree != 0:
-            pad_size = self.comp_comm_parallel_degree-(capacity % self.comp_comm_parallel_degree)
+            pad_size = self.comp_comm_parallel_degree - (capacity % self.comp_comm_parallel_degree)
             capacity += pad_size
             pad_tensor = self.stride_slice_dp(expert_input, (0, 0, 0, 0),
                                               (self.expert_dim, self.dp_group, pad_size, self.hidden_size),

mindspore/parallel/_transformer/op_parallel_config.py

@@ -147,10 +147,11 @@ class _PipeLineConfig(_Config):
             >>> config=_PipeLineConfig(pipeline_stage=1, micro_batch_num=1)
     """
 
-    def __init__(self, pipeline_stage=1, micro_batch_num=1):
+    def __init__(self, pipeline_stage=1, micro_batch_num=1, pipeline_segment=1):
         Validator.check_positive_int(pipeline_stage, "pipeline_stage")
         Validator.check_positive_int(micro_batch_num, "micro_batch_num")
         self.pipeline_stage = pipeline_stage
+        self.pipeline_segment = pipeline_segment
         self.micro_batch_num = micro_batch_num
 
     @property
@@ -163,6 +164,16 @@ class _PipeLineConfig(_Config):
         self._pipeline_stage = value
         context.set_auto_parallel_context(pipeline_stages=value)
 
+    @property
+    def pipeline_segment(self):
+        return self._pipeline_segment
+
+    @pipeline_segment.setter
+    def pipeline_segment(self, value):
+        Validator.check_positive_int(value, "pipeline_segment")
+        self._pipeline_segment = value
+        context.set_auto_parallel_context(pipeline_segments=value)
+
     @property
     def micro_batch_num(self):
         return self._micro_batch_num

mindspore/parallel/_transformer/transformer.py

@@ -226,7 +226,8 @@ class TransformerOpParallelConfig(_Config):
             >>> config=TransformerOpParallelConfig(data_parallel=1, model_parallel=1, recompute=recompute_config)
     """
 
-    def __init__(self, data_parallel=1, model_parallel=1, expert_parallel=1, pipeline_stage=1, micro_batch_num=1,
+    def __init__(self, data_parallel=1, model_parallel=1, expert_parallel=1, pipeline_stage=1, pipeline_segment=1,
+                 micro_batch_num=1,
                  recompute=default_transformer_recompute_config,
                  optimizer_shard=False, gradient_aggregation_group=4, vocab_emb_dp=True):
         self.recompute = recompute
@@ -234,7 +235,8 @@ class TransformerOpParallelConfig(_Config):
         self.gradient_aggregation_group = gradient_aggregation_group
         self._embed_dp_mp_config = EmbeddingOpParallelConfig(data_parallel=data_parallel, model_parallel=model_parallel,
                                                              vocab_emb_dp=vocab_emb_dp)
-        self._pp_config = _PipeLineConfig(pipeline_stage=pipeline_stage, micro_batch_num=micro_batch_num)
+        self._pp_config = _PipeLineConfig(pipeline_stage=pipeline_stage, micro_batch_num=micro_batch_num,
+                                          pipeline_segment=pipeline_segment)
         self._moe_config = MoEParallelConfig(data_parallel=data_parallel, model_parallel=model_parallel,
                                              expert_parallel=expert_parallel)
 
@@ -309,6 +311,14 @@ class TransformerOpParallelConfig(_Config):
     def pipeline_stage(self, value):
         self._pp_config.pipeline_stage = value
 
+    @property
+    def pipeline_segment(self):
+        return self._pp_config.pipeline_segment
+
+    @pipeline_segment.setter
+    def pipeline_segment(self, value):
+        self._pp_config.pipeline_segment = value
+
     @property
     def optimizer_shard(self):
         return self._optimizer_shard
@@ -429,6 +439,7 @@ class FeedForward(Cell):
             >>> print(output.shape)
             (2, 20, 15)
     """
+
     @_LogActionOnce(logger=logger, key='FeedForward',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
@@ -622,6 +633,7 @@ class AttentionMask(Cell):
               [1. 1. 1. 0]
               [0. 0. 0. 0]]]
     """
+
     @_LogActionOnce(logger=logger, key='AttentionMask',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(seq_length=Validator.check_positive_int,
@@ -710,6 +722,7 @@ class VocabEmbedding(Cell):
             >>> print(table.shape)
             (30, 30)
     """
+
     @_LogActionOnce(logger=logger, key='VocabEmbedding',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(vocab_size=Validator.check_positive_int,
@@ -866,6 +879,7 @@ class MultiHeadAttention(Cell):
             >>> print(past[1].shape)
             (2, 3, 20, 5)
     """
+
     @_LogActionOnce(logger=logger, key='MultiHeadAttention',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
@@ -1203,7 +1217,8 @@ class MultiHeadAttention(Cell):
     def _get_batch_size_from_query(self, query):
         r"""Get the batch size from query tensor"""
         # For the incremental prediction, the seq length for the input is 1.
-        if len(F.shape(query)) == 2 and ((self.use_past and self.is_first_iteration) or (not self.use_past)):
+        incr_infer = self.use_past and self.is_first_iteration
+        if len(F.shape(query)) == 2 and ((incr_infer) or (not self.use_past)):
             return F.shape(query)[0] // self.src_seq_length
         return F.shape(query)[0]
 
@@ -1459,6 +1474,7 @@ class TransformerEncoderLayer(Cell):
             >>> print(past[1].shape)
             (2, 2, 16, 4)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerEncoderLayer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
@@ -1848,6 +1864,7 @@ class TransformerDecoderLayer(Cell):
             >>> print(past[3].shape)
             (2, 2, 20, 32)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerDecoderLayer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
@@ -2379,6 +2396,7 @@ class TransformerEncoder(Cell):
             >>> print(past[0][1].shape)
             (2, 2, 16, 4)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerEncoder',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(batch_size=Validator.check_positive_int,
@@ -2613,6 +2631,7 @@ class TransformerDecoder(Cell):
             >>> print(past[0][3].shape)
             (2, 2, 20, 32)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerDecoder',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(batch_size=Validator.check_positive_int,
@@ -2882,6 +2901,7 @@ class Transformer(Cell):
             >>> print(de_past[0][3].shape)
             (2, 2, 20, 32)
     """
+
     @_LogActionOnce(logger=logger, key='Transformer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(batch_size=Validator.check_positive_int,

mindspore/parallel/_utils.py

@@ -100,13 +100,14 @@ def _slice_parameter(parameter, phase, layout):
         parameter.sliced = True
         return
     if not parameter.sliced:
-        new_tensor = _load_tensor_by_layout(parameter, layout)
+        rank = get_rank()
+        new_tensor = _load_tensor_by_layout(parameter, layout, rank)
         parameter.set_data(new_tensor, True)
 
 
-def _slice_tensor(tensor, layout):
+def _slice_tensor(tensor, layout, rank_id):
     """Slice python tensor obj according to the layout."""
-    new_tensor = _load_tensor_by_layout(tensor, layout)
+    new_tensor = _load_tensor_by_layout(tensor, layout, rank_id)
     return new_tensor
 
 
@@ -136,14 +137,17 @@ def _to_full_shapes(shapes, device_num):
                                  "dataset strategy item size {}".format(len(shape), len(dataset_strategy[index])))
             new_shape = ()
             for i, item in enumerate(shape):
-                new_shape += (item * dataset_strategy[index][i],)
+                if item > 0:
+                    new_shape += (item * dataset_strategy[index][i],)  # static shape
+                else:
+                    new_shape += (item,)  # dynamic shape
             new_shapes.append(new_shape)
         return new_shapes
     for shape in shapes:
         new_shape = ()
         for i, item in enumerate(shape):
-            if i == 0:
-                new_shape += (item * device_num,)
+            if i == 0 and item > 0:
+                new_shape += (item * device_num,)  # only for static shape
             else:
                 new_shape += (item,)
         new_shapes.append(new_shape)
@@ -201,7 +205,7 @@ def _to_full_tensor(elem, global_device_num, global_rank, scaling_sens=None):
                 slice_index += (s,)
             new_tensor_numpy = np.zeros(new_shape, dtype_to_nptype(type_))
             new_tensor_numpy[slice_index] = data.asnumpy()
-        new_tensor = Tensor(new_tensor_numpy)
+        new_tensor = Tensor(new_tensor_numpy, dtype=type_)
         lst.append(new_tensor)
     if scaling_sens:
         lst.append(Tensor(scaling_sens, mstype.float32))

mindspore/parallel/algo_parameter_config.py

@@ -229,7 +229,7 @@ def set_algo_parameters(**kwargs):
     """
     Set parameters in the algorithm for parallel strategy searching. See a typical use in
     `test_auto_parallel_resnet.py
-    <https://gitee.com/mindspore/mindspore/blob/r2.1/tests/ut/python/parallel/test_auto_parallel_resnet.py>`_.
+    <https://gitee.com/mindspore/mindspore/blob/r2.2/tests/ut/python/parallel/test_auto_parallel_resnet.py>`_.
 
     Note:
         The attribute name is required. This interface works ONLY in AUTO_PARALLEL mode.
@@ -239,10 +239,10 @@ def set_algo_parameters(**kwargs):
             Default: ``True`` . For example with 8 devices available, if set ``True`` , strategy (4, 1) will not be
             included in ReLU's candidate strategies, because strategy (4, 1) only utilizes 4 devices.
         elementwise_op_strategy_follow (bool): Whether the elementwise operator has the consistent strategies as its
-            subsequent operators. Default: ``False`` . For the example of ReLU followed by Add, where ReLU is
-            elementwise operator, if this flag is set ``True`` , then the searched strategy by the algorithm
-            guarantees that strategies of these two operators are consistent, e.g., ReLU's strategy (8, 1) and Add's
-            strategy ((8, 1), (8, 1)).
+            subsequent operators. Elementwise operators refer to operators that operate on input element by element,
+            such as Add, ReLU, etc. Default: ``False`` . For the example of ReLU followed by Add, if this flag is set
+            ``True`` , then the searched strategy by the algorithm guarantees that strategies of these two operators
+            are consistent, e.g., ReLU's strategy (8, 1) and Add's strategy ((8, 1), (8, 1)).
         enable_algo_approxi (bool): Whether to enable the approximation in the algorithms. Default: ``False`` . Due to
             large solution space in searching parallel strategy for large DNN model, the algorithm takes fairly long
             time in this case. To mitigate it, if this flag is set ``True`` , an approximation is made to discard some
@@ -261,8 +261,87 @@ def set_algo_parameters(**kwargs):
         ValueError: If context keyword is not recognized.
 
     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 numpy as np
         >>> import mindspore as ms
+        >>> import mindspore.dataset as ds
+        >>> from mindspore import nn, ops, train
+        >>> from mindspore.communication import init
+        >>> from mindspore.common.initializer import initializer
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.AUTO_PARALLEL,
+        >>>                              search_mode="sharding_propagation")
+        >>> init()
+        >>> ms.set_algo_parameters(fully_use_devices=True)
         >>> ms.set_algo_parameters(elementwise_op_strategy_follow=True)
+        >>> ms.set_algo_parameters(enable_algo_approxi=True)
+        >>> ms.set_algo_parameters(algo_approxi_epsilon=0.2)
+        >>> ms.set_algo_parameters(tensor_slice_align_enable=True)
+        >>> ms.set_algo_parameters(tensor_slice_align_size=8)
+        >>>
+        >>> # Define the network structure.
+        >>> class Dense(nn.Cell):
+        ...     def __init__(self, in_channels, out_channels):
+        ...         super().__init__()
+        ...         self.weight = ms.Parameter(initializer("normal", [in_channels, out_channels], ms.float32))
+        ...         self.bias = ms.Parameter(initializer("normal", [out_channels], ms.float32))
+        ...         self.matmul = ops.MatMul()
+        ...         self.add = ops.Add()
+        ...
+        ...     def construct(self, x):
+        ...         x = self.matmul(x, self.weight)
+        ...         x = self.add(x, self.bias)
+        ...         return x
+        >>>
+        >>> class FFN(nn.Cell):
+        ...     def __init__(self):
+        ...         super().__init__()
+        ...         self.flatten = ops.Flatten()
+        ...         self.dense1 = Dense(28*28, 64)
+        ...         self.relu = ops.ReLU()
+        ...         self.dense2 = Dense(64, 10)
+        ...
+        ...     def construct(self, x):
+        ...         x = self.flatten(x)
+        ...         x = self.dense1(x)
+        ...         x = self.relu(x)
+        ...         x = self.dense2(x)
+        ...         return x
+        >>> net = FFN()
+        >>> net.dense1.matmul.shard(((2, 1), (1, 2)))
+        >>>
+        >>> # Create dataset.
+        >>> step_per_epoch = 16
+        >>> def get_dataset(*inputs):
+        ...     def generate():
+        ...         for _ in range(step_per_epoch):
+        ...             yield inputs
+        ...     return generate
+        >>>
+        >>> input_data = np.random.rand(1, 28, 28).astype(np.float32)
+        >>> label_data = np.random.rand(1).astype(np.int32)
+        >>> fake_dataset = get_dataset(input_data, label_data)
+        >>> dataset = ds.GeneratorDataset(fake_dataset, ["input", "label"])
+        >>> # Train network.
+        >>> optimizer = nn.Momentum(net.trainable_params(), 1e-3, 0.1)
+        >>> loss_fn = nn.CrossEntropyLoss()
+        >>> loss_cb = train.LossMonitor()
+        >>> model = ms.Model(network=net, loss_fn=loss_fn, optimizer=optimizer)
+        >>> model.train(epoch=2, train_dataset=dataset, callbacks=[loss_cb])
     """
     for key, value in kwargs.items():
         if key not in set_algo_parameters_config_func_map:
@@ -282,6 +361,7 @@ def get_algo_parameters(attr_key):
         attr_key (str): The key of the attribute. The keys include: "fully_use_devices",
             "elementwise_op_strategy_follow", "enable_algo_approxi", "algo_approxi_epsilon",
             "tensor_slice_align_enable","tensor_slice_align_size".
+            See :func:`mindspore.set_algo_parameters` for more details about the meaning of the attributes.
 
     Returns:
         Return attribute value according to the key.