mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

mindspore/nn/wrap/cell_wrapper.py

@@ -7,6 +7,7 @@
 #
 # http://www.apache.org/licenses/LICENSE-2.0
 #
+#
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -85,8 +86,8 @@ class WithLossCell(Cell):
         loss_fn (Cell): The loss function used to compute loss.
 
     Inputs:
-        - **data** (Tensor) - Tensor of shape :math:`(N, \ldots)`.
-        - **label** (Tensor) - Tensor of shape :math:`(N, \ldots)`.
+        - **data** (Tensor) - Tensor of shape :math:`(N, \ldots)`. The dtype of `data` must be float16 or float32.
+        - **label** (Tensor) - Tensor of shape :math:`(N, \ldots)`. The dtype of `label` must be float16 or float32.
 
     Outputs:
         Tensor, a tensor means the loss value, the shape of which is usually :math:`()`.
@@ -328,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`` .
 
@@ -529,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:
@@ -571,122 +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.
-
-    Note:
-        The output of the input network must be a single tensor.
-
-    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:
-        Tensor. The output of the input `network` .
-
-    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.
@@ -736,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.
@@ -753,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:
@@ -776,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)
@@ -966,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

@@ -31,7 +31,6 @@ from mindspore.ops import composite as C
 from mindspore.ops import operations as P
 from mindspore.ops.operations.nn_ops import AllFinite
 from mindspore.common import dtype as mstype
-from mindspore.common.api import jit
 from mindspore._c_expression import MSContext
 from mindspore.run_check._check_version import AscendEnvChecker
 from mindspore import log as logger
@@ -93,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
@@ -377,7 +376,7 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         self.loss_scaling_manager = None
         self._ascend_check_overflow_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
 
-        self.enable_allfinite = False
+        self.enable_allfinite = True
         runtime_conf = os.environ.get('MS_DEV_RUNTIME_CONF')
         global_jit_config = context.get_jit_config()
         if runtime_conf is not None and ("all_finite:True" in runtime_conf or "all_finite:true" in runtime_conf):
@@ -389,7 +388,8 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         elif global_jit_config:
             logger.debug("Current global jit config is: {}".format(global_jit_config["jit_level"]))
             self.enable_allfinite = global_jit_config["jit_level"] == "O0" or global_jit_config["jit_level"] == "O1"
-
+        if "RANK_TABLE_FILE" in os.environ:
+            self.enable_allfinite = False
         if self.ascend_910b_target:
             checker = AscendEnvChecker(None)
             if not checker.check_custom_version():
@@ -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
 
@@ -548,7 +548,6 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         overflow = self.logic_not(overall_finite)
         return overflow
 
-    @jit
     def get_overflow_status(self, status, compute_output):
         """
         Get floating-point overflow status.