mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/nn/probability/distribution/_utils/custom_ops.py

@@ -62,12 +62,12 @@ def log_generic(input_x):
     if not checktype(dtype(input_x), mstype.float_):
         input_x = cast(input_x, mstype.float32)
     nan = F.fill(dtype(input_x), shape(input_x), np.nan)
-    inf = F.fill(dtype(input_x), shape(input_x), np.inf)
+    neginf = F.fill(dtype(input_x), shape(input_x), -np.inf)
     neg_x = less(input_x, 0.0)
     nonpos_x = lessequal(input_x, 0.0)
     log_x = log(input_x)
     result = select(
-        nonpos_x, (-1.0) * inf, log_x)
+        nonpos_x, neginf, log_x)
     return select(neg_x, nan, result)
 
 

mindspore/nn/probability/distribution/beta.py

@@ -31,9 +31,9 @@ class Beta(Distribution):
     A Beta distributio is a continuous distribution with the range :math:`[0, 1]` and the probability density function:
 
     .. math::
-        f(x, \alpha, \beta) = x^\alpha (1-x)^{\beta - 1} / B(\alpha, \beta),
+        f(x, \alpha, \beta) = x^\alpha (1-x)^{\beta - 1} / B(\alpha, \beta)
 
-    where :math:`B` is the Beta function.
+    Where :math:`B` is the Beta function.
 
     Args:
         concentration1 (int, float, list, numpy.ndarray, Tensor): The concentration1,

mindspore/nn/probability/distribution/categorical.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2020-2024 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,7 +15,6 @@
 """Categorical Distribution"""
 import numpy as np
 from mindspore import context
-from mindspore.common import Tensor
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
 from mindspore.ops import composite as C
@@ -328,10 +327,9 @@ class Categorical(Distribution):
         value_clipped = self.clip_by_value(value, 0.0, num_classes - 1)
         value_clipped = self.cast(value_clipped, self.index_type)
         # create index from 0 ... NumOfLabels
-        start = Tensor(0, self.index_type)
-        end = self.cast(self.shape(value)[0], self.index_type)
-        delta = Tensor(1, self.index_type)
-        index = self.reshape(ops.range(start, end, delta), (-1, 1))
+        index = self.reshape(
+            ops.arange(0, self.shape(value)[0], 1, dtype=self.index_type), (-1, 1)
+        )
         index = self.concat((index, value_clipped))
 
         # index into logit_pmf, fill in out_of_bound places with -inf

mindspore/nn/probability/distribution/cauchy.py

@@ -30,9 +30,9 @@ class Cauchy(Distribution):
     and the probability density function:
 
     .. math::
-        f(x, a, b) = 1 / \pi b(1 - ((x - a)/b)^2),
+        f(x, a, b) = 1 / \pi b(1 - ((x - a)/b)^2)
 
-    where :math:`a, b` are loc and scale parameter respectively.
+    Where :math:`a, b` are loc and scale parameter respectively.
 
     Args:
         loc (int, float, list, numpy.ndarray, Tensor): The location of the Cauchy distribution. Default: ``None`` .

mindspore/nn/probability/distribution/exponential.py

@@ -30,7 +30,7 @@ class Exponential(Distribution):
     and the probability density function:
 
     .. math::
-        f(x, \lambda) = \lambda \exp(-\lambda x),
+        f(x, \lambda) = \lambda \exp(-\lambda x)
 
     where :math:`\lambda` is the rate of the distribution.
 
@@ -152,7 +152,7 @@ class Exponential(Distribution):
         if self.rate is not None:
             check_greater_zero(self.rate, 'rate')
 
-        self.minval = np.finfo(np.float).tiny
+        self.minval = np.finfo(np.float_).tiny
 
         # ops needed for the class
         self.exp = exp_generic

mindspore/nn/probability/distribution/geometric.py

@@ -152,7 +152,7 @@ class Geometric(Distribution):
         if self._probs is not None:
             check_prob(self.probs)
 
-        self.minval = np.finfo(np.float).tiny
+        self.minval = np.finfo(np.float_).tiny
 
         # ops needed for the class
         self.exp = exp_generic

mindspore/nn/probability/distribution/gumbel.py

@@ -32,9 +32,9 @@ class Gumbel(TransformedDistribution):
     and the probability density function:
 
     .. math::
-        f(x, a, b) = 1 / b \exp(\exp(-(x - a) / b) - x),
+        f(x, a, b) = 1 / b \exp(\exp(-(x - a) / b) - x)
 
-    where :math:`a, b` are loc and scale parameter respectively.
+    Where :math:`a, b` are loc and scale parameter respectively.
 
     Args:
         loc (int, float, list, numpy.ndarray, Tensor): The location of Gumbel distribution.

mindspore/nn/probability/distribution/logistic.py

@@ -170,7 +170,7 @@ class Logistic(Distribution):
         self.neg = P.Neg()
 
         self.threshold = np.log(np.finfo(np.float32).eps) + 1.
-        self.tiny = np.finfo(np.float).tiny
+        self.tiny = np.finfo(np.float_).tiny
         self.sd_const = np.pi / np.sqrt(3)
 
     def _softplus(self, x):

mindspore/nn/probability/distribution/poisson.py

@@ -229,11 +229,11 @@ class Poisson(Distribution):
         rate = self._check_param_type(rate)
         log_rate = self.log(rate)
         zeros = F.fill(self.dtypeop(value), self.shape(value), 0.0)
-        inf = F.fill(self.dtypeop(value), self.shape(value), np.inf)
+        neginf = F.fill(self.dtypeop(value), self.shape(value), -np.inf)
         safe_x = self.select(self.less(value, zeros), zeros, value)
         y = log_rate * safe_x - self.lgamma(safe_x + 1.)
         comp = self.equal(value, safe_x)
-        log_unnormalized_prob = self.select(comp, y, (-1) * inf)
+        log_unnormalized_prob = self.select(comp, y, neginf)
         log_normalization = self.exp(log_rate)
         return log_unnormalized_prob - log_normalization
 

mindspore/nn/probability/distribution/uniform.py

@@ -31,9 +31,9 @@ class Uniform(Distribution):
     and the probability density function:
 
     .. math::
-        f(x, a, b) = 1 / (b - a),
+        f(x, a, b) = 1 / (b - a)
 
-    where :math:`a, b` are the lower and upper bound respectively.
+    Where :math:`a, b` are the lower and upper bound respectively.
 
     Args:
         low (int, float, list, numpy.ndarray, Tensor): The lower bound of the distribution. Default: ``None`` .

mindspore/nn/reinforcement/_tensors_queue.py

@@ -37,7 +37,7 @@ class TensorsQueue(Cell):
         name (string): the name of this TensorsQueue. Default: "TQ".
 
     Raises:
-        TypeError: If `dtype` is not mindpore number type.
+        TypeError: If `dtype` is not mindspore number type.
         ValueError: If `size` is less than 0.
         ValueError: If `shapes` size is less than 1.
 
@@ -69,6 +69,7 @@ class TensorsQueue(Cell):
         self.tensors_q_clear = rl_ops.TensorsQueueClear()
         self.tensors_q_close = rl_ops.TensorsQueueClose()
         self.tensors_q_size = rl_ops.TensorsQueueSize()
+        self.__is_tensors_queue__ = True
 
     def put(self, element):
         """
@@ -103,6 +104,17 @@ class TensorsQueue(Cell):
         element = self.tensors_q_pop(self.handle_)
         return element
 
+    def __graph_pop__(self):
+        """
+        Get one element int the front of the TensorsQueue, and remove it.
+        This is only used in graph mode.
+
+        Returns:
+            tuple(Tensors), the element in TensorsQueue.
+        """
+        element = self.tensors_q_pop(self.handle_)
+        return self.handle_, element
+
     def size(self):
         """
         Get the used/available size of the TensorsQueue, and remove it.

mindspore/nn/wrap/__init__.py

@@ -24,7 +24,7 @@ from mindspore.nn.wrap.cell_wrapper import ForwardValueAndGrad, TrainOneStepCell
     GradAccumulationCell
 from mindspore.nn.wrap.loss_scale import TrainOneStepWithLossScaleCell,\
     DynamicLossScaleUpdateCell, FixedLossScaleUpdateCell
-from mindspore.nn.wrap.grad_reducer import DistributedGradReducer
+from mindspore.nn.wrap.grad_reducer import DistributedGradReducer, PipelineGradReducer
 from mindspore.nn.layer.timedistributed import TimeDistributed
 
 
@@ -40,6 +40,7 @@ __all__ = [
     "GetNextSingleOp",
     "TrainOneStepWithLossScaleCell",
     "DistributedGradReducer",
+    "PipelineGradReducer",
     "ParameterUpdate",
     "DynamicLossScaleUpdateCell",
     "FixedLossScaleUpdateCell",

mindspore/nn/wrap/cell_wrapper.py

@@ -17,6 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 
+import os
 from types import FunctionType, MethodType
 
 from mindspore import log as logger
@@ -27,6 +28,7 @@ from mindspore import _checkparam as validator
 from mindspore import ops, nn
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter, ParameterTuple
+from mindspore.common.tensor import Tensor
 from mindspore.ops.primitive import _primexpr
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
@@ -34,6 +36,7 @@ from mindspore.ops import operations as P
 from mindspore.ops.operations.comm_ops import _VirtualDataset
 from mindspore.nn.cell import Cell
 from mindspore.nn.wrap.grad_reducer import DistributedGradReducer
+from mindspore.utils import ExitByRequest
 
 _get_datatype = C.MultitypeFuncGraph("_get_datatype")
 
@@ -99,7 +102,7 @@ class WithLossCell(Cell):
         >>> from mindspore import Tensor, nn
         >>> import numpy as np
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=False)
         >>> net_with_criterion = nn.WithLossCell(net, loss_fn)
@@ -132,7 +135,7 @@ class WithLossCell(Cell):
         Examples:
             >>> from mindspore import nn
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=False)
             >>> net_with_criterion = nn.WithLossCell(net, loss_fn)
@@ -175,7 +178,7 @@ class WithGradCell(Cell):
         >>> import mindspore as ms
         >>> from mindspore import nn
         >>> # Defined a network without loss function, taking LeNet5 as an example.
-        >>> # Refer to https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # Refer to https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> grad_net = nn.WithGradCell(net, loss_fn)
@@ -346,7 +349,7 @@ class TrainOneStepCell(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> optim = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
@@ -412,6 +415,11 @@ class TrainOneStepCell(Cell):
                 group = server_group_name
             self.grad_reducer = DistributedGradReducer(self.weights, self.mean, self.degree, group=group)
         self._get_attr_from_cell(network)
+        self.use_graceful_exit = os.environ.get("MS_ENABLE_GRACEFUL_EXIT") == "1"
+        if self.use_graceful_exit:
+            self.graceful_exit = ExitByRequest()
+            self.exit_param = Parameter(Tensor(False, mstype.bool_), name="graceful_exit")  # update by reduce value
+            self.init_param = Parameter(Tensor([0], mstype.int32), name="graceful_init")  # update by config file
 
     def construct(self, *inputs):
         if not self.sense_flag:
@@ -420,6 +428,8 @@ class TrainOneStepCell(Cell):
         sens = F.fill(loss.dtype, loss.shape, self.sens)
         grads = self.grad(self.network, self.weights)(*inputs, sens)
         grads = self.grad_reducer(grads)
+        if self.use_graceful_exit:
+            grads = self.graceful_exit.exit_by_request(grads, self.init_param, self.exit_param)
         loss = F.depend(loss, self.optimizer(grads))
         if self.return_grad:
             grad_with_param_name = {}
@@ -433,6 +443,8 @@ class TrainOneStepCell(Cell):
         loss = self.network(*inputs)
         grads = self.grad_no_sens(self.network, self.weights)(*inputs)
         grads = self.grad_reducer(grads)
+        if self.use_graceful_exit:
+            grads = self.graceful_exit.exit_by_request(grads, self.init_param, self.exit_param)
         loss = F.depend(loss, self.optimizer(grads))
         if self.return_grad:
             grad_with_param_name = {}
@@ -454,7 +466,7 @@ class GetNextSingleOp(Cell):
         queue_name (str): Queue name to fetch the data.
 
     Outputs:
-        tuple[Tensor], the data get from Dataset.
+        tuple[Tensor], the data gets from Dataset.
 
     Supported Platforms:
         ``Ascend`` ``GPU``
@@ -586,7 +598,7 @@ class MicroBatchInterleaved(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> net = nn.MicroBatchInterleaved(net, 2)
     """
@@ -619,7 +631,7 @@ class MicroBatchInterleaved(Cell):
 
 class PipelineCell(Cell):
     """
-    Wrap the network with Micro Batch.
+    Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training.
 
     Note:
         micro_size must be greater or equal to pipeline stages.
@@ -634,7 +646,7 @@ class PipelineCell(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> net = nn.PipelineCell(net, 4)
     """
@@ -685,7 +697,7 @@ class GradAccumulationCell(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> net = nn.GradAccumulationCell(net, 4)
     """
@@ -726,7 +738,7 @@ class GradAccumulationCell(Cell):
 
 def _pipeline_clear_grad(accu_grad, grad):
     accu_grad = F.depend(accu_grad, grad)
-    zeros = F.tensor_mul(accu_grad, 0.0)
+    zeros = F.zeros_like(accu_grad)
     return F.assign(accu_grad, zeros)
 
 
@@ -740,6 +752,7 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         self.hyper_map = ops.HyperMap()
         self.opt_shard = _get_enable_parallel_optimizer()
         self._get_attr_from_cell(network)
+        self.enable_tft = False
 
     def construct(self, *inputs):
         if not self.sense_flag:
@@ -772,6 +785,27 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         return loss
 
 
+class AllreduceGraph(Cell):
+    """
+    A allreduce graph to broadcast parameters.
+    """
+    def __init__(self, inputs, group_name):
+        super(AllreduceGraph, self).__init__()
+        self.input_num = len(inputs)
+        self.inputs = inputs
+        self.allreduces = []
+        self.assigns = []
+        for _ in range(self.input_num):
+            self.allreduces.append(ops.AllReduce(op="sum", group=group_name))
+            self.assigns.append(ops.Assign())
+
+    def construct(self):
+        for i in range(self.input_num):
+            res = self.allreduces[i](self.inputs[i])
+            self.assigns[i](self.inputs[i], res)
+        return self.inputs
+
+
 class VirtualDatasetCellTriple(Cell):
     """
     Wrap the network with virtual dataset to convert data parallel layout to model parallel layout.
@@ -790,7 +824,7 @@ class VirtualDatasetCellTriple(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> net = nn.VirtualDatasetCellTriple(net)
     """
@@ -833,7 +867,7 @@ class WithEvalCell(Cell):
     Examples:
         >>> import mindspore.nn as nn
         >>> # Define a forward network without loss function, taking LeNet5 as an example.
-        >>> # Refer to https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # Refer to https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits()
         >>> eval_net = nn.WithEvalCell(net, loss_fn)

mindspore/nn/wrap/grad_reducer.py

@@ -18,17 +18,40 @@ from __future__ import absolute_import
 from mindspore import context
 from mindspore import log as logger
 from mindspore.nn.cell import Cell
+from mindspore.nn.layer import Identity
 from mindspore.communication.management import GlobalComm, get_group_size
 from mindspore.common.sparse_tensor import RowTensorInner
-from mindspore.ops import functional as F, composite as C
+from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.ops.operations.comm_ops import AllReduce, AllGather
 from mindspore.parallel._auto_parallel_context import auto_parallel_context
 import mindspore.common.dtype as mstype
 from mindspore.common.sparse_tensor import Tensor
 from mindspore.common.api import jit
-
+from mindspore.common.parameter import Parameter
+from mindspore.parallel._utils import _get_enable_parallel_optimizer
 
 reduce_opt = C.MultitypeFuncGraph("reduce_opt")
+grad_scale = C.MultitypeFuncGraph("grad_scale")
+shard_grad_scale = C.MultitypeFuncGraph("shard_grad_scale")
+reciprocal = P.Reciprocal()
+
+
+@grad_scale.register("Tensor", "Tensor", "Tensor")
+def tensor_grad_scale_pipeline(scale, grad, accu_grad):
+    accu_grad = F.depend(accu_grad, grad)
+    new_grad = accu_grad * reciprocal(scale)
+    accu_grad = F.depend(accu_grad, new_grad)
+    zeros = F.tensor_mul(accu_grad, 0.0)
+    new_grad = F.depend(new_grad, F.assign(accu_grad, zeros))
+    return new_grad
+
+
+@shard_grad_scale.register("Tensor", "Tensor", "Tensor")
+def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
+    new_grad = grad * reciprocal(scale)
+    accu_grad = F.depend(accu_grad, new_grad)
+    new_grad = F.depend(new_grad, F.assign(accu_grad, F.zeros_like(accu_grad)))
+    return new_grad
 
 
 def _init_allreduce_operators(length, split_indices, group=GlobalComm.WORLD_COMM_GROUP):
@@ -53,9 +76,6 @@ def _init_allreduce_operators(length, split_indices, group=GlobalComm.WORLD_COMM
     for i in range(length):
         op = AllReduce('sum', group)
         op_fusion_id = fusion[i]
-        # When running in ge and enabled all_reduce_fusion_config, hccl will check the allreduce' fusion id to be -1
-        if context.get_context("enable_ge") and context.get_auto_parallel_context("all_reduce_fusion_config"):
-            op_fusion_id = -1
         op.add_prim_attr('fusion', op_fusion_id)
         op.add_prim_attr('index', index[i])
         op_list = op_list + (op,)
@@ -315,14 +335,14 @@ class DistributedGradReducer(Cell):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/rank_table.html>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/mpirun.html>`_ .
 
             For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/docs/en/master/model_train/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -466,3 +486,123 @@ class DistributedGradReducer(Cell):
         if context.get_context('mode') == context.GRAPH_MODE and parallel_mode in (
                 context.ParallelMode.SEMI_AUTO_PARALLEL, context.ParallelMode.AUTO_PARALLEL):
             raise RuntimeError("{} can not use DistributedGradReducer in graph mode".format(parallel_mode))
+
+
+class PipelineGradReducer(Cell):
+    """
+    PipelineGradReducer is a gradient reducer for pipeline parallelism.
+
+    Args:
+        parameters (list): the parameters to be updated.
+        scale_sense (float): the scale sense of the gradient. Default: 1.0.
+
+    Raise:
+        RuntimeError: If the mode is not graph mode.
+        RuntimeError: If the parallel mode is not semi auto parallel or auto parallel.
+
+    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/docs/en/master/model_train/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/docs/en/master/model_train/parallel/mpirun.html>`_ .
+
+            This example should be run with multiple devices.
+
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import nn, ops, Tensor
+        >>> from mindspore.communication import init
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> ms.reset_auto_parallel_context()
+        >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL, pipeline_stages=2)
+        >>> init()
+        >>> ms.set_seed(1)
+        >>>
+        >>> class Network(nn.Cell):
+        ...     def __init__(self, in_features, out_features, sens=1.0):
+        ...         super().__init__()
+        ...         self.layer1 = nn.Dense(in_features, 16)
+        ...         self.relu1 = nn.ReLU()
+        ...         self.layer2 = nn.Dense(16, 16)
+        ...         self.relu2 = nn.ReLU()
+        ...         self.layer3 = nn.Dense(16, out_features)
+        ...
+        ...     def construct(self, x):
+        ...         x = self.layer1(x)
+        ...         x = self.relu1(x)
+        ...         x = self.layer2(x)
+        ...         x = self.relu2(x)
+        ...         logits = self.layer3(x)
+        ...         return logits
+        >>>
+        >>> size, in_features, out_features = 16, 32, 10
+        >>> net = Network(in_features, out_features)
+        >>> net.layer1.pipeline_stage = 0
+        >>> net.relu1.pipeline_stage = 0
+        >>> net.layer2.pipeline_stage = 0
+        >>> net.relu2.pipeline_stage = 1
+        >>> net.layer3.pipeline_stage = 1
+        >>> loss_fn = nn.CrossEntropyLoss()
+        >>> optimizer = nn.SGD(net.trainable_params(), 1e-2)
+        >>> net_with_loss = nn.PipelineCell(nn.WithLossCell(net, loss_fn), 2)
+        >>> net_with_loss.set_train()
+        >>> def forward_fn(inputs, target):
+        ...     loss = net_with_loss(inputs, target)
+        ...     return loss
+        >>>
+        >>> grad_fn = ops.value_and_grad(forward_fn, None, net_with_loss.trainable_params())
+        >>> pp_grad_reducer = nn.PipelineGradReducer(optimizer.parameters)
+        >>>
+        >>> @ms.jit
+        >>> def train_one_step(inputs, target):
+        ...     loss, grads = grad_fn(inputs, target)
+        ...     grads = pp_grad_reducer(grads)
+        ...     optimizer(grads)
+        ...     return loss, grads
+        >>>
+        >>> inputs = Tensor(np.ones([size, in_features]).astype(np.float32))
+        >>> label = Tensor(np.ones([size, out_features]).astype(np.float32))
+        >>> loss, _ = train_one_step(inputs, label)
+        >>> print(loss)
+        46.36721
+    """
+    def __init__(self, parameters, scale_sense=1.0):
+        super(PipelineGradReducer, self).__init__(auto_prefix=False)
+        self._check_mode()
+        self.accu_grads = parameters.clone(prefix="accu_grads", init="zeros")
+        self.grad_reducer = Identity()
+        self.degree = Tensor(1, mstype.float32)
+        self.scale_sense = Parameter(scale_sense, name='scale_sense')
+        self.hyper_map = C.HyperMap()
+        self.opt_shard = _get_enable_parallel_optimizer()
+
+    @jit
+    def construct(self, grads):
+        new_grads = None
+        if self.opt_shard:
+            grads = self.grad_reducer(grads)
+            new_grads = self.hyper_map(F.partial(shard_grad_scale, self.scale_sense * self.degree),
+                                       grads, self.accu_grads)
+        else:
+            accu_grads = self.grad_reducer(self.accu_grads)
+            new_grads = self.hyper_map(F.partial(grad_scale, self.scale_sense * self.degree), grads, accu_grads)
+        return new_grads
+
+    def _check_mode(self):
+        """check parallel mode"""
+        mode = context.get_context('mode')
+        if mode != context.GRAPH_MODE:
+            raise RuntimeError(f"PipelineGradReducer only support graph mode, but get {mode}")
+        parallel_mode = context.get_auto_parallel_context('parallel_mode')
+        if parallel_mode not in (context.ParallelMode.SEMI_AUTO_PARALLEL, context.ParallelMode.AUTO_PARALLEL):
+            raise RuntimeError(f"{parallel_mode} can not use PipelineGradReducer in graph mode")