mindstudio-probe 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

msprobe/pytorch/debugger/precision_debugger.py

@@ -19,8 +19,9 @@ import torch
 from msprobe.core.common.const import Const, FileCheckConst, MsgConst
 from msprobe.core.common.exceptions import MsprobeException
 from msprobe.core.common.file_utils import FileChecker
-from msprobe.core.common.utils import get_real_step_or_rank
+from msprobe.core.common.utils import get_real_step_or_rank, check_init_step
 from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.common.utils import check_save_param
 from msprobe.pytorch.debugger.debugger_config import DebuggerConfig
 from msprobe.pytorch.dump.module_dump.module_dump import ModuleDumper
 from msprobe.pytorch.grad_probe.grad_monitor import GradientMonitor
@@ -158,6 +159,28 @@ class PrecisionDebugger:
             return
         cls._instance.gm.monitor(model)
 
+    @classmethod
+    def save(cls, variable, name, save_backward=True):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if instance.task not in [Const.TENSOR, Const.STATISTICS] or instance.config.level != Const.LEVEL_DEBUG:
+            return
+        try:
+            check_save_param(variable, name, save_backward)
+        except ValueError:
+            return
+        instance.service.save(variable, name, save_backward)
+
+    @classmethod
+    def set_init_step(cls, step):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        check_init_step(step)
+        instance.service.init_step = step
+        instance.service.loop = 0
+
 
 def module_dump(module, dump_name):
     if not isinstance(module, torch.nn.Module):

msprobe/pytorch/dump/module_dump/module_dump.py

@@ -17,7 +17,7 @@ import torch
 from msprobe.core.common.const import Const
 from msprobe.core.data_dump.scope import BaseScope
 from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.hook_module.api_registry import api_register
+from msprobe.pytorch.hook_module.api_register import get_api_register
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
@@ -26,13 +26,14 @@ class ModuleDumper:
     def __init__(self, service):
         self.service = service
         self.hook_handle_list = []
+        self.api_register = get_api_register()
 
     def start_module_dump(self, module, dump_name):
-        api_register.api_originality()
+        self.api_register.restore_all_api()
         self.register_hook(module, dump_name)
 
     def stop_module_dump(self):
-        api_register.api_modularity()
+        self.api_register.register_all_api()
         for hook_handle in self.hook_handle_list:
             if isinstance(hook_handle, torch.utils.hooks.RemovableHandle):
                 hook_handle.remove()

msprobe/pytorch/dump/module_dump/module_processer.py

@@ -16,14 +16,17 @@
 from functools import wraps
 
 import torch
+from torch.utils.hooks import BackwardHook
+
 from msprobe.core.common.const import Const
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.core.data_dump.scope import BaseScope, ModuleRangeScope, MixRangeScope
 from msprobe.pytorch.common.log import logger
-from torch.utils.checkpoint import checkpoint as origin_checkpoint
-from torch.utils.checkpoint import set_checkpoint_early_stop
-from torch.utils.hooks import BackwardHook
+from msprobe.pytorch.common.utils import replace_last_occurrence, is_float8_tensor
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
+if torch_version_above_or_equal_2:
+    from torch.utils.checkpoint import checkpoint as origin_checkpoint, set_checkpoint_early_stop
 
 
 def checkpoint_without_early_stop(*args, **kwargs):
@@ -32,7 +35,8 @@ def checkpoint_without_early_stop(*args, **kwargs):
 
 
 def replace_checkpoint():
-    torch.utils.checkpoint.checkpoint = checkpoint_without_early_stop
+    if torch_version_above_or_equal_2:
+        torch.utils.checkpoint.checkpoint = checkpoint_without_early_stop
 
 
 class ModuleProcesser:
@@ -45,29 +49,8 @@ class ModuleProcesser:
         self.scope = scope if isinstance(scope, (ModuleRangeScope, MixRangeScope)) else None
         BackwardHook.setup_input_hook = ModuleProcesser.clone_return_value(BackwardHook.setup_input_hook)
         BackwardHook.setup_output_hook = ModuleProcesser.clone_return_value(BackwardHook.setup_output_hook)
-        BackwardHook.setup_output_hook = ModuleProcesser.filter_tensor_and_tuple(BackwardHook.setup_output_hook)
         replace_checkpoint()
 
-    @staticmethod
-    def filter_tensor_and_tuple(func):
-        @wraps(func)
-        def wrap_by_filter_tensor_and_tuple(*args, **kwargs):
-            # setup_output_hook传入非tensor数据，工具后续dump会报错，处理方式是解析非tensor数据的属性，对tensor属性挂hook
-            # setup_output_hook定义为setup_output_hook(self, args)，因此处理第二个位置参数，即*args[1]
-            if not isinstance(args[1], (torch.Tensor, tuple)):
-                for item_str in dir(args[1]):
-                    item = getattr(args[1], item_str)
-                    # 处理tensor或者只包含tensor的元组
-                    if isinstance(item, torch.Tensor) or \
-                            (isinstance(item, tuple) and all(isinstance(x, torch.Tensor) for x in item)):
-                        args_new = (args[0], item)
-                        result = func(*args_new, **kwargs)
-                        setattr(args[1], item_str, result)
-                return args[1]
-            return func(*args, **kwargs)
-
-        return wrap_by_filter_tensor_and_tuple
-
     @staticmethod
     def clone_return_value(func):
         @wraps(func)
@@ -78,14 +61,15 @@ class ModuleProcesser:
         return clone_return_value_func
 
     @staticmethod
+    @recursion_depth_decorator("ModuleDump: ModuleProcesser.clone_if_tensor", max_depth=Const.DUMP_MAX_DEPTH)
     def clone_if_tensor(result):
-        if isinstance(result, torch.Tensor):
+        if isinstance(result, torch.Tensor) and not is_float8_tensor(result):
             return result.clone()
-        elif isinstance(result, tuple):
+        elif type(result) is tuple:
             return tuple(ModuleProcesser.clone_if_tensor(x) for x in result)
-        elif isinstance(result, list):
+        elif type(result) is list:
             return list(ModuleProcesser.clone_if_tensor(x) for x in result)
-        elif isinstance(result, dict):
+        elif type(result) is dict:
             return {k: ModuleProcesser.clone_if_tensor(v) for k, v in result.items()}
         else:
             return result
@@ -103,7 +87,7 @@ class ModuleProcesser:
         return hasattr(module, '_backward_hooks') and \
             len(module._backward_hooks) > 0 and \
             module._is_full_backward_hook is False
-    
+
     @staticmethod
     def get_modules_and_names(models):
         modules_and_names_with_index = {}
@@ -129,9 +113,11 @@ class ModuleProcesser:
             for name, module in modules_and_names:
                 if module == model:
                     continue
+                if module.__class__.__name__ == "FullyShardedDataParallel":
+                    continue
                 module_index = (index + Const.SEP) if index != "-1" else ""
-                prefix_name = (BaseScope.Module_Type_Module + Const.SEP + module_index + 
-                                name + Const.SEP + module.__class__.__name__ + Const.SEP)
+                prefix_name = (BaseScope.Module_Type_Module + Const.SEP + module_index +
+                               name + Const.SEP + module.__class__.__name__ + Const.SEP)
                 pre_forward_hook, forward_hook, backward_hook, forward_hook_torch_version_below_2 = build_hook(
                     BaseScope.Module_Type_Module,
                     prefix_name
@@ -203,9 +189,9 @@ class ModuleProcesser:
             if not hasattr(module, "mindstudio_reserved_name") or not module.mindstudio_reserved_name:
                 module.mindstudio_reserved_name = []
             module.mindstudio_reserved_name.append(full_name)
-            forward_full_name = full_name.replace(Const.BACKWARD, Const.FORWARD)
-            ModuleProcesser.module_node[full_name] = ModuleProcesser.module_node[forward_full_name].replace(
-                Const.FORWARD, Const.BACKWARD) if ModuleProcesser.module_node[forward_full_name] else None
+            forward_full_name = replace_last_occurrence(full_name, Const.BACKWARD, Const.FORWARD)
+            ModuleProcesser.module_node[full_name] = replace_last_occurrence(
+                ModuleProcesser.module_node.get(forward_full_name), Const.FORWARD, Const.BACKWARD)
             ModuleProcesser.api_parent_node = None
             if self.scope:
                 self.scope.begin_module(full_name)

msprobe/pytorch/free_benchmark/common/utils.py

@@ -16,7 +16,7 @@
 
 import torch
 from msprobe.core.common.exceptions import FreeBenchmarkException
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark.common.enums import DeviceType
 
 

msprobe/pytorch/free_benchmark/compare/single_benchmark.py

@@ -16,7 +16,7 @@
 import math
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
 from msprobe.pytorch.free_benchmark.common.utils import TorchC

msprobe/pytorch/free_benchmark/perturbed_layers/npu/add_noise.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode
@@ -95,13 +95,13 @@ class AddNoiseLayer(NpuBaseLayer):
         except Exception:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"when calculate maximun value, tensor is changed to float32."
+                f"when calculating the maximum value, the tensor is changed to float32."
             )
             max_val = TorchC.max(TorchC.abs(tensor_obj.to(torch.float32))).item()
         if max_val < abs_tol:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"Maximun value is less than the minimun threshold. Cancel add noise."
+                f"maximum value is less than the minimum threshold. Cancel adding noise."
             )
             return False
         return True

msprobe/pytorch/free_benchmark/perturbed_layers/npu/bit_noise.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode
@@ -100,13 +100,13 @@ class BitNoiseLayer(NpuBaseLayer):
         except Exception:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"when calculate maximun value, tensor is changed to float32."
+                f"when calculate the maximum value, the tensor is changed to float32."
             )
             max_val = TorchC.max(TorchC.abs(tensor_obj.to(torch.float32))).item()
         if max_val < abs_tol:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"Maximun value is less than the minimun threshold. Cancel add noise."
+                f"maximum value is less than the minimum threshold. Cancel adding noise."
             )
             return False
         return True

msprobe/pytorch/free_benchmark/perturbed_layers/npu/change_value.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode
 from msprobe.pytorch.free_benchmark.common.params import DataParams

msprobe/pytorch/free_benchmark/perturbed_layers/npu/improve_precision.py

@@ -15,7 +15,7 @@
 
 import torch
 from msprobe.core.common.const import Const
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import CommonField
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode

msprobe/pytorch/free_benchmark/result_handlers/check_handler.py

@@ -49,6 +49,6 @@ class CheckerHandler(FuzzHandler):
         except Exception as e:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when campare the result exception raise {e}"
+                f"when comparing the results, an exception is raised: {e}"
             )
         return data_params.original_result

msprobe/pytorch/function_factory.py

@@ -27,6 +27,11 @@ from msprobe.pytorch.bench_functions.rotary_mul import npu_rotary_mul, npu_rotar
 from msprobe.pytorch.bench_functions.scaled_mask_softmax import npu_scaled_masked_softmax, \
     npu_scaled_masked_softmax_backward
 from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward
+from msprobe.pytorch.bench_functions.apply_adam import npu_apply_adam
+from msprobe.pytorch.bench_functions.group_norm_silu import npu_group_norm_silu
+from msprobe.pytorch.bench_functions.mish import npu_mish
+from msprobe.pytorch.bench_functions.moe_gating_top_k_softmax import npu_moe_gating_top_k_softmax
+from msprobe.pytorch.bench_functions.sort_v2 import npu_sort_v2
 from msprobe.pytorch.common.utils import logger
 
 
@@ -65,7 +70,7 @@ class Register(dict):
 
         def add_register_item(key, value):
             if key in self._dict:
-                logger.warning(f"{value.__name__} has been registered before, so we will overriden it.")
+                logger.warning(f"{value.__name__} has been registered before, so we will override it.")
             self[key] = value
             return value
 
@@ -79,7 +84,8 @@ class Register(dict):
 npu_custom_functions = Register()
 npu_custom_functions([
     npu_apply_adam_w, npu_confusion_transpose, npu_fast_gelu, npu_layer_norm_eval, npu_linear, npu_fusion_attention,
-    npu_rms_norm, npu_rotary_mul, npu_scaled_masked_softmax, npu_swiglu, gpu_fusion_attention
+    npu_rms_norm, npu_rotary_mul, npu_scaled_masked_softmax, npu_swiglu, gpu_fusion_attention, npu_apply_adam,
+    npu_group_norm_silu, npu_mish, npu_moe_gating_top_k_softmax, npu_sort_v2
 ])
 
 # register for npu custom backward bench functions

msprobe/pytorch/grad_probe/grad_monitor.py

@@ -46,7 +46,7 @@ class GradientMonitor:
         if not os.path.exists(self._output_path):
             create_directory(self._output_path)
         else:
-            logger.warning(f"the file in {self._output_path} will be recoverd")
+            logger.warning(f"the file in {self._output_path} will be deleted")
         self._step = -1
         self._param2name = defaultdict(str)
 
@@ -97,7 +97,7 @@ class GradientMonitor:
                 create_directory(output_dirpath)
             output_path = os.path.join(output_dirpath, f"grad_summary_{self._step}.csv")
             if os.path.exists(output_path):
-                logger.warning(f"{output_path} will be recoverd")
+                logger.warning(f"{output_path} will be deleted")
                 remove_path(output_path)
             header_result = GradStatCsv.generate_csv_header(self._level_adp, self._bounds)
             output_lines.insert(0, header_result)

msprobe/pytorch/hook_module/api_register.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2025-2025, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
+import functools
+import os
+
+import torch
+import torch.distributed as dist
+
+from msprobe.core.common.const import Const
+from msprobe.core.data_dump.api_registry import ApiRegistry
+from msprobe.pytorch.common.utils import (
+    torch_without_guard_version, is_gpu, torch_device_guard, parameter_adapter
+)
+from msprobe.pytorch.function_factory import npu_custom_functions
+from msprobe.pytorch.hook_module.hook_module import HOOKModule
+
+
+torch_version_above_2 = torch.__version__.split('+')[0] > '2.0'
+
+_api_types = {
+    Const.PT_FRAMEWORK: {
+        Const.PT_API_TYPE_FUNCTIONAL: (torch.nn.functional, (torch.nn.functional,)),
+        Const.PT_API_TYPE_TENSOR: (torch.Tensor, (torch.Tensor,)),
+        Const.PT_API_TYPE_TORCH: (torch, (torch,)),
+        Const.PT_API_TYPE_VF: (torch._C._VariableFunctionsClass, (torch._VF,)),
+        Const.PT_API_TYPE_DIST: (dist, (dist, dist.distributed_c10d))
+    }
+}
+if not is_gpu:
+    import torch_npu
+    if torch_without_guard_version:
+        _api_types.get(Const.PT_FRAMEWORK).update(
+            {
+                Const.PT_API_TYPE_NPU: (torch.ops.npu, (torch_npu, torch.ops.npu))
+            }
+        )
+    else:
+        _api_types.get(Const.PT_FRAMEWORK).update(
+            {Const.PT_API_TYPE_NPU: (torch_npu._C._VariableFunctionsClass, (torch_npu,))}
+        )
+        _api_types.get(Const.PT_FRAMEWORK).update(
+            {
+                Const.PT_API_TYPE_NPU_DIST: (torch_npu.distributed, (torch_npu.distributed,
+                                                                     torch_npu.distributed.distributed_c10d))
+            }
+        )
+
+_inner_used_api = {}
+_supported_api_list_path = (os.path.join(os.path.dirname(os.path.realpath(__file__)), Const.SUPPORT_API_FILE_NAME),)
+_cuda_func_mapping = {"npu_fusion_attention": "gpu_fusion_attention"}
+
+
+@parameter_adapter
+def tensor_module_forward(module, *args, **kwargs):
+    return module.api_func(*args, **kwargs)
+
+
+def dist_module_forward(module, *args, **kwargs):
+    handle = module.api_func(*args, **kwargs)
+    if kwargs.get("async_op") or module.api_name in ["isend", "irecv"]:
+        if handle and hasattr(handle, 'wait'):
+            handle.wait()
+    if module.api_name == "batch_isend_irecv":
+        if isinstance(handle, list):
+            for req in handle:
+                req.wait()
+    return handle
+
+
+def npu_module_forward(module, *args, **kwargs):
+    if not module.need_hook:
+        if module.api_name not in npu_custom_functions:
+            raise Exception(f'There is not bench function {module.api_name}')
+        if module.device == Const.CUDA_LOWERCASE:
+            module.api_name = _cuda_func_mapping.get(module.api_name, module.api_name)
+        if module.device in [Const.CUDA_LOWERCASE, Const.CPU_LOWERCASE]:
+            return npu_custom_functions[module.api_name](*args, **kwargs)
+    return module.api_func(*args, **kwargs)
+
+
+forward_methods = {
+    "Tensor": tensor_module_forward,
+    "Distributed": dist_module_forward,
+    "NPU": npu_module_forward
+}
+
+
+class ApiTemplate(HOOKModule):
+    def __init__(self, api_name, api_func, prefix, hook_build_func, need_hook=True, device=Const.CPU_LOWERCASE):
+        self.api_name = api_name
+        self.api_func = api_func
+        self.prefix = prefix
+        self.prefix_api_name = prefix + Const.SEP + str(api_name.split(Const.SEP)[-1]) + Const.SEP
+        self.need_hook = need_hook
+        self.device = device
+        if self.need_hook:
+            super().__init__(hook_build_func)
+        if prefix == Const.DIST_API_TYPE_PREFIX:
+            self.op_is_distributed = True
+
+    @torch_device_guard
+    def forward(self, *args, **kwargs):
+        exec_func = forward_methods.get(self.prefix)
+        exec_func = functools.partial(exec_func, self) if exec_func else self.api_func
+        return exec_func(*args, **kwargs)
+
+
+api_register = None
+
+
+def get_api_register(return_new=False):
+    if return_new:
+        return ApiRegistry(_api_types, _inner_used_api, _supported_api_list_path, ApiTemplate)
+
+    global api_register
+    if api_register is None:
+        api_register = ApiRegistry(_api_types, _inner_used_api, _supported_api_list_path, ApiTemplate)
+    return api_register

msprobe/pytorch/hook_module/hook_module.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -21,6 +21,8 @@ import torch
 import torch.nn as nn
 import torch.utils.hooks as full_hooks
 
+from msprobe.pytorch.common.utils import is_float8_tensor
+
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
 
@@ -28,28 +30,27 @@ class HOOKModule(nn.Module):
     module_count = defaultdict(int)
     inner_stop_hook = {}
 
-    def __init__(self, build_hook) -> None:
+    def __init__(self, hook_build_func) -> None:
         super(HOOKModule, self).__init__()
         self.has_overflow = False
-        self.prefix = ""
         self.current_thread = threading.current_thread().ident
         if self.current_thread not in HOOKModule.inner_stop_hook:
             HOOKModule.inner_stop_hook[self.current_thread] = False
         self.stop_hook = HOOKModule.inner_stop_hook.get(self.current_thread, False)
 
         if not self.stop_hook:
-            if hasattr(self, "prefix_op_name_"):
-                self.prefix = self.prefix_op_name_
-
             self.forward_data_collected = False
-            forward_pre_hook, forward_hook, backward_hook, _ = build_hook(self.prefix)
-            if torch_version_above_or_equal_2:
-                self.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
-                self.register_forward_hook(forward_hook, with_kwargs=True)
-            else:
-                self.register_forward_pre_hook(forward_pre_hook)
-                self.register_forward_hook(forward_hook)
-            self.register_backward_hook(backward_hook)
+
+            prefix = self.prefix_api_name if hasattr(self, "prefix_api_name") else ""
+            if callable(hook_build_func):
+                forward_pre_hook, forward_hook, backward_hook, _ = hook_build_func(prefix)
+                if torch_version_above_or_equal_2:
+                    self.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
+                    self.register_forward_hook(forward_hook, with_kwargs=True)
+                else:
+                    self.register_forward_pre_hook(forward_pre_hook)
+                    self.register_forward_hook(forward_hook)
+                self.register_backward_hook(backward_hook)
 
     def __call__(self, *args, **kwargs):
         changed = False
@@ -111,6 +112,10 @@ class HOOKModule(nn.Module):
                         return result
                 else:
                     return result
+
+            if is_float8_tensor(var) or not (var.requires_grad and torch.is_grad_enabled()):
+                return result
+
             grad_fn = var.grad_fn
             if grad_fn is not None:
                 for hook in non_full_backward_hooks:

msprobe/pytorch/hook_module/register_optimizer_hook.py

@@ -32,8 +32,9 @@ def register_optimizer_hook(data_collector):
     def patch_clip_grad(func):
         def wrapper(*args, **kwargs):
             data_collector.optimizer_status = Const.CLIP_GRAD
-            func(*args, **kwargs)
+            result = func(*args, **kwargs)
             data_collector.optimizer_status = Const.END_PREFIX + Const.CLIP_GRAD
+            return result
 
         return wrapper