mindspore 2.4.10__cp310-cp310-win_amd64.whl → 2.6.0__cp310-cp310-win_amd64.whl

mindspore/_extends/pijit/pijit_func_white_list.py

@@ -29,7 +29,7 @@ from mindspore.common.api import jit
 from mindspore.common.tensor import Tensor
 from mindspore.common._register_for_tensor import Registry
 from mindspore._c_expression import MetaFuncGraph_, function_id
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._extends.parse.resources import convert_object_map
 from mindspore import _checkparam as validator
 from mindspore import Parameter, ParameterTuple
@@ -49,11 +49,12 @@ from mindspore.train.data_sink import _init_sink_dataset
 from mindspore.train.summary import SummaryRecord
 from mindspore.train._utils import _exec_datagraph
 from mindspore.train.summary.writer import BaseWriter
-from mindspore.train.serialization import _exec_save, load, export_split_mindir, obfuscate_model, _parse_ckpt_proto, \
+from mindspore.train.serialization import _exec_save, load, export_split_mindir, _parse_ckpt_proto, \
     _generate_front_info_for_param_data_file, _get_data_file, _encrypt_data, _split_save, _save_mindir_together, \
     _load_into_param_dict
 from mindspore.parallel import _cost_model_context
 from mindspore.parallel._offload_context import offload_context
+from mindspore.parallel._utils import _is_in_data_parallel_mode
 from mindspore.run_check._check_version import check_version_and_env_config
 from mindspore.dataset.callback.ds_callback import DSCallback, WaitedDSCallback
 from mindspore.dataset.transforms.c_transforms import TensorOperation as CTensorOperation, OneHot as COneHot, \
@@ -127,7 +128,7 @@ from mindspore.dataset.vision.transforms import AdjustBrightness, AdjustContrast
     RandomVerticalFlipWithBBox as VRandomVerticalFlipWithBBox, Rescale as VRescale, Resize as VResize, ResizedCrop, \
     ResizeWithBBox as VResizeWithBBox, Rotate as VRotate, SlicePatches as VSlicePatches, Solarize, ToTensor,\
     TrivialAugmentWide, UniformAugment as VUniformAugment, VerticalFlip as VVerticalFlip
-from mindspore.profiler.profiling import Profiler
+from mindspore.profiler.profiler import Profiler
 from mindspore.communication._hccl_management import get_rank_size, get_rank_id
 from mindspore.communication._comm_helper import _create_group_helper, _destroy_group_helper
 from mindspore.communication.management import _set_rank_from_mpi, init as cinit, release as crelease
@@ -360,6 +361,7 @@ FUNC_KEY_DICT_ITEMS = 22  # dict.items
 FUNC_KEY_PRIMITIVE_ASSIGN = 23  # mindspore.ops.assign, Primitive("Assign")
 FUNC_KEY_TENSOR_SETITEM = 24  # Tensor.__setitem__
 FUNC_KEY_TENSOR_ASSIGN_VALUE = 25  # Tensor.assign_value
+FUNC_KEY_TENSOR_IS_CONTIGUOUS = 26  # Tensor.is_contiguous
 
 # Initialized only once. This map will initialize by c++ when start pijit.
 # key is customer if fuzzy match. (Primitive, constexpr, primexpr, MetaFuncGraph)
@@ -376,19 +378,19 @@ _func_map = {
     constexpr_key: FUNC_KEY_CONSTEXPR,
     primexpr_key: FUNC_KEY_PRIMEXPR,
     meta_func_graph_key: FUNC_KEY_META_FUNCG_RAPH,
-    id(GraphCell.__call__): FUNC_KEY_GRAPH_CELL,
+    function_id(GraphCell.__call__): FUNC_KEY_GRAPH_CELL,
     id(psjit_code): FUNC_KEY_PSJIT_CODE,
-    id(_get_cache_prim): FUNC_KEY_GET_CACHE_PRIM,
-    id(Registry.get): FUNC_KEY_REGISTRY_GET,
+    function_id(_get_cache_prim): FUNC_KEY_GET_CACHE_PRIM,
+    function_id(Registry.get): FUNC_KEY_REGISTRY_GET,
 
     # tensor side-effect
     primitive_assign_key: FUNC_KEY_PRIMITIVE_ASSIGN,
-    id(F.assign): FUNC_KEY_PRIMITIVE_ASSIGN,
-    id(Tensor.assign_value): FUNC_KEY_TENSOR_ASSIGN_VALUE,
-    id(Tensor.__setitem__): FUNC_KEY_TENSOR_SETITEM,
+    function_id(F.assign): FUNC_KEY_PRIMITIVE_ASSIGN,
+    function_id(Tensor.assign_value): FUNC_KEY_TENSOR_ASSIGN_VALUE,
+    function_id(Tensor.__setitem__): FUNC_KEY_TENSOR_SETITEM,
 
     # Tensor method
-    id(Tensor.astype): FUNC_KEY_TENSOR_ASTYPE,
+    function_id(Tensor.astype): FUNC_KEY_TENSOR_ASTYPE,
 
     # types.BuiltinFunctionType
     function_id(isinstance): FUNC_KEY_BUILTIN_FUNC,
@@ -448,6 +450,7 @@ _func_map = {
     function_id(str.isalnum): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.isidentifier): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.isprintable): FUNC_KEY_BUILTIN_FUNC,
+    function_id(str.replace): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.format): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.format_map): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.__format__): FUNC_KEY_BUILTIN_FUNC,
@@ -472,7 +475,7 @@ _func_map = {
     function_id(Tensor_.getitem_index_info): FUNC_KEY_BUILTIN_FUNC,
     function_id(Tensor_.get_bytes): FUNC_KEY_BUILTIN_FUNC,
     function_id(Tensor_.is_init): FUNC_KEY_BUILTIN_FUNC,
-    function_id(Tensor_.is_contiguous): FUNC_KEY_BUILTIN_FUNC,
+    function_id(Tensor_.is_contiguous): FUNC_KEY_TENSOR_IS_CONTIGUOUS,
     function_id(Tensor_.stride): FUNC_KEY_BUILTIN_FUNC,
     # Tensor_.asnumpy need real tensor value
 
@@ -488,6 +491,7 @@ _func_map = {
     function_id(validator.check_number_range): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(validator.check_is_int): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(validator.check_is_number): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(validator.check_positive_int_sequence): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(np_version_valid): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(_is_initialized): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(_set_elegant_exit_handle): FUNC_KEY_PIJIT_CONSTEXPR,
@@ -496,7 +500,9 @@ _func_map = {
     function_id(get_rank_size): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(get_rank_id): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(offload_context): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(_is_in_data_parallel_mode): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(check_version_and_env_config): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(Tensor.tolist): FUNC_KEY_PIJIT_CONSTEXPR,
 
     # inner function
     function_id(type_size_in_bytes): FUNC_KEY_BUILTIN_FUNC,
@@ -530,7 +536,6 @@ _func_map = {
     function_id(_exec_save): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(load): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(export_split_mindir): FUNC_KEY_PIJIT_FORBIDDEN,
-    function_id(obfuscate_model): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(_parse_ckpt_proto): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(_generate_front_info_for_param_data_file): FUNC_KEY_PIJIT_FORBIDDEN,
     function_id(_get_data_file): FUNC_KEY_PIJIT_FORBIDDEN,

mindspore/_extends/pijit/tensor_func_list.py

@@ -0,0 +1,27 @@
+# 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.
+# You may obtain a copy of the License at
+#
+# 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.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Store and get tensor method"""
+from mindspore import Tensor
+from mindspore._c_expression import function_id
+
+tensor_method_id_to_name = {}
+for method_name in dir(Tensor):
+    method_id = function_id(getattr(Tensor, method_name))
+    tensor_method_id_to_name[method_id] = method_name
+
+
+def get_tensor_method_name(id):
+    """Get method name by function id"""
+    return tensor_method_id_to_name.get(id, None)

mindspore/_extends/utils.py

@@ -28,7 +28,7 @@ def cell_attr_register(fn=None, attrs=None):
 
     Args:
         fn (function): __init__ function of cell.
-        attrs (list(string) | string): attr list.
+        attrs (list(str) | str): attr list.
 
     Returns:
         function, original function.

mindspore/amp.py

@@ -69,6 +69,12 @@ def _enable_all_finite():
         if not checker.check_custom_version():
             logger.debug("Disable AllFinite due to version check failure.")
             return False
+    else:
+        return False
+
+    if "RANK_TABLE_FILE" in os.environ:
+        return False
+
     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):
@@ -82,7 +88,7 @@ def _enable_all_finite():
     if global_jit_config:
         logger.debug("Current global jit config is: {}".format(global_jit_config["jit_level"]))
         return global_jit_config["jit_level"] == "O0" or global_jit_config["jit_level"] == "O1"
-    return False
+    return True
 
 
 def _grad_unscale(scale, grad):
@@ -93,12 +99,12 @@ def _grad_scale(scale, grad):
     return grad * scale.astype(grad.dtype)
 
 
-@jit
+@jit(backend="ms_backend")
 def _grad_scale_map(scale_value, inputs):
     return _hypermap(_partial(_grad_scale, scale_value), inputs)
 
 
-@jit
+@jit(backend="ms_backend")
 def _grad_unscale_map(scale_value, inputs):
     return _hypermap(_partial(_grad_unscale, scale_value), inputs)
 
@@ -110,7 +116,7 @@ def _overflow(inputs):
     return 1 - status.all()
 
 
-@jit
+@jit(backend="ms_backend")
 def _all_finite(inputs, check_overflow_mode, enable_allfinite):
     """all finite check"""
     if _ascend_target():
@@ -319,7 +325,7 @@ class StaticLossScaler(LossScaler):
 
 class DynamicLossScaler(LossScaler):
     r"""
-    Dynamic Loss scale class.
+    Manager for dynamically adjusting the loss scaling factor.
 
     Dynamic loss scaling tries to determine the largest loss scale value that
     will keep gradients finite. It does this by increasing the loss scale every

mindspore/boost/__init__.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 # ============================================================================
 """
-Boost provide auto accelerating for network, such as Less BN, Gradient Freeze, Gradient
-accumulation and so on.
+Boost is able to automatically optimize network performance, e.g., by reducing BN, gradient freezing,
+and accumulating gradients to achieve network acceleration.
 
 Note:
     This feature is a beta feature, and we are still improving its functionality.

mindspore/boost/base.py

@@ -21,15 +21,12 @@ import math
 import copy
 import numpy as np
 from scipy import linalg as la
-from mindspore.context import ParallelMode
 import mindspore.nn as nn
 from mindspore.nn.optim import LARS
 from mindspore import log as logger
 from mindspore.common import Parameter
-from mindspore.communication.management import get_group_size
+from mindspore.communication.management import get_rank, get_group_size
 from mindspore.train.serialization import load_checkpoint
-from mindspore.parallel._utils import _get_global_rank
-from mindspore.parallel._auto_parallel_context import auto_parallel_context
 from mindspore.boost.less_batch_normalization import CommonHeadLastFN
 
 
@@ -329,7 +326,7 @@ def _get_local_pca_mat_path(weight_load_dir, pca_mat_path, n_component, device_n
     if os.path.exists(save_pca_end_path):
         os.remove(save_pca_end_path)
 
-    rank = _get_global_rank()
+    rank = get_rank()
     local_pca_mat_path = full_pca_mat_path[:-4] + "_rank_" + str(rank) + ".npy"
     if os.path.exists(local_pca_mat_path):
         os.remove(local_pca_mat_path)
@@ -498,8 +495,7 @@ def _save_local_pca_mat(pca_mat, full_pca_mat_path, n_component):
         full_pca_mat_path (str): the path of full pca mat.
         n_component (int): pca component.
     """
-    parallel_mode = auto_parallel_context().get_parallel_mode()
-    rank_size = 1 if parallel_mode == ParallelMode.STAND_ALONE else get_group_size()
+    rank_size = get_group_size()
     local_dim = math.ceil(n_component // rank_size)
     for rank_id in range(rank_size):
         start_index = rank_id * local_dim

mindspore/boost/boost_cell_wrapper.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
+# Copyright 2021-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.
@@ -15,12 +15,13 @@
 """Boost Mode Cell Wrapper."""
 from __future__ import absolute_import
 
+import os
 import numpy as np
 from mindspore.nn.wrap import TrainOneStepCell
 import mindspore.context as context
 from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_global_rank, _get_device_num, _get_gradients_mean
-from mindspore.communication.management import get_group_size, create_group
+from mindspore.communication.management import get_rank, get_group_size, create_group
 from mindspore.nn.cell import Cell
 from mindspore.nn import SequentialCell
 from mindspore.common import Tensor
@@ -38,6 +39,10 @@ from mindspore.boost.adasum import AdaSum
 from mindspore.boost.dim_reduce import DimReduce
 from mindspore.boost.grad_accumulation import gradient_accumulation_op, gradient_clear_op
 from mindspore.boost.base import _load_local_pca_mat
+from mindspore.ops.operations.nn_ops import AllFinite
+from mindspore._c_expression import MSContext
+from mindspore.run_check._check_version import AscendEnvChecker
+from mindspore import log as logger
 
 __all__ = ["BoostTrainOneStepCell", "BoostTrainOneStepWithLossScaleCell"]
 
@@ -90,6 +95,27 @@ def _tensor_grad_overflow(grad):
 def _tensor_grad_overflow_row_tensor(grad):
     return grad_overflow(grad.values)
 
+_ascend_grad_overflow = C.MultitypeFuncGraph("_ascend_grad_overflow")
+ascend_grad_overflow = P.IsFinite()
+
+
+@_ascend_grad_overflow.register("Tensor")
+def _tensor_ascend_grad_overflow(grad):
+    status = ascend_grad_overflow(grad)
+    base = Tensor(1.0, dtype=mstype.float32)
+    output = base - status.all()
+    output = P.Reshape()(output, ((-1,)))
+    return output
+
+
+@_ascend_grad_overflow.register("RowTensor")
+def _tensor_ascend_grad_overflow_row_tensor(grad):
+    status = ascend_grad_overflow(grad.values)
+    base = Tensor(1.0, dtype=mstype.float32)
+    output = base - status.all()
+    output = P.Reshape()(output, ((1,)))
+    return output
+
 
 class _OutputToFloat16(Cell):
     "Wrap cell for amp. Cast network output back to float16"
@@ -362,7 +388,7 @@ class BoostTrainOneStepCell(TrainOneStepCell):
         gamma = self.auto_boost.gamma
         alpha = self.auto_boost.alpha
         sigma = self.auto_boost.sigma
-        _rank = _get_global_rank()
+        _rank = get_rank()
         _rank_size = 1 if self.parallel_mode == ParallelMode.STAND_ALONE else get_group_size()
         n_components = self.auto_boost.n_components
         timeout = self.auto_boost.timeout
@@ -483,7 +509,11 @@ class BoostTrainOneStepWithLossScaleCell(BoostTrainOneStepCell):
         self.allreduce = P.AllReduce()
         self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
         self.gpu_target = (context.get_context("device_target") == "GPU")
+        self.ascend_910a_target = (MSContext.get_instance().get_ascend_soc_version() == 'ascend910')
+        self.ascend_910b_target = (MSContext.get_instance().get_ascend_soc_version() in ['ascend910b', 'ascend910_93'])
         self.loss_scaling_manager = None
+        self._ascend_check_overflow_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
+
         self.base0 = Tensor(0, mstype.int32)
         self.reduce_all = P.ReduceAll(keep_dims=False)
         self.logic_not = P.LogicalNot()
@@ -512,6 +542,26 @@ class BoostTrainOneStepWithLossScaleCell(BoostTrainOneStepCell):
             else:
                 raise TypeError("The scale_sense must be Cell or Tensor, but got {}".format(type(scale_sense)))
 
+        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):
+            logger.debug("Enable AllFinite through the environment variable MS_DEV_RUNTIME_CONF.")
+            self.enable_allfinite = True
+        elif runtime_conf is not None and ("all_finite:False" in runtime_conf or "all_finite:false" in runtime_conf):
+            logger.debug("Disable AllFinite through the environment variable MS_DEV_RUNTIME_CONF.")
+            self.enable_allfinite = False
+        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():
+                logger.debug("Disable AllFinite due to version check failure.")
+                self.enable_allfinite = False
+
     def construct(self, *inputs):
         weights = self.weights
         loss = self.network(*inputs)
@@ -523,7 +573,7 @@ class BoostTrainOneStepWithLossScaleCell(BoostTrainOneStepCell):
             cond, scaling_sens = self._enhanced_amp_process_overflow_status(grads)
         else:
             scaling_sens = self.scale_sense
-            status, scaling_sens = self._start_overflow_check(loss, scaling_sens)
+            status = Tensor([0] * 8, mstype.int32)
             scaling_sens_filled = C.ones_like(loss) * F.cast(scaling_sens, F.dtype(loss))
 
             grads = self.grad(self.network, weights)(*inputs, scaling_sens_filled)
@@ -646,54 +696,99 @@ class BoostTrainOneStepWithLossScaleCell(BoostTrainOneStepCell):
             compute_input = F.depend(compute_input, clear_status)
         return status, compute_input
 
+    def _check_overflow_status_on_infnan_mode(self, grad_overflow_check_func, compute_output):
+        """check overflow status on infnan mode."""
+        flag_sum = self.hyper_map(F.partial(grad_overflow_check_func), compute_output)
+        flag_sum = P.AddN()(flag_sum)
+        # convert flag_sum to scalar
+        flag_sum = P.Reshape()(flag_sum, (()))
+        return flag_sum
+
+    def _get_distributed_overflow_status_on_infnan_mode(self, grad_overflow_check_func, compute_output):
+        """converge the distributed overflow status on infnan mode."""
+        flag_sum = self._check_overflow_status_on_infnan_mode(grad_overflow_check_func, compute_output)
+
+        if self.is_distributed:
+            # sum overflow flag over devices
+            flag_reduce = self.allreduce(flag_sum)
+            overflow = self.less_equal(self.base, flag_reduce)
+        else:
+            overflow = self.less_equal(self.base, flag_sum)
+        return overflow
+
+    def _get_distributed_overflow_status_on_infnan_enable_allfinite(self, compute_output):
+        """check overflow status on infnan kernel mode."""
+        overflow = AllFinite()(compute_output)
+
+        if self.is_distributed:
+            overflow = P.Cast()(overflow, mstype.int8)
+            overflow = P.Cast()(self.allreduce(overflow), mstype.bool_)
+        return overflow
+
+    def _get_gpu_overflow_status(self, compute_output):
+        """get overflow status of gpu."""
+        overflow = self._get_distributed_overflow_status_on_infnan_mode(_grad_overflow, compute_output)
+        return overflow
+
+    def _get_ascend_overflow_status_on_infnan_mode(self, compute_output):
+        """get overflow status of ascend on infnan mode."""
+        overflow = False
+        if self.enable_allfinite:
+            overflow = self._get_distributed_overflow_status_on_infnan_enable_allfinite(compute_output)
+        else:
+            overflow = self._get_distributed_overflow_status_on_infnan_mode(_ascend_grad_overflow, compute_output)
+        return overflow
+
+    def _get_ascend_overflow_status_on_saturation_mode(self, status, compute_output):
+        """get overflow status of ascend on saturation mode"""
+        status = F.depend(status, compute_output)
+        get_status = NPUGetFloatStatusV2()(status)
+
+        if self.is_distributed:
+            # sum overflow flag over devices
+            flag_reduce = self.allreduce(get_status)
+            # get_status not equal to [0]*8 means overflow
+            flag = self.equal(self.base0, flag_reduce)
+            status = F.depend(status, flag)
+            # distributed needs to skip allreduce to avoid its overflow affecting the next step
+            clear_status = NPUClearFloatStatusV2()(status)
+            flag = F.depend(flag, clear_status)
+            overall_finite = self.reduce_all(flag)
+        else:
+            status = F.depend(status, get_status)
+            clear_status = NPUClearFloatStatusV2()(status)
+            get_status = F.depend(get_status, clear_status)
+            flag = self.equal(self.base0, get_status)
+            overall_finite = self.reduce_all(flag)
+        overflow = self.logic_not(overall_finite)
+        return overflow
+
+
     def _get_overflow_status(self, status, compute_output):
         """
         Get floating-point overflow status.
 
-        Get overflow results after executing the target process for overflow detection.
+        Get overflow results after executing the target process for overflow detection. User-defined training network
+        based on this class can also call this interface to process the overflow.
 
-        Inputs:
-            - **status** (object) - A status instance used to detect the overflow.
-            - **compute_output** - Overflow detection should be performed on a certain computation. Set `compute_output`
-              as the output of the computation, to ensure overflow status is acquired before executing the
-              computation.
+        Args:
+            status (object): To control the execution sequence with start_overflow_check, it should be set as the first
+              output of start_overflow_check.
+            compute_output: Overflow detection should be performed in a certain computation process. Set
+              `compute_output` as the output of the computation process.
 
-        Outputs:
+        Returns:
             bool, whether the overflow occurs or not.
         """
-        if not self.gpu_target:
-            status = F.depend(status, compute_output)
-            get_status = NPUGetFloatStatusV2()(status)
-
-            if self.is_distributed:
-                # sum overflow flag over devices
-                flag_reduce = self.allreduce(get_status)
-                # get_status not equal to [0]*8 means overflow
-                flag = self.equal(self.base0, flag_reduce)
-                status = F.depend(status, flag)
-                # distributed needs to skip allreduce to avoid its overflow affecting the next step
-                clear_status = NPUClearFloatStatusV2()(status)
-                flag = F.depend(flag, clear_status)
-                overall_finite = self.reduce_all(flag)
-            else:
-                status = F.depend(status, get_status)
-                clear_status = NPUClearFloatStatusV2()(status)
-                get_status = F.depend(get_status, clear_status)
-                flag = self.equal(self.base0, get_status)
-                overall_finite = self.reduce_all(flag)
-            overflow = self.logic_not(overall_finite)
-        else:
-            flag_sum = self.hyper_map(F.partial(_grad_overflow), compute_output)
-            flag_sum = P.AddN()(flag_sum)
-            # convert flag_sum to scalar
-            flag_sum = P.Reshape()(flag_sum, (()))
-
-            if self.is_distributed:
-                # sum overflow flag over devices
-                flag_reduce = self.allreduce(flag_sum)
-                overflow = self.less_equal(self.base, flag_reduce)
+        if self.gpu_target:
+            overflow = self._get_gpu_overflow_status(compute_output)
+        elif self.ascend_910b_target:
+            if self._ascend_check_overflow_mode == "SATURATION_MODE":
+                overflow = self._get_ascend_overflow_status_on_saturation_mode(status, compute_output)
             else:
-                overflow = self.less_equal(self.base, flag_sum)
+                overflow = self._get_ascend_overflow_status_on_infnan_mode(compute_output)
+        else:  # ascend_910a_target
+            overflow = self._get_ascend_overflow_status_on_saturation_mode(status, compute_output)
         return overflow
 
     def _process_loss_scale(self, overflow):

mindspore/common/__init__.py

@@ -15,7 +15,8 @@
 """Top-level reference to dtype of common module."""
 from __future__ import absolute_import
 from mindspore.common import dtype
-from mindspore.common.api import ms_function, ms_memory_recycle, ms_class, jit, jit_class, _no_grad, flops_collection
+from mindspore.common.api import ms_memory_recycle, jit, jit_class, _no_grad, \
+    flops_collection, set_recursion_limit
 from mindspore.common.dtype import Type, int8, byte, int16, short, int32, intc, int64, intp, \
     uint8, ubyte, uint16, ushort, uint32, uintc, uint64, uintp, float16, half, \
     float32, single, float64, bfloat16, double, bool_, float_, list_, tuple_, int_, \
@@ -38,6 +39,7 @@ from mindspore.common import generator
 from mindspore.common.generator import (
     Generator, default_generator, seed, manual_seed, initial_seed, get_rng_state, set_rng_state)
 from mindspore.ops.function.array_func import is_tensor, from_numpy
+from mindspore.common._grad_function import _Function
 
 # symbols from dtype
 __all__ = [
@@ -69,18 +71,19 @@ __all__ = [
 
 __all__.extend([
     "tensor", "Tensor", "RowTensor", "SparseTensor", "COOTensor", "CSRTensor",  # tensor
-    "ms_function", "ms_class", 'jit', 'jit_class', '_no_grad',  # api
+    'jit', 'jit_class', '_no_grad',  # api
     "Parameter", "ParameterTuple",  # parameter
     "dtype",
     "set_seed", "get_seed", "manual_seed", # random seed
     "set_dump",
     "ms_memory_recycle",
+    "set_recursion_limit",
     "mutable", "JitConfig",
     "flops_collection",
     "lazy_inline", "load_mindir", "save_mindir",
     "no_inline",
     "Symbol",
     "recompute",
-    "is_tensor", "from_numpy",
+    "is_tensor", "from_numpy", "_Function"
 ])
 __all__.extend(generator.__all__)

mindspore/common/_grad_function.py

@@ -0,0 +1,56 @@
+# 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.
+# You may obtain a copy of the License at
+#
+# 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.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Defines custom autograd function with functional form."""
+
+from typing import Any
+from mindspore._c_expression import FunctionBase as FunctionBase_
+from mindspore.common.tensor import Tensor
+
+__all__ = ['_Function']
+
+class _Function(FunctionBase_):
+    """
+    A Class provides the ability to custom autograd function.
+
+    Note:
+        It is only supported in pynative mode.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+    """
+    @staticmethod
+    def forward(ctx: Any, *args: Any, **kwars: Any) -> Any:
+        raise NotImplementedError("forward function should be customized.")
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any) -> Any:
+        raise NotImplementedError("backward function should be customized.")
+
+    @classmethod
+    def apply(cls, *args, **kwargs):
+        return super().apply(cls, *args, **kwargs)
+
+    def save_for_backward(self, *tensors: Tensor):
+        self.saved_tensors = tensors
+
+    def mark_dirty(self, *args: Tensor):
+        self.dirty_tensors = args
+
+    def mark_non_differentiable(self, *args: Tensor):
+        self.non_differentiable = args
+
+    def set_materialize_grads(self, value: bool):
+        self.materialize_grads = value