mindspore 2.5.0__cp310-cp310-win_amd64.whl → 2.6.0rc1__cp310-cp310-win_amd64.whl

mindspore/nn/probability/distribution/poisson.py

@@ -36,7 +36,7 @@ class Poisson(Distribution):
     where :math:`\lambda` is the rate of the distribution.
 
     Args:
-        rate (list, numpy.ndarray, Tensor): The rate of the Poisson distribution. Default: ``None`` .
+        rate (list, numpy.ndarray, Tensor): The rate of the Poisson distribution. :math:`\lambda` in the formula. Default: ``None`` .
         seed (int): The seed used in sampling. The global seed is used if it is ``None`` . Default: ``None`` .
         dtype (mindspore.dtype): The type of the event samples. Default: ``mstype.float32`` .
         name (str): The name of the distribution. Default: ``'Poisson'`` .

mindspore/nn/probability/distribution/uniform.py

@@ -36,8 +36,10 @@ class Uniform(Distribution):
     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`` .
-        high (int, float, list, numpy.ndarray, Tensor): The upper bound of the distribution. Default: ``None`` .
+        low (int, float, list, numpy.ndarray, Tensor): The lower bound of the distribution.
+            :math:`a` in the formula. Default: ``None`` .
+        high (int, float, list, numpy.ndarray, Tensor): The upper bound of the distribution.
+            :math:`b` in the formula. Default: ``None`` .
         seed (int): The seed uses in sampling. The global seed is used if it is ``None`` . Default: ``None`` .
         dtype (mindspore.dtype): The type of the event samples. Default: ``mstype.float32`` .
         name (str): The name of the distribution. Default: ``'Uniform'`` .
@@ -49,7 +51,7 @@ class Uniform(Distribution):
 
     Raises:
         ValueError: When high <= low.
-        TypeError: When the input `dtype` is not a subclass of float.
+        TypeError: When the input `dtype` is not a float or a subclass of float.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``

mindspore/nn/reinforcement/_tensors_queue.py

@@ -34,7 +34,7 @@ class TensorsQueue(Cell):
         dtype (mindspore.dtype): the data type in the TensorsQueue. Each tensor should have the same dtype.
         shapes (tuple[int64]): the shape of each element in TensorsQueue.
         size (int): the size of the TensorsQueue.
-        name (string): the name of this TensorsQueue. Default: "TQ".
+        name (str): the name of this TensorsQueue. Default: "TQ".
 
     Raises:
         TypeError: If `dtype` is not mindspore number type.

mindspore/nn/reinforcement/tensor_array.py

@@ -34,7 +34,7 @@ class TensorArray(Cell):
         element_shape (tuple[int]): the shape of each tensor in a TensorArray.
         dynamic_size (bool): if ``true`` , the size of TensorArray can be increased. Default: ``True`` .
         size (int): if dynamic_size=False, `size` means the max_size of the TensorArray.
-        name (string): the name of this TensorArray. Default: ``"TA"`` .
+        name (str): the name of this TensorArray. Default: ``"TA"`` .
 
     Supported Platforms:
         ``GPU`` ``CPU``

mindspore/nn/wrap/__init__.py

@@ -20,8 +20,8 @@ Use the Wrapper to combine the loss or build the training steps.
 from __future__ import absolute_import
 
 from mindspore.nn.wrap.cell_wrapper import ForwardValueAndGrad, TrainOneStepCell, WithLossCell, WithGradCell, \
-    WithEvalCell, ParameterUpdate, GetNextSingleOp, VirtualDatasetCellTriple, MicroBatchInterleaved, PipelineCell, \
-    GradAccumulationCell
+    WithEvalCell, ParameterUpdate, GetNextSingleOp, VirtualDatasetCellTriple, GradAccumulationCell, \
+    MicroBatchInterleaved, PipelineCell
 from mindspore.nn.wrap.loss_scale import TrainOneStepWithLossScaleCell,\
     DynamicLossScaleUpdateCell, FixedLossScaleUpdateCell
 from mindspore.nn.wrap.grad_reducer import DistributedGradReducer, PipelineGradReducer
@@ -34,16 +34,16 @@ __all__ = [
     "TrainOneStepCell",
     "WithLossCell",
     "WithGradCell",
-    "MicroBatchInterleaved",
-    "PipelineCell",
     "WithEvalCell",
     "GetNextSingleOp",
     "TrainOneStepWithLossScaleCell",
     "DistributedGradReducer",
-    "PipelineGradReducer",
     "ParameterUpdate",
     "DynamicLossScaleUpdateCell",
     "FixedLossScaleUpdateCell",
     "VirtualDatasetCellTriple",
-    "GradAccumulationCell"
+    "GradAccumulationCell",
+    "MicroBatchInterleaved",
+    "PipelineCell",
+    "PipelineGradReducer"
     ]

mindspore/nn/wrap/cell_wrapper.py

@@ -329,9 +329,11 @@ class TrainOneStepCell(Cell):
     Args:
         network (Cell): The training network. The network only supports single output.
         optimizer (Union[Cell]): Optimizer for updating the network parameters.
-        sens (numbers.Number): The scaling number to be filled as the input of backpropagation. Default value is
+        sens (numbers.Number, optional): The scaling number to be filled as the input of backpropagation.
+            Default value is
             ``None`` , which is ``1.0`` .
-        return_grad (bool): Whether to return gradient. If ``True``, it will return the gradient in the form of a dict
+        return_grad (bool, optional): Whether to return gradient. If ``True``,
+            it will return the gradient in the form of a dict
             while returning loss. The key of the dict is the parameter name corresponding to the gradient, and value
             is the gradient value. Default value is ``False`` .
 
@@ -530,6 +532,20 @@ class _VirtualDatasetCell(Cell):
         return self._backbone(*output)
 
 
+def _pipeline_clear_grad(accu_grad, grad):
+    accu_grad = F.depend(accu_grad, grad)
+    zeros = F.zeros_like(accu_grad)
+    return F.assign(accu_grad, zeros)
+
+def grad_scale(scale, grad):
+    """grad_scale"""
+    new_grad = scale * grad
+    grad = ops.depend(grad, new_grad)
+    zeros = F.zeros_like(grad)
+    new_grad = ops.depend(new_grad, F.assign(grad, zeros))
+    return new_grad
+
+
 @_primexpr
 def _check_shape_value_on_axis_divided_by_target_value(input_shape, micro_size):
     if F.isconstant(input_shape[0]) is False:
@@ -572,119 +588,13 @@ class _MicroBatch(Cell):
         return micro_inputs
 
 
-class MicroBatchInterleaved(Cell):
-    """
-    This function splits the input at the 0th into interleave_num pieces and then performs
-    the computation of the wrapped cell. Application scenario: When there is model parallelism in semi-automatic mode
-    and network, if the first slice data is calculating forward, the second slice data will execute the
-    communication operators at the same time, to achieve the performance acceleration of communication and computing
-    concurrency.
-
-    Args:
-        network (Cell): The target network to wrap.
-        interleave_num (int, optional): split num of batch size. Default: ``2`` .
-
-    Inputs:
-        tuple[Tensor]. It's the same with the input of the `network` .
-
-    Outputs:
-        The wrapped input. The output of the input `network` should be a Tensor.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU``
-
-    Examples:
-        >>> import mindspore.nn as nn
-        >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
-        >>> net = LeNet5()
-        >>> net = nn.MicroBatchInterleaved(net, 2)
-    """
-    def __init__(self, network, interleave_num=2):
-        super(MicroBatchInterleaved, self).__init__(auto_prefix=False)
-        if not isinstance(interleave_num, int):
-            raise TypeError("For 'MicroBatchInterleaved', the argument 'interleave_num' must be integer, "
-                            "but got the type : {}.".format(type(interleave_num)))
-        if interleave_num <= 0:
-            raise ValueError("For 'MicroBatchInterleaved', the argument 'interleave_num' must be large than 0, "
-                             "but got {}.".format(interleave_num))
-        self.network = network
-        self.interleave_num = interleave_num
-        self.interleave_inputs = nn.CellList()
-        self.add = P.Add().add_prim_attr("micro_interleaved_add_flag", True)
-        for _ in range(interleave_num):
-            interleave_data = _MicroBatch(interleave_num)
-            interleave_data.strided_slice.add_prim_attr("strided_slice_flag", True)
-            interleave_data.strided_slice.add_prim_attr("interleave_num", interleave_num)
-            self.interleave_inputs.append(interleave_data)
-        self._get_attr_from_cell(network)
-
-    def construct(self, *inputs):
-        output = 0.0
-        for i in range(self.interleave_num):
-            interleave_input = self.interleave_inputs[i](i, *inputs)
-            output = self.add(output, self.network(*interleave_input))
-        return output
-
-
-class PipelineCell(Cell):
-    """
-    Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training.
-
-    Note:
-        micro_size must be greater or equal to pipeline stages.
-
-    Args:
-        network (Cell): The target network to wrap.
-        micro_size (int): MicroBatch size.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU``
-
-    Examples:
-        >>> import mindspore.nn as nn
-        >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
-        >>> net = LeNet5()
-        >>> net = nn.PipelineCell(net, 4)
-    """
-    def __init__(self, network, micro_size):
-        super(PipelineCell, self).__init__(auto_prefix=False)
-        self.network = network
-        self.micro_inputs = nn.CellList()
-        self.micro_size = micro_size
-        self.add_list = []
-        if not isinstance(network, Cell):
-            raise TypeError("For 'PipelineCell', the argument 'network' must cell type, "
-                            "but got the type : {}.".format(type(network)))
-        if not isinstance(micro_size, int):
-            raise TypeError("For 'PipelineCell', the argument 'micro_size' must be integer, "
-                            "but got the type : {}.".format(type(micro_size)))
-        if micro_size <= 0:
-            raise ValueError("For 'PipelineCell', the argument 'micro_size' must be large than 0, "
-                             "but got {}.".format(micro_size))
-        for i in range(micro_size):
-            micro_input = _MicroBatch(micro_size)
-            self.micro_inputs.append(micro_input)
-            self.add = P.Add().add_prim_attr("pipeline_end", i)
-            self.add_list.append(self.add)
-        self._get_attr_from_cell(network)
-
-    def construct(self, *inputs):
-        ret = None
-        for i in range(self.micro_size):
-            micro_input = self.micro_inputs[i](i, *inputs)
-            output = self.network(*micro_input)
-            if ret is not None:
-                ret = self.add_list[i](ret, output)
-            else:
-                ret = output
-        return ret
-
 class GradAccumulationCell(Cell):
     """
     Wrap the network with Micro Batch to enable the grad accumulation in semi_auto_parallel/auto_parallel mode.
 
+    Note:
+        The api will be deprecated, please use the api :class:`mindspore.parallel.nn.GradAccumulation` instead.
+
     Args:
         network (Cell): The target network to wrap.
         micro_size (int): MicroBatch size.
@@ -734,12 +644,6 @@ class GradAccumulationCell(Cell):
         return ret
 
 
-def _pipeline_clear_grad(accu_grad, grad):
-    accu_grad = F.depend(accu_grad, grad)
-    zeros = F.zeros_like(accu_grad)
-    return F.assign(accu_grad, zeros)
-
-
 class _TrainGradAccuStepCell(TrainOneStepCell):
     """
     Wraps the network with an optimizer in pipeline mode.
@@ -751,6 +655,13 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         self.opt_shard = _get_enable_parallel_optimizer()
         self._get_attr_from_cell(network)
         self.enable_tft = False
+        if not self.sense_flag:
+            micro_size = 1.0
+            for _, cell in network.cells_and_names():
+                if hasattr(cell, 'micro_size'):
+                    micro_size = cell.micro_size
+                    break
+            self.sens = 1 / micro_size
 
     def construct(self, *inputs):
         if not self.sense_flag:
@@ -774,8 +685,10 @@ class _TrainGradAccuStepCell(TrainOneStepCell):
         grads = self.grad_no_sens(self.network, self.weights)(*inputs)
         accu_grads = ops.depend(self.accu_grads, grads)
         if self.opt_shard:
+            grads = self.hyper_map(F.partial(grad_scale, self.sens), grads)
             succ = self.optimizer(grads)
         else:
+            accu_grads = self.hyper_map(F.partial(grad_scale, self.sens), accu_grads)
             succ = self.optimizer(accu_grads)
         loss = ops.depend(loss, succ)
         clear = self.hyper_map(_pipeline_clear_grad, accu_grads, grads)
@@ -964,3 +877,151 @@ class _BroadCastCell(Cell):
         params = self.broadcast(params)
         new_params = self.map_(F.partial(_cast_datatype), datatypes, params)
         return new_params
+
+
+class PipelineCell(Cell):
+    """
+    Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training.
+
+    Note:
+        - micro_size must be greater or equal to pipeline stages.
+        - The api will be deprecated, please use the api :class:`mindspore.parallel.nn.Pipeline` instead.
+
+    Args:
+        network (Cell): The target network to wrap.
+        micro_size (int): MicroBatch size.
+        stage_config (dict, optional): The stage configuration for each cell's execution in pipeline parallel.
+            Default ``None``.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        >>> import mindspore.nn as nn
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = nn.PipelineCell(net, 4)
+    """
+    def __init__(self, network, micro_size, stage_config=None):
+        super(PipelineCell, self).__init__(auto_prefix=False)
+        self.network = network
+        self.micro_inputs = nn.CellList()
+        self.micro_size = micro_size
+        self.add_list = []
+        if not isinstance(network, Cell):
+            raise TypeError("For 'PipelineCell', the argument 'network' must cell type, "
+                            "but got the type : {}.".format(type(network)))
+        if not isinstance(micro_size, int):
+            raise TypeError("For 'PipelineCell', the argument 'micro_size' must be integer, "
+                            "but got the type : {}.".format(type(micro_size)))
+        if micro_size <= 0:
+            raise ValueError("For 'PipelineCell', the argument 'micro_size' must be large than 0, "
+                             "but got {}.".format(micro_size))
+        for i in range(micro_size):
+            micro_input = _MicroBatch(micro_size)
+            self.micro_inputs.append(micro_input)
+            self.add = P.Add().add_prim_attr("pipeline_end", i)
+            self.add_list.append(self.add)
+        self._get_attr_from_cell(network)
+
+        # prase stage_config
+        config_dict = {}
+        if stage_config is not None:
+            for cell_name, stage_num in stage_config.items():
+                config_cell_name = cell_name
+                config_stage_num = stage_num
+                config_dict[config_cell_name] = config_stage_num
+
+        # set cell.stage_config
+            for cell_name, cell in self.network.cells_and_names():
+                for config_cell_name, config_stage_num in config_dict.copy().items():
+                    if not cell_name or not config_cell_name:
+                        continue
+                    if cell_name == config_cell_name:
+                        setattr(cell, "pipeline_stage", config_stage_num)
+                        del config_dict[config_cell_name]
+
+            for config_cell_name, config_stage_num in config_dict.copy().items():
+                if str(network) == config_cell_name:
+                    setattr(network, "pipeline_stage", config_stage_num)
+                    del config_dict[config_cell_name]
+
+            # if there are any config elements left, print them
+            if config_dict:
+                for config_cell_name, config_stage_num in config_dict.items():
+                    print("pipeline_cell stage_config set pipeline_stage fail!")
+                    print("config cell name:" + str(config_cell_name) +
+                          " config stage num:" + str(config_stage_num))
+                print("network:" + str(self.network))
+                print("cell name available:")
+                for cell_name, cell in self.network.cells_and_names():
+                    print(cell_name)
+                raise KeyError("For 'PipelineCell', the argument 'stage_config' : {} is not "
+                               "found in 'network' : {}".format(config_dict, network))
+
+    def construct(self, *inputs):
+        ret = None
+        for i in range(self.micro_size):
+            micro_input = self.micro_inputs[i](i, *inputs)
+            output = self.network(*micro_input)
+            if ret is not None:
+                ret = self.add_list[i](ret, output)
+            else:
+                ret = output
+        return ret
+
+
+class MicroBatchInterleaved(Cell):
+    """
+    This function splits the input at the 0th into interleave_num pieces and then performs
+    the computation of the wrapped cell. Application scenario: When there is model parallelism in semi-automatic mode
+    and network, if the first slice data is calculating forward, the second slice data will execute the
+    communication operators at the same time, to achieve the performance acceleration of communication and computing
+    concurrency.
+
+    Args:
+        network (Cell): The target network to wrap.
+        interleave_num (int, optional): split num of batch size. Default: ``2`` .
+
+    Inputs:
+        tuple[Tensor]. It's the same with the input of the `network` .
+
+    Outputs:
+        The wrapped input. The output of the input `network` should be a Tensor.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        >>> import mindspore.nn as nn
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = nn.MicroBatchInterleaved(net, 2)
+    """
+    def __init__(self, network, interleave_num=2):
+        super(MicroBatchInterleaved, self).__init__(auto_prefix=False)
+        if not isinstance(interleave_num, int):
+            raise TypeError("For 'MicroBatchInterleaved', the argument 'interleave_num' must be integer, "
+                            "but got the type : {}.".format(type(interleave_num)))
+        if interleave_num <= 0:
+            raise ValueError("For 'MicroBatchInterleaved', the argument 'interleave_num' must be large than 0, "
+                             "but got {}.".format(interleave_num))
+        self.network = network
+        self.interleave_num = interleave_num
+        self.interleave_inputs = nn.CellList()
+        self.add = P.Add().add_prim_attr("micro_interleaved_add_flag", True)
+        for _ in range(interleave_num):
+            interleave_data = _MicroBatch(interleave_num)
+            interleave_data.strided_slice.add_prim_attr("strided_slice_flag", True)
+            interleave_data.strided_slice.add_prim_attr("interleave_num", interleave_num)
+            self.interleave_inputs.append(interleave_data)
+        self._get_attr_from_cell(network)
+
+    def construct(self, *inputs):
+        output = 0.0
+        for i in range(self.interleave_num):
+            interleave_input = self.interleave_inputs[i](i, *inputs)
+            output = self.add(output, self.network(*interleave_input))
+        return output

mindspore/nn/wrap/grad_reducer.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2025 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.
@@ -18,7 +18,6 @@ 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, operations as P
@@ -28,30 +27,13 @@ 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.nn.layer import Identity
 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()
+__all__ = ['DistributedGradReducer']
 
 
-@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
+reduce_opt = C.MultitypeFuncGraph("reduce_opt")
 
 
 def _init_allreduce_operators(length, split_indices, group=GlobalComm.WORLD_COMM_GROUP):
@@ -335,14 +317,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/docs/en/master/model_train/parallel/rank_table.html>`_
+            <https://www.mindspore.cn/tutorials/en/master/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>`_ .
+            <https://www.mindspore.cn/tutorials/en/master/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/docs/en/master/model_train/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/tutorials/en/master/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -427,7 +409,8 @@ class DistributedGradReducer(Cell):
             self.degree = degree
         self.degree = Tensor(1.0 / self.degree, mstype.float32)
 
-        self.allreduce_filter = tuple((x.layerwise_parallel is False) and (x.is_in_shard is False) for x in parameters)
+        self.allreduce_filter = tuple((x.layerwise_parallel is False) and
+                                      (not x.param_info.is_in_pynative_shard) for x in parameters)
         is_parallel_optimizer = context.get_auto_parallel_context("enable_parallel_optimizer")
         split_indices = auto_parallel_context().get_all_reduce_fusion_split_indices()
         if is_parallel_optimizer and split_indices:
@@ -447,7 +430,7 @@ class DistributedGradReducer(Cell):
         self.mode = context.get_context("mode")
         self.enable_tuple_broaden = True
 
-    @jit
+    @jit(backend="ms_backend")
     def construct(self, grads):
         """
         Under certain circumstances, the data precision of grads could be mixed with float16 and float32. Thus, the
@@ -488,13 +471,39 @@ class DistributedGradReducer(Cell):
             raise RuntimeError("{} can not use DistributedGradReducer in graph mode".format(parallel_mode))
 
 
+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
+
+
 class PipelineGradReducer(Cell):
     """
     PipelineGradReducer is a gradient reducer for pipeline parallelism.
 
+    Note:
+        The api will be deprecated, please use the api :class:`mindspore.parallel.nn.PipelineGradReducer` instead.
+
     Args:
         parameters (list): the parameters to be updated.
-        scale_sense (float): the scale sense of the gradient. Default: 1.0.
+        scale_sense (float, optional): the scale sense of the gradient. Default: ``1.0``.
 
     Raise:
         RuntimeError: If the mode is not graph mode.
@@ -509,11 +518,11 @@ class PipelineGradReducer(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/docs/en/master/model_train/parallel/rank_table.html>`_
+            <https://www.mindspore.cn/tutorials/en/master/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>`_ .
+            <https://www.mindspore.cn/tutorials/en/master/parallel/mpirun.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -554,7 +563,7 @@ class PipelineGradReducer(Cell):
         >>> 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 = nn.Pipeline(nn.WithLossCell(net, loss_fn), 2)
         >>> net_with_loss.set_train()
         >>> def forward_fn(inputs, target):
         ...     loss = net_with_loss(inputs, target)
@@ -576,7 +585,7 @@ class PipelineGradReducer(Cell):
         >>> print(loss)
         46.36721
     """
-    def __init__(self, parameters, scale_sense=1.0):
+    def __init__(self, parameters, scale_sense=1.0, opt_shard=None):
         super(PipelineGradReducer, self).__init__(auto_prefix=False)
         self._check_mode()
         self.accu_grads = parameters.clone(prefix="accu_grads", init="zeros")
@@ -584,7 +593,10 @@ class PipelineGradReducer(Cell):
         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()
+        if opt_shard is None:
+            self.opt_shard = _get_enable_parallel_optimizer()
+        else:
+            self.opt_shard = opt_shard
 
     @jit
     def construct(self, grads):
@@ -603,6 +615,3 @@ class PipelineGradReducer(Cell):
         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")

mindspore/nn/wrap/loss_scale.py

@@ -92,8 +92,8 @@ class DynamicLossScaleUpdateCell(Cell):
     Dynamic Loss scale update cell.
 
     For loss scaling training, the initial loss scaling value will be set to be `loss_scale_value`.
-    In each training step, the loss scaling value will be decreased by `loss_scale`/`scale_factor`
-    when there is an overflow. And it will be increased by `loss_scale` * `scale_factor` if there is no
+    In each training step, the loss scaling value will be decreased by :math:`loss\_scale/scale\_factor`
+    when there is an overflow. And it will be increased by :math:`loss\_scale * scale\_factor` if there is no
     overflow for a continuous `scale_window` steps.
 
     `get_update_cell` method of :class:`mindspore.amp.DynamicLossScaleManager` will return this class. It will be called
@@ -506,7 +506,7 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         overflow = AllFinite()(compute_output)
 
         if self.is_distributed:
-            overflow = P.Cast()(overflow, mstype.int8)
+            overflow = P.Cast()(overflow, mstype.float32)
             overflow = P.Cast()(self.allreduce(overflow), mstype.bool_)
         return overflow
 

mindspore/numpy/array_creations.py

@@ -30,7 +30,7 @@ from mindspore.ops.primitive import constexpr, _primexpr
 from mindspore.ops.function.random_func import _get_seed
 from mindspore.nn.layer.basic import tril as nn_tril
 from mindspore.nn.layer.basic import triu as nn_triu
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 
 from mindspore.numpy.utils import _check_input_for_asarray, _deep_list, _deep_tensor_to_nparray, \
     _check_input_tensor, _convert_64_to_32, _get_dtype_from_scalar, \
@@ -133,7 +133,7 @@ def asarray_const(a, dtype=None):
             elif dtype == mstype.int64:
                 dtype = mstype.int32
         if a.size == 0:
-            a = Tensor_(a)
+            a = Tensor(a)
 
     if isinstance(a, onp.ndarray) and dtype is None:
         if a.dtype is onp.dtype('object'):
@@ -945,7 +945,7 @@ def identity(n, dtype=mstype.float32):
 @constexpr
 def empty_compile(dtype, shape):
     """Returns an empty Tensor."""
-    return Tensor_(dtype, shape)
+    return Tensor(dtype=dtype, shape=shape)
 
 
 def empty(shape, dtype=mstype.float32):

mindspore/numpy/array_ops.py

@@ -1945,7 +1945,7 @@ def repeat(a, repeats, axis=None):
         [3 4]]
     """
     a = _to_tensor(a)
-    return a.repeat(repeats, axis)
+    return a.repeat_interleave(repeats, axis)
 
 
 def rot90(a, k=1, axes=(0, 1)):