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

msprobe/core/data_dump/json_writer.py

@@ -17,9 +17,11 @@ import csv
 import os
 import copy
 import threading
+import traceback
+from datetime import datetime, timezone, timedelta
 
 from msprobe.core.common.const import Const, FileCheckConst
-from msprobe.core.common.file_utils import change_mode, FileOpen, save_json, load_json
+from msprobe.core.common.file_utils import change_mode, FileOpen, save_json, load_json, check_path_before_create
 from msprobe.core.common.log import logger
 from msprobe.core.common.decorator import recursion_depth_decorator
 
@@ -35,11 +37,15 @@ class DataWriter:
         self.free_benchmark_file_path = None
         self.dump_tensor_data_dir = None
         self.debug_file_path = None
+        self.dump_error_info_path = None
         self.flush_size = 1000
+        self.larger_flush_size = 20000
         self.cache_data = {}
         self.cache_stack = {}
         self.cache_construct = {}
         self.cache_debug = {}
+        self.stat_stack_list = []
+        self._error_log_initialized = False
 
     @staticmethod
     def write_data_to_csv(result: list, result_header: tuple, file_path: str):
@@ -56,13 +62,54 @@ class DataWriter:
         if is_new_file:
             change_mode(file_path, FileCheckConst.DATA_FILE_AUTHORITY)
 
+    @recursion_depth_decorator("JsonWriter: DataWriter._replace_stat_placeholders")
+    def _replace_stat_placeholders(self, data, stat_result):
+        if isinstance(data, dict):
+            keys = list(data.keys())  # 获取当前所有键
+            for key in keys:  # 递归所有变量
+                value = data[key]
+                if key == Const.TENSOR_STAT_INDEX and isinstance(value, int):
+                    if value >= 0:
+                        idx = value
+                    else:
+                        return
+                    stat_values = stat_result[idx] if idx < len(stat_result) else [None] * 4
+
+                    new_entries = {
+                        Const.TYPE: data["type"],
+                        Const.DTYPE: data["dtype"],
+                        Const.SHAPE: data["shape"],
+                        Const.MAX: stat_values[0],
+                        Const.MIN: stat_values[1],
+                        Const.MEAN: stat_values[2],
+                        Const.NORM: stat_values[3],
+                    }
+                    del data[key]
+
+                    # 重构字典顺序
+                    updated_dict = {}
+                    # 通过插入排序后字段保证字段写入json的有序
+                    updated_dict.update(new_entries)
+                    # 遍历原字典其他字段（排除已删除的tensor_stat_index）
+                    for k in data:
+                        if k not in new_entries:
+                            updated_dict[k] = data[k]
+                    data.clear()
+                    data.update(updated_dict)
+                else:
+                    self._replace_stat_placeholders(value, stat_result)
+        elif isinstance(data, (list, tuple)):
+            for item in data:
+                self._replace_stat_placeholders(item, stat_result)
+
     def reset_cache(self):
         self.cache_data = {}
         self.cache_stack = {}
         self.cache_construct = {}
+        self.cache_debug = {}
 
     def initialize_json_file(self, **kwargs):
-        if self.debug_file_path and not self.cache_debug:
+        if kwargs["level"] == Const.LEVEL_DEBUG and not self.cache_debug:
             # debug level case only create debug.json
             debug_dict = copy.deepcopy(kwargs)
             debug_dict.update({"dump_data_dir": self.dump_tensor_data_dir, Const.DATA: {}})
@@ -85,12 +132,46 @@ class DataWriter:
         self.dump_tensor_data_dir = dump_path_aggregation.dump_tensor_data_dir
         self.free_benchmark_file_path = dump_path_aggregation.free_benchmark_file_path
         self.debug_file_path = dump_path_aggregation.debug_file_path
+        self.dump_error_info_path = dump_path_aggregation.dump_error_info_path
 
     def flush_data_periodically(self):
         dump_data = self.cache_data.get(Const.DATA)
-        if dump_data and isinstance(dump_data, dict) and len(dump_data) % self.flush_size == 0:
+
+        if not dump_data or not isinstance(dump_data, dict):
+            return
+
+        length = len(dump_data)
+
+        threshold = self.flush_size if length < self.larger_flush_size else self.larger_flush_size
+
+        if length % threshold == 0:
             self.write_json()
 
+    def write_error_log(self, message: str):
+        """
+        写错误日志：
+          - 第一次调用时以 'w' 模式清空文件，之后都用 'a' 模式追加
+          - 添加时间戳
+          - 在 message 后写入当前的调用栈（方便追踪日志来源）
+        """
+        try:
+            mode = "w" if not self._error_log_initialized else "a"
+            self._error_log_initialized = True
+
+            check_path_before_create(self.dump_error_info_path)
+
+            with FileOpen(self.dump_error_info_path, mode) as f:
+                cst_timezone = timezone(timedelta(hours=8), name="CST")
+                timestamp = datetime.now(cst_timezone).strftime("%Y-%m-%d %H:%M:%S %z")
+                f.write(f"[{timestamp}] {message}\n")
+                f.write("Call stack (most recent call last):\n")
+
+                f.write("".join(traceback.format_stack()[:-1]))  # 去掉自己这一层
+                f.write("\n")
+        except Exception as e:
+            # 如果连写日志都失败了，就打印到 stderr
+            logger.warning(f"[FallbackError] Failed to write error log: {e}")
+
     def update_data(self, new_data):
         with lock:
             if not isinstance(new_data, dict) or len(new_data.keys()) != 1:
@@ -107,9 +188,13 @@ class DataWriter:
             else:
                 dump_data.update(new_data)
 
-    def update_stack(self, new_data):
+    def update_stack(self, name, stack_data):
         with lock:
-            self.cache_stack.update(new_data)
+            api_list = self.cache_stack.get(stack_data)
+            if api_list is None:
+                self.cache_stack.update({stack_data: [name]})
+            else:
+                api_list.append(name)
 
     def update_construct(self, new_data):
         with lock:
@@ -124,7 +209,11 @@ class DataWriter:
         save_json(file_path, self.cache_data, indent=1)
 
     def write_stack_info_json(self, file_path):
-        save_json(file_path, self.cache_stack, indent=1)
+        num, new_cache_stack = 0, {}
+        for key, value in self.cache_stack.items():
+            new_cache_stack[num] = [value, key]
+            num += 1
+        save_json(file_path, new_cache_stack, indent=1)
 
     def write_construct_info_json(self, file_path):
         save_json(file_path, self.cache_construct, indent=1)
@@ -132,8 +221,56 @@ class DataWriter:
     def write_debug_info_json(self, file_path):
         save_json(file_path, self.cache_debug, indent=1)
 
+    def append_stat_to_buffer(self, stat_vector):
+        """
+        直接使用 Python list 存储 stat_vector,
+        将 stat_vector 存入 self.stat_stack_list 的方式
+        """
+        self.stat_stack_list.append(stat_vector)
+        return len(self.stat_stack_list) - 1
+
+    def get_buffer_values_max(self, index):
+        if 0 <= index < len(self.stat_stack_list) and len(self.stat_stack_list[index]) >= 1:
+            return self.stat_stack_list[index][0]
+        else:
+            logger.warning(f"stat_stack_list[{index}] The internal data is incomplete,"
+                           f" and the maximum value cannot be obtained.")
+            return None
+
+    def get_buffer_values_min(self, index):
+        if 0 <= index < len(self.stat_stack_list) and len(self.stat_stack_list[index]) >= 1:
+            return self.stat_stack_list[index][1]
+        else:
+            logger.warning(f"stat_stack_list[{index}] Internal data is incomplete"
+                           f" and minimum values cannot be obtained.")
+            return None
+
+    def flush_stat_stack(self):
+        """
+        在 flush 阶段，将所有存储的统计值从设备搬到 CPU，
+        这里返回一个列表，每个元素是 [Max, Min, Mean, Norm] 的数值列表
+        """
+        if not self.stat_stack_list:
+            return []
+        result = [
+            [
+                x.item() if hasattr(x, "item") else x
+                for x in stat_values
+            ]
+            for stat_values in self.stat_stack_list
+        ]
+        self.stat_stack_list = []
+        return result
+
     def write_json(self):
         with lock:
+            # 在写 JSON 前，统一获取统计值
+            stat_result = self.flush_stat_stack()
+            # 遍历 cache_data，将占位符替换为最终统计值
+            if stat_result:
+                self._replace_stat_placeholders(self.cache_data, stat_result)
+                if self.cache_debug:
+                    self._replace_stat_placeholders(self.cache_debug, stat_result)
             if self.cache_data:
                 self.write_data_json(self.dump_file_path)
             if self.cache_stack:
@@ -143,24 +280,3 @@ class DataWriter:
             if self.cache_debug:
                 self.write_debug_info_json(self.debug_file_path)
 
-    def fill_stack_tensor_data(self):
-        self.process_stat_data_recursive(self.cache_data)
-
-    @recursion_depth_decorator("AsyncDump: DataWriter.process_stat_data_recursive", max_depth=Const.DUMP_MAX_DEPTH)
-    def process_stat_data_recursive(self, data):
-        if isinstance(data, dict):
-            if "tensor_stat" in data.keys():
-                tensor_stat = data["tensor_stat"]
-                if len(tensor_stat) != Const.TENSOR_STAT_LEN or len(tensor_stat[0]) != len(tensor_stat[1]):
-                    logger.warning("Some bad data in async dump")
-                else:
-                    tensor_stat_index, tensor_stat_data = tensor_stat[0], tensor_stat[1]
-                    for index, stat in zip(tensor_stat_index, tensor_stat_data):
-                        data.update({index: stat.item()})
-                del data["tensor_stat"]
-            else:
-                for key in data.keys():
-                    self.process_stat_data_recursive(data[key])
-        elif isinstance(data, (list, tuple)):
-            for i in data:
-                self.process_stat_data_recursive(i)

msprobe/core/debugger/precision_debugger.py

@@ -0,0 +1,144 @@
+# 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 os
+
+from msprobe.core.common.const import Const, FileCheckConst, MsgConst
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.common.file_utils import FileChecker, load_json
+from msprobe.core.common.utils import get_real_step_or_rank, check_init_step
+from msprobe.core.common_config import CommonConfig
+
+
+class BasePrecisionDebugger:
+    _instance = None
+    tasks_not_need_debugger = [Const.GRAD_PROBE]
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(BasePrecisionDebugger, cls).__new__(cls)
+            cls._instance.config = None
+            cls._instance.enable_dataloader = False
+            cls._instance.initialized = False
+            cls.service = None
+            cls.first_start = False
+        return cls._instance
+
+    def __init__(
+            self,
+            config_path=None,
+            task=None,
+            dump_path=None,
+            level=None,
+            step=None
+    ):
+        if self.initialized:
+            return
+        self.initialized = True
+        self._check_input_params(config_path, task, dump_path, level)
+        self.common_config, self.task_config = self._parse_config_path(config_path, task)
+        self.task = self.common_config.task
+        if step is not None:
+            self.common_config.step = get_real_step_or_rank(step, Const.STEP)
+
+    @staticmethod
+    def _check_input_params(config_path, task, dump_path, level):
+        if not config_path:
+            config_path = os.path.join(os.path.dirname(__file__), "../../config.json")
+        if config_path is not None:
+            if not isinstance(config_path, str):
+                raise MsprobeException(
+                    MsprobeException.INVALID_PARAM_ERROR, f"config_path must be a string")
+            file_checker = FileChecker(
+                file_path=config_path, path_type=FileCheckConst.FILE, file_type=FileCheckConst.JSON_SUFFIX)
+            file_checker.common_check()
+
+        if task is not None and task not in Const.TASK_LIST:
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"task must be one of {Const.TASK_LIST}")
+
+        if dump_path is not None:
+            if not isinstance(dump_path, str):
+                raise MsprobeException(
+                    MsprobeException.INVALID_PARAM_ERROR, f"dump_path must be a string")
+
+        if level is not None and level not in Const.LEVEL_LIST:
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"level must be one of {Const.LEVEL_LIST}")
+
+    @staticmethod
+    def _get_task_config(task, json_config):
+        raise NotImplementedError("Subclass must implement _get_task_config")
+
+    @classmethod
+    def forward_backward_dump_end(cls):
+        instance = cls._instance
+        instance.stop()
+
+    @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
+
+    @classmethod
+    def register_custom_api(cls, module, api, api_prefix=None):
+        if not api_prefix:
+            api_prefix = getattr(module, "__name__", "Custom")
+        if not isinstance(api_prefix, str):
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, "api_prefix must be string")
+        if not hasattr(module, api):
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"module {str(module)} does not have {api}")
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        instance.service.register_custom_api(module, api, api_prefix)
+
+    @classmethod
+    def restore_custom_api(cls, module, api):
+        if not hasattr(module, api):
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"module {str(module)} does not have {api}")
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        instance.service.restore_custom_api(module, api)
+
+    @classmethod
+    def _get_instance(cls):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if instance.task in BasePrecisionDebugger.tasks_not_need_debugger:
+            instance = None
+        return instance
+
+    def _parse_config_path(self, json_file_path, task):
+        if not json_file_path:
+            json_file_path = os.path.join(os.path.dirname(__file__), "../../config.json")
+        json_config = load_json(json_file_path)
+        common_config = CommonConfig(json_config)
+        if task:
+            task_config = self._get_task_config(task, json_config)
+        else:
+            if not common_config.task:
+                common_config.task = Const.STATISTICS
+            task_config = self._get_task_config(common_config.task, json_config)
+        return common_config, task_config

msprobe/core/grad_probe/constant.py

@@ -52,7 +52,7 @@ class GradConst:
     BOUNDS_MINIMUM = -2**63
     BOUNDS_MAXIMUM = 2**63 - 1
 
-    # file safty
+    # file safety
     DATA_DIR_AUTHORITY = 0o750
     DATA_FILE_AUTHORITY = 0o640
     DIRECTORY_LENGTH = 4096

msprobe/core/grad_probe/grad_compare.py

@@ -121,7 +121,7 @@ class GradComparator:
         similarities = {}
         logger.info(f"{len(steps)} steps will be compared")
         grad_weight_order = cls._get_grad_weight_order(path1, path2)
-        for step in tqdm(steps, desc="culculate similarities (by step)"):
+        for step in tqdm(steps, desc="calculate similarities (by step)"):
             grad_files = cls._get_matched_grad_files(path1, path2, step)
             same_count_summary = 0
             total_count_summary = 0

msprobe/core/grad_probe/utils.py

@@ -82,7 +82,7 @@ class ListCache(list):
         if len(self) == 0:
             return
         if not self._output_file:
-            logger.warning("dumpfile path is not setted")
+            logger.warning("dumpfile path is not set.")
         write_csv(self, self._output_file)
         logger.info(f"write {len(self)} items to {self._output_file}.")
         self.clear()

msprobe/core/hook_manager.py

@@ -0,0 +1,242 @@
+# Copyright (c) 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.
+
+
+from abc import ABC, abstractmethod
+import os
+
+from msprobe.core.common.runtime import Runtime
+from msprobe.core.common.utils import Const
+from msprobe.core.data_dump.data_processor.base import (ModuleBackwardInputsOutputs, ModuleForwardInputsOutputs)
+
+
+class HookSet:
+    def __init__(self, forward_hook=None, forward_pre_hook=None, backward_hook=None, backward_pre_hook=None):
+        self.forward_hook = forward_hook
+        self.forward_pre_hook = forward_pre_hook
+        self.backward_hook = backward_hook
+        self.backward_pre_hook = backward_pre_hook
+
+
+class BaseHookManager(ABC):
+    inner_switch = False
+    hook_handle_dict = {}
+    params_grad_info = {}
+
+    def __init__(self, data_collector, config, attl_manager=None):
+        self.data_collector = data_collector
+        self.config = config
+        self.attl_manager = attl_manager
+
+    @property
+    def _pid(self):
+        return os.getpid()
+
+    @property
+    @abstractmethod
+    def _is_recompute(self):
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _no_grad_context():
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _add_count(name):
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _process_kwargs_and_output(module, hook_type, kwargs_or_output, output_or_kwargs):
+        pass
+
+    @staticmethod
+    def _clear_input_kwargs(module):
+        if hasattr(module, 'msprobe_input_kwargs'):
+            del module.msprobe_input_kwargs
+
+    @abstractmethod
+    def build_hook(self):
+        pass
+
+    @abstractmethod
+    def _get_params_dict(self, module):
+        pass
+
+    @abstractmethod
+    def _need_exchange(self, module):
+        pass
+
+    def _register_param_hook(self, name, module, params_dict):
+        ori_name = name.rsplit(Const.SEP, 2)[0]
+        grad_name = ori_name + Const.SEP + Const.PARAMS_GRAD
+        # 首次执行前向hook时，添加params_grad_name属性，并注册参数hook
+        setattr(module, 'params_grad_name', grad_name)
+         # data_mode为forward时，不注册参数hook
+        if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
+            for param_name, param in params_dict.items():
+                if param.requires_grad:
+                    name = ori_name + Const.SEP + param_name
+                    old_handle = BaseHookManager.hook_handle_dict.get(name)
+                    if old_handle and hasattr(old_handle, "remove"):
+                        old_handle.remove()
+                    handle = param.register_hook(self._build_grad_hook(module, ori_name, param_name))
+                    BaseHookManager.hook_handle_dict[name] = handle
+
+    def _init_params_grad_info(self, module, params_dict):
+        '''
+        初始化参数梯度信息, 在前向hook结束后, 将参数梯度信息写入cache_data中用于占位
+        '''
+        if not params_dict:
+            return
+        if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
+            grad_name = module.params_grad_name if hasattr(module, 'params_grad_name') else None
+            # 判断是否已经在cache_data中进行了占位, 若没有则先写入cache_data中
+            if not BaseHookManager.params_grad_info.get(grad_name):
+                data_info = {grad_name: {key: [None] for key, value in params_dict.items() if value.requires_grad}}
+                # 当模块中的参数有requires_grad属性为True时，才会进行梯度计算，此时才需要占位
+                if data_info.get(grad_name):
+                    # 将grad_name的data_info先写入cache_data中, 梯度计算后再更新
+                    self.data_collector.handle_data(grad_name, data_info,
+                                                    flush=self.data_collector.data_processor.is_terminated)
+                # 记录当前模块的参数梯度信息已占位
+                BaseHookManager.params_grad_info[grad_name] = True
+
+    def _should_execute_hook(self, hook_type, module, is_forward):
+        is_module_hook = hook_type == Const.MODULE
+        if is_module_hook and not Runtime.is_running:
+            return False
+        elif not is_module_hook and is_forward and not Runtime.is_running:
+            return False
+        elif not is_module_hook and not is_forward and not module.forward_data_collected:
+            return False
+        if BaseHookManager.inner_switch:
+            return False
+        if not self.data_collector or self.data_collector.data_processor.is_terminated:
+            return False
+        return True
+
+    def _build_grad_hook(self, module, ori_name, param_name):
+        def hook_fn(grad):
+            if not self._should_execute_hook(Const.MODULE, module, False):
+                return
+            BaseHookManager.inner_switch = True
+            self.data_collector.params_data_collect(ori_name, param_name, self._pid, grad)
+            BaseHookManager.inner_switch = False
+            return
+        return hook_fn
+
+    def _build_forward_pre_hook(self, hook_type, full_name, api_name):
+        def forward_pre_hook(module, args, kwargs=None):
+            if hook_type == Const.MODULE:
+                return
+            if not self._should_execute_hook(hook_type, module, True):
+                return
+            if kwargs is None:
+                kwargs = module.msprobe_input_kwargs if hasattr(module, 'msprobe_input_kwargs') else {}
+            with self._no_grad_context():
+                BaseHookManager.inner_switch = False
+                module.forward_data_collected = True
+                self._add_count(api_name)
+                module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
+                self.data_collector.update_api_or_module_name(full_name)
+                if getattr(self.config, "online_run_ut", False):
+                    BaseHookManager.inner_switch = False
+                    return
+                self.data_collector.forward_input_data_collect(
+                    full_name,
+                    module,
+                    self._pid,
+                    module_input_output,
+                    self._is_recompute
+                )
+                BaseHookManager.inner_switch = False
+        return forward_pre_hook
+
+    def _build_forward_hook(self, hook_type, full_name):
+        def forward_hook(module, args, kwargs_or_output, output_or_kwargs=None):
+            if not self._should_execute_hook(hook_type, module, True):
+                self._clear_input_kwargs(module)
+                return None
+            kwargs, output = self._process_kwargs_and_output(module, hook_type, kwargs_or_output, output_or_kwargs)
+            BaseHookManager.inner_switch = True
+            self.data_collector.update_api_or_module_name(full_name)
+            module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+            with self._no_grad_context():
+                if getattr(self.config, "online_run_ut", False):
+                    if self.data_collector.scope and not self.data_collector.scope.check(full_name):
+                        return None
+                    if self.attl_manager:
+                        self.attl_manager.attl_send(full_name, args, kwargs, output)
+                    BaseHookManager.inner_switch = False
+                    return None
+                if hook_type == Const.MODULE:
+                    params_dict = self._get_params_dict(module)
+                    setattr(module_input_output, Const.PARAMS, params_dict)
+                    if params_dict:
+                        self._register_param_hook(full_name, module, params_dict)
+                    self.data_collector.update_api_or_module_name(full_name)
+                    self.data_collector.forward_data_collect(
+                        full_name,
+                        module,
+                        self._pid,
+                        module_input_output,
+                        self._is_recompute
+                    )
+                    self._init_params_grad_info(module, params_dict)
+                else:
+                    self.data_collector.forward_output_data_collect(
+                        full_name,
+                        module,
+                        self._pid,
+                        module_input_output,
+                        self._is_recompute
+                    )
+                self._clear_input_kwargs(module)
+
+                if self.data_collector.if_return_forward_new_output():
+                    forward_new_output = self.data_collector.get_forward_new_output()
+                    BaseHookManager.inner_switch = False
+                    return forward_new_output
+
+                BaseHookManager.inner_switch = False
+                return output
+        return forward_hook
+
+    def _build_backward_hook(self, hook_type, full_name):
+        def backward_hook(module, grad_input, grad_output):
+            if not self._should_execute_hook(hook_type, module, False):
+                return
+            BaseHookManager.inner_switch = True
+            self.data_collector.update_api_or_module_name(full_name)
+            if getattr(self.config, "online_run_ut", False):
+                BaseHookManager.inner_switch = False
+                return
+            need_exchange = self._need_exchange(module) if hook_type == Const.MODULE else True
+            if need_exchange:
+                module_input_output = ModuleBackwardInputsOutputs(grad_input=grad_output, grad_output=grad_input)
+            else:
+                module_input_output = ModuleBackwardInputsOutputs(grad_input=grad_input, grad_output=grad_output)
+            self.data_collector.backward_data_collect(
+                    full_name,
+                    module,
+                    self._pid,
+                    module_input_output,
+                    self._is_recompute
+                )
+            BaseHookManager.inner_switch = False
+        return backward_hook