mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/ops/operations/comm_ops.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2020-2023 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,15 +13,14 @@
 # limitations under the License.
 # ============================================================================
 
-"""Communication APIs."""
+"""Communication APIs.
+"""
 from __future__ import absolute_import
 from __future__ import division
 
 from mindspore.common import Tensor
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
-from mindspore.communication.management import get_rank, get_group_size, GlobalComm, _get_group
-from mindspore.communication._comm_helper import _check_mpi_envs
+from mindspore import _checkparam as validator
+from mindspore.communication.management import get_rank, get_group_size, GlobalComm, _get_group, _host_distribute
 from mindspore.common import dtype as mstype
 from mindspore.ops.primitive import PrimitiveWithInfer, PrimitiveWithCheck, Primitive, prim_attr_register
 from mindspore.common.api import context
@@ -30,7 +29,6 @@ from mindspore.common.api import context
 class ReduceOp:
     """
     Operation options for reducing tensors. This is an enumerated type, not an operator.
-    Mainly used in data parallel mode.
 
     The main calling methods are as follows:
 
@@ -46,17 +44,25 @@ class ReduceOp:
     - MIN: Take the minimum.
     - PROD: Take the product.
 
-    Note:
-        For more, refer to example. This needs to run in an environment with multiple graphics cards.
-        The user needs to preset
-        communication environment variables before running the following example, please check the details on the
-        official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
-
     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 `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            for more details.
+
+            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with multiple devices.
+
+        >>> import numpy as np
+        >>> import mindspore
         >>> from mindspore.communication import init
         >>> from mindspore import Tensor, ops, nn
         >>> from mindspore.ops import ReduceOp
@@ -96,25 +102,21 @@ def check_collective_target_dtype(data_name, data_dtype, prim_name):
 def check_hcom_group_valid(group, prim_name=None):
     """Check if hcom group is valid."""
     msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    if not _check_mpi_envs() and context.get_context("mode") == context.PYNATIVE_MODE and \
+    if not _host_distribute() and context.get_context("mode") == context.PYNATIVE_MODE and \
             group != GlobalComm.WORLD_COMM_GROUP:
         raise RuntimeError(f"{msg_prefix} 'group' only support 'hccl_world_group' in pynative mode, but got "
                            f"'group': {group}. Please start by using mpi-run.")
 
 
-class AllReduce(PrimitiveWithInfer):
+class AllReduce(Primitive):
     """
     Reduces the tensor data across all devices in such a way that all devices will get the same final result.
 
     Note:
-        The operation of AllReduce does not support "prod" currently.
-        The tensors must have the same shape and format in all processes of the collection. The user needs to preset
-        communication environment variables before running the following example, please check the details on the
-        official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
+        The tensors must have the same shape and format in all processes of the collection.
 
     Args:
-        op (str): Specifies an operation used for element-wise reductions, like sum, max, and min.
+        op (str): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
                   On the CPU, only 'sum' is supported. Default: ReduceOp.SUM.
         group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
 
@@ -128,13 +130,24 @@ class AllReduce(PrimitiveWithInfer):
     Raises:
         TypeError: If any of `op` and `group` is not a str,
                    or fusion is not an integer, or the input's dtype is bool.
-        ValueError: If the `op` is "prod".
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> # This example should be run with two devices. Refer to the tutorial > Distributed Training on mindspore.cn
+        .. 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 `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            for more details.
+
+            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with multiple devices.
+
         >>> import numpy as np
         >>> from mindspore.communication import init
         >>> from mindspore import Tensor
@@ -174,23 +187,14 @@ class AllReduce(PrimitiveWithInfer):
         self.add_prim_attr('index', 0)
         self.add_prim_attr('no_eliminate', True)
 
-    def infer_shape(self, x_shape):
-        return x_shape
-
-    def infer_dtype(self, x_dtype):
-        check_collective_target_dtype('x', x_dtype, self.name)
-        return x_dtype
-
 
 class AllGather(PrimitiveWithInfer):
     """
     Gathers tensors from the specified communication group.
 
     Note:
-        The tensors must have the same shape and format in all processes of the collection. The user needs to preset
-        communication environment variables before running the following example. Please check the details on the
-        official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
+        - The tensors must have the same shape and format in all processes of the collection.
+        - Currently only supports GRAPH_MODE and it should be called in Cell.
 
     Args:
         group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
@@ -211,7 +215,19 @@ class AllGather(PrimitiveWithInfer):
         ``Ascend`` ``GPU``
 
     Examples:
-        >>> # This example should be run with two devices. Refer to the tutorial > Distributed Training on mindspore.cn
+        .. 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 `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            for more details.
+
+            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with 2 devices.
+
         >>> import numpy as np
         >>> import mindspore as ms
         >>> import mindspore.ops as ops
@@ -245,13 +261,16 @@ class AllGather(PrimitiveWithInfer):
         validator.check_value_type('group', _get_group(group), (str,), self.name)
         self.rank = get_rank(_get_group(group))
         self.rank_size = get_group_size(_get_group(group))
-        validator.check('rank', self.rank, 'rank_size', self.rank_size, Rel.LT, self.name)
+        validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
         self.add_prim_attr('rank_size', self.rank_size)
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('mean_flag', False)
         self.add_prim_attr('no_eliminate', True)
 
+    def __call__(self, tensor):
+        raise NotImplementedError
+
     def infer_shape(self, x_shape):
         validator.check_positive_int(len(x_shape), "x shape", self.name)
         if x_shape[0] > 0:
@@ -262,9 +281,6 @@ class AllGather(PrimitiveWithInfer):
         check_collective_target_dtype('x', x_dtype, self.name)
         return x_dtype
 
-    def __call__(self, tensor):
-        raise NotImplementedError
-
 
 class _MiniStepAllGather(PrimitiveWithInfer):
     """
@@ -282,13 +298,14 @@ class _MiniStepAllGather(PrimitiveWithInfer):
         validator.check_value_type('group', _get_group(group), (str,), self.name)
         self.rank = get_rank(_get_group(group))
         self.rank_size = get_group_size(_get_group(group))
-        validator.check('rank', self.rank, 'rank_size', self.rank_size, Rel.LT, self.name)
+        validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
         self.add_prim_attr('rank_size', self.rank_size)
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('fusion', 1)
         self.grad_accumulation_step = grad_accumulation_step
         self.mean_flag = mean_flag
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, z_shape):
         validator.check_positive_int(len(x_shape), "x shape", self.name)
@@ -313,15 +330,17 @@ class _MicroStepAllGather(PrimitiveWithInfer):
     @prim_attr_register
     def __init__(self, group=GlobalComm.WORLD_COMM_GROUP, mean_flag=None):
         validator.check_value_type('group', _get_group(group), (str,), self.name)
-        self.rank = get_rank(_get_group(group))
-        self.rank_size = get_group_size(_get_group(group))
-        validator.check('rank', self.rank, 'rank_size', self.rank_size, Rel.LT, self.name)
-        self.add_prim_attr('rank_size', self.rank_size)
-        self.add_prim_attr('group', _get_group(group))
-        self.add_prim_attr('fusion', 1)
-        self.add_prim_attr('do_mirror', False)
-        self.mean_flag = mean_flag
-        self.add_prim_attr('order_enforce_skip', True)
+        self.rank_size = 1
+        if group != "":
+            self.rank = get_rank(_get_group(group))
+            self.rank_size = get_group_size(_get_group(group))
+            validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
+            self.add_prim_attr('rank_size', self.rank_size)
+            self.add_prim_attr('group', _get_group(group))
+            self.add_prim_attr('fusion', 1)
+            self.add_prim_attr('do_mirror', False)
+            self.mean_flag = mean_flag
+            self.add_prim_attr('order_enforce_skip', True)
 
     def infer_shape(self, x_shape, z_shape):
         validator.check_positive_int(len(x_shape), "x shape", self.name)
@@ -365,9 +384,9 @@ class _HostAllGather(PrimitiveWithInfer):
         if group is None:
             raise ValueError(f"For '{self.name}', the 'group' cannot be None, but got {group}.")
         validator.check_value_type('group', group, (tuple, list), self.name)
-        validator.check_int(len(group), 2, Rel.GE, "group size", self.name)
+        validator.check_int(len(group), 2, validator.GE, "group size", self.name)
         for r in group:
-            validator.check_int_range(r, 0, 7, Rel.INC_BOTH, "rank_id", self.name)
+            validator.check_int_range(r, 0, 7, validator.INC_BOTH, "rank_id", self.name)
             validator.check_value_type("rank_id", r, (int,), self.name)
         self.group_size = len(group)
         self.add_prim_attr('group', group)
@@ -388,31 +407,49 @@ class _HostAllGather(PrimitiveWithInfer):
         raise NotImplementedError
 
 
-class ReduceScatter(PrimitiveWithInfer):
-    """
+class ReduceScatter(Primitive):
+    r"""
     Reduces and scatters tensors from the specified communication group.
+    For more details about it, please refer to `Distributed Set Communication Primitives - ReduceScatter \
+    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#reducescatter>`_ .
 
     Note:
-        The back propagation of the op is not supported yet. Stay tuned for more.
-        The tensors must have the same shape and format in all processes of the collection. The user needs to preset
-        communication environment variables before running the following example, please check the details on the
-        official website of `Communication Operator API \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
+        The tensors must have the same shape and format in all processes of the collection.
 
     Args:
         op (str): Specifies an operation used for element-wise reductions,
-                  like SUM, MAX, AVG. Default: ReduceOp.SUM.
+                  like SUM and MAX. Default: ReduceOp.SUM.
         group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
 
+    Inputs:
+        - **input_x** (Tensor) - Input Tensor, suppose it has a shape :math:`(N, *)`, where `*`
+          means any number of additional dimensions. N must be divisible by rank_size.
+          rank_size refers to the number of cards in the communication group.
+
+    Outputs:
+        Tensor, it has the same dtype as `input_x` with a shape of :math:`(N/rank\_size, *)`.
+
     Raises:
         TypeError: If any of operation and group is not a string.
-        ValueError: If the first dimension of the input cannot be divided by the rank size.
+        ValueError: If the first dimension of the input cannot be divided by the rank_size.
 
     Supported Platforms:
         ``Ascend`` ``GPU``
 
     Examples:
-        >>> # This example should be run with two devices. Refer to the tutorial > Distributed Training on mindspore.cn
+        .. 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 `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            for more details.
+
+            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with 2 devices.
+
         >>> import mindspore as ms
         >>> from mindspore import Tensor
         >>> from mindspore.communication import init
@@ -453,19 +490,6 @@ class ReduceScatter(PrimitiveWithInfer):
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('no_eliminate', True)
 
-    def infer_shape(self, x_shape):
-        if self.rank_size == 0:
-            raise ValueError(f"For '{self.name}', the 'rank_size' can not be zero, but got {self.rank_size}.")
-        if x_shape[0] % self.rank_size != 0:
-            raise ValueError(f"For '{self.name}', the first dimension of 'x_shape' must be divided by 'rank_size', "
-                             f"but got 'x_shape[0]': {x_shape[0]}, 'rank_size': {self.rank_size}.")
-        x_shape[0] = int(x_shape[0] / self.rank_size)
-        return x_shape
-
-    def infer_dtype(self, x_dtype):
-        check_collective_target_dtype('x', x_dtype, self.name)
-        return x_dtype
-
     def __call__(self, tensor):
         raise NotImplementedError
 
@@ -499,9 +523,9 @@ class _HostReduceScatter(PrimitiveWithInfer):
             raise ValueError(f"For '{self.name}', the 'group' cannot be None, but got {group}.")
         validator.check_value_type('op', op, (type(ReduceOp.SUM),), self.name)
         validator.check_value_type('group', group, (tuple, list), self.name)
-        validator.check_int(len(group), 2, Rel.GE, "group size", self.name)
+        validator.check_int(len(group), 2, validator.GE, "group size", self.name)
         for r in group:
-            validator.check_int_range(r, 0, 7, Rel.INC_BOTH, "rank_id", self.name)
+            validator.check_int_range(r, 0, 7, validator.INC_BOTH, "rank_id", self.name)
             validator.check_value_type("rank_id", r, (int,), self.name)
         self.op = op
         self.group_size = len(group)
@@ -529,10 +553,7 @@ class Broadcast(PrimitiveWithInfer):
     Broadcasts the tensor to the whole group.
 
     Note:
-        The tensors must have the same shape and format in all processes of the collection. The user needs to preset
-        communication environment variables before running the following example, please check the details on the
-        official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
+        The tensors must have the same shape and format in all processes of the collection.
 
     Args:
         root_rank (int): Source rank. Required in all processes except the one
@@ -543,7 +564,7 @@ class Broadcast(PrimitiveWithInfer):
         - **input_x** (tuple[Tensor]) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
 
     Outputs:
-        Tensor, has the same shape of the input, i.e., :math:`(x_1, x_2, ..., x_R)`.
+        tuple[Tensor], Tensor has the same shape of the input, i.e., :math:`(x_1, x_2, ..., x_R)`.
         The contents depend on the data of the `root_rank` device.
 
     Raises:
@@ -553,10 +574,19 @@ class Broadcast(PrimitiveWithInfer):
         ``Ascend`` ``GPU``
 
     Examples:
-        >>> # This example should be run with multiple processes.
-        >>> # Please refer to the Programming Guide > Distributed Training -> Distributed Parallel Usage Example
-        >>> # on mindspore.cn and focus on the contents of these three parts: Configuring Distributed Environment
-        >>> # Variables, Calling the Collective Communication Library, Running The Script.
+        .. 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 `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            for more details.
+
+            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with multiple devices.
+
         >>> import mindspore as ms
         >>> from mindspore import Tensor
         >>> from mindspore.communication import init
@@ -599,7 +629,7 @@ class Broadcast(PrimitiveWithInfer):
         if not isinstance(x_dtype, tuple):
             raise TypeError(f"For '{self.name}', the 'input_x' must be a tuple, but got {type(x_dtype).__name__}!")
         for _ele in x_dtype:
-            check_collective_target_dtype('x', _ele, self.name)
+            check_collective_target_dtype('input_x', _ele, self.name)
         return x_dtype
 
 
@@ -669,11 +699,11 @@ class NeighborExchange(Primitive):
         The user needs to preset
         communication environment variables before running the following example, please check the details on the
         official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
+        <https://www.mindspore.cn/docs/en/r2.0/api_python/mindspore.ops.html#communication-operator>`_.
 
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/zh-CN/r1.10/parallel/communicate_ops.html#注意事项>`_.
+        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to.
@@ -747,14 +777,9 @@ class AlltoAll(PrimitiveWithInfer):
     - The gather phase: Each process concatenates the received blocks along the concat_dimension.
 
     Note:
-        The tensors must have the same shape and format in all processes of the collection. The user needs to preset
-        communication environment variables before running the following example, please check the details on the
-        official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
-
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/zh-CN/r1.10/parallel/communicate_ops.html#注意事项>`_.
+        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
 
     Args:
         split_count (int): On each process, divide blocks into split_count number.
@@ -769,11 +794,9 @@ class AlltoAll(PrimitiveWithInfer):
         Tensor. If the shape of input tensor is :math:`(x_1, x_2, ..., x_R)`, then the shape of output tensor is
         :math:`(y_1, y_2, ..., y_R)`, where:
 
-        :math:`y_{split\_dim} = x_{split\_dim} / split\_count`
-
-        :math:`y_{concat\_dim} = x_{concat\_dim} * split\_count`
-
-        :math:`y\_other = x\_other`.
+        - :math:`y_{split\_dim} = x_{split\_dim} / split\_count`
+        - :math:`y_{concat\_dim} = x_{concat\_dim} * split\_count`
+        - :math:`y_{other} = x_{other}`.
 
     Raises:
         TypeError: If group is not a string.
@@ -782,7 +805,19 @@ class AlltoAll(PrimitiveWithInfer):
         ``Ascend``
 
     Examples:
-        >>> # This example should be run with 8 devices. Refer to the tutorial > Distributed Training on mindspore.cn
+        .. 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 `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            for more details.
+
+            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with 8 devices.
+
         >>> import os
         >>> import mindspore as ms
         >>> from mindspore import Tensor
@@ -843,21 +878,19 @@ class AlltoAll(PrimitiveWithInfer):
 
 
 class NeighborExchangeV2(Primitive):
-    """
-    NeighborExchangeV2 is a collective operation.
+    r"""
+    NeighborExchangeV2 is a collective communication operation.
 
-    NeighborExchangeV2 sends data from the local rank to ranks in the send_rank_ids,
-    as while receive data from recv_rank_ids.
+    NeighborExchangeV2 sends data from the local rank to ranks in the `send_rank_ids`,
+    as while receive data from `recv_rank_ids`. Please refer to
+    `Distributed Set Communication Primitives - NeighborExchangeV2 \
+    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#neighborexchangev2>`_
+    to learn about how the data is exchanged between neighborhood devices.
 
     Note:
-        The user needs to preset
-        communication environment variables before running the following example, please check the details on the
-        official website of `Communication Operator API \
-        <https://www.mindspore.cn/docs/en/r1.10/api_python/mindspore.ops.html#communication-operator>`_.
-
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
-        in the same subnet, please check the `Distributed Set Communication Primitives - Notes \
-        <https://www.mindspore.cn/tutorials/experts/zh-CN/r1.10/parallel/communicate_ops.html#注意事项>`_.
+        in the same subnet, please check the `details \
+        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to. 8 rank_ids represents 8 directions, if one
@@ -865,17 +898,44 @@ class NeighborExchangeV2(Primitive):
         recv_rank_ids (list(int)): Ranks which the data is received from. 8 rank_ids represents 8 directions,
                                    if one direction is not recv from , set it -1.
         send_lens (list(int)): Data lens which send to the send_rank_ids, 4 numbers represent the lens of
-                               [top, bottom, left, right].
+                               [send_top, send_bottom, send_left, send_right].
         recv_lens (list(int)): Data lens which received from recv_rank_ids, 4 numbers represent the lens of
-                               [top, bottom, left, right].
+                               [recv_top, recv_bottom, recv_left, recv_right].
         data_format (str): Data format, only support NCHW now.
-        group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
+        group (str, optional): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP", which means
+                     "hccl_world_group" in Ascend, and "nccl_world_group" in GPU.
+
+    Inputs:
+        - **input_x** (Tensor) - The Tensor before being exchanged. It has a shape of :math:`(N, C, H, W)`.
+
+    Outputs:
+        The Tensor after being exchanged. If input shape is :math:`(N, C, H, W)`, output shape is
+        :math:`(N, C, H+recv\_top+recv\_bottom, W+recv\_left+recv\_right)`.
+
+    Raises:
+        TypeError: If `group` is not a string or any one of `send_rank_ids`,
+            `recv_rank_ids`, `send_lens`, `recv_lens` is not a list.
+        ValueError: If `send_rank_ids` or `recv_rank_ids` has value less than -1 or has repeated values.
+        ValueError: If `send_lens`, `recv_lens` has value less than 0.
+        ValueError: If `data_format` is not "NCHW".
 
     Supported Platforms:
         ``Ascend``
 
-    Example:
-        >>> # This example should be run with 2 devices. Refer to the tutorial > Distributed Training on mindspore.cn
+    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 `Ascend tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            for more details.
+
+            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+
+            This example should be run with 2 devices.
+
         >>> import os
         >>> import mindspore as ms
         >>> from mindspore import Tensor
@@ -971,6 +1031,7 @@ class _MirrorMiniStepOperator(PrimitiveWithInfer):
         self.mean_flag = mean_flag
         self.grad_accumulation_step = grad_accumulation_step
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, z_shape):
         return x_shape
@@ -1076,6 +1137,7 @@ class _VirtualAssignAdd(PrimitiveWithInfer):
     def __init__(self):
         """Initialize _VirtualAssignAdd."""
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, y_shape):
         return x_shape
@@ -1126,6 +1188,7 @@ class _MirrorMicroStepOperator(PrimitiveWithInfer):
         self.dev_num = dev_num
         self.mean_flag = mean_flag
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, z_shape):
         return x_shape

mindspore/ops/operations/control_ops.py

@@ -14,8 +14,10 @@
 # ============================================================================
 
 """control_ops"""
+
+from __future__ import absolute_import
 from mindspore.ops.primitive import PrimitiveWithInfer, prim_attr_register
-from mindspore._checkparam import Validator as validator, Rel
+from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 
 
@@ -74,7 +76,7 @@ class GeSwitch(PrimitiveWithInfer):
         raise NotImplementedError
 
     def infer_shape(self, data, pred):
-        validator.check_int_range(len(pred), 0, 1, Rel.INC_BOTH, "pred rank", self.name)
+        validator.check_int_range(len(pred), 0, 1, validator.INC_BOTH, "pred rank", self.name)
         return data, data
 
     def infer_dtype(self, data_type, pred_type):