mindstudio-probe 1.2.2__py3-none-any.whl → 8.1.0__py3-none-any.whl

msprobe/mindspore/dump/dump_tool_factory.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");
@@ -14,15 +14,18 @@
 # limitations under the License.
 
 from msprobe.mindspore.common.const import Const
+from msprobe.core.common.log import logger
+from msprobe.mindspore.common.utils import is_graph_mode_cell_dump_allowed
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
 from msprobe.mindspore.dump.kernel_graph_dump import KernelGraphDump
 from msprobe.mindspore.dump.kernel_kbyk_dump import KernelKbykDump
+from msprobe.mindspore.dump.graph_mode_cell_dump import GraphModeCellDump
 
 
 class DumpToolFactory:
     tools = {
         Const.CELL: {
-            Const.GRAPH_KBYK_MODE: None,
+            Const.GRAPH_KBYK_MODE: GraphModeCellDump,
             Const.GRAPH_GE_MODE: None,
             Const.PYNATIVE_MODE: None
         },
@@ -39,14 +42,21 @@ class DumpToolFactory:
     }
 
     @staticmethod
-    def create(config: DebuggerConfig):
-        if len(config.data_mode) != 1 or config.data_mode[0] not in Const.GRAPH_DATA_MODE_LIST:
-            raise Exception("data_mode must be one of all, input, output.")
+    def create(config: DebuggerConfig, model=None):
+        if config.level == Const.CELL:
+            if not is_graph_mode_cell_dump_allowed(config):
+                raise Exception("Cell dump is not supported in graph mode.")
+            if len(config.data_mode) != 1 or config.data_mode[0] not in Const.GRAPH_CELL_DUMP_DATA_MODE_LIST:
+                raise Exception("data_mode must be one of all, forward, backward.")
+        else:
+            if len(config.data_mode) != 1 or config.data_mode[0] not in Const.GRAPH_DATA_MODE_LIST:
+                raise Exception("data_mode must be one of all, input, output.")
         tool = DumpToolFactory.tools.get(config.level)
         if not tool:
             raise Exception("Valid level is needed.")
         tool = tool.get(config.execution_mode)
         if not tool:
-            raise Exception(f"Data dump is not supported in {config.execution_mode} mode "
-                            f"when dump level is {config.level}.")
-        return tool(config)
+            logger.error(f"Data dump is not supported in {config.execution_mode} mode "
+                         f"when dump level is {config.level}.")
+            raise ValueError
+        return tool(config, model) if tool == GraphModeCellDump else tool(config)

msprobe/mindspore/dump/graph_mode_cell_dump.py

@@ -0,0 +1,139 @@
+# 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
+
+import mindspore as ms
+from mindspore import hal, ops, Tensor
+from mindspore.ops.primitive import _run_op
+
+from msprobe.core.common.const import Const as CoreConst
+from msprobe.core.common.runtime import Runtime
+from msprobe.mindspore.common.const import Const
+from msprobe.mindspore.common.log import logger
+from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
+import msprobe.mindspore.dump.cell_dump_process as cellDumperWithDumpGradient
+import msprobe.mindspore.dump.cell_dump_with_insert_gradient as cellDumperWithInsertGradient
+
+tensordump_flag = True
+try:
+    from mindspore._c_expression import _tensordump_set_step
+except ImportError:
+    tensordump_flag = False
+
+
+class GraphModeCellDump:
+    task = CoreConst.STATISTICS
+
+    def __init__(self, config: DebuggerConfig, model, strict=True):
+        self.net = model
+        self.white_list = []
+        self.black_list = []
+        self.execution_mode = config.execution_mode
+        self.dump_path = config.dump_path if config.dump_path else "./"
+        self.rank = config.rank
+        self.step = config.step
+        self.scope = config.scope
+        self.list = config.list
+        self.data_mode = config.data_mode
+        self.file_format = config.file_format
+        GraphModeCellDump.task = config.task
+        self.summary_mode = config.summary_mode
+        self.check_config(strict)
+        self.set_step()
+
+    @staticmethod
+    def step():
+        # 更新TensorDump Step
+        if GraphModeCellDump.task == CoreConst.TENSOR:
+            hal.synchronize()
+            temp_tensor = ms.Tensor([1], dtype=ms.float32)
+            step_flag = "<tensordump-update-step>"
+            _run_op(ops.TensorDump(), "TensorDump", (step_flag, temp_tensor))
+            ops.tensordump(step_flag, temp_tensor)
+
+    def check_config(self, strict):
+        if not self.net:
+            raise Exception("The model is empty and cell dump is not enabled.")
+
+        if strict:
+            if self.rank:
+                raise Exception("In graph mode, cell dump does not currently support specifying rank.")
+            if self.scope:
+                raise Exception("In graph mode, cell dump does not currently support specifying scope.")
+            if self.list:
+                raise Exception("In graph mode, cell dump does not currently support specifying list.")
+            if len(self.data_mode) != 1 or self.data_mode[0] not in Const.GRAPH_CELL_DUMP_DATA_MODE_LIST:
+                raise Exception("In graph mode and cell dump, data_mode must be one of all, forword, backword.")
+            if self.file_format != []:
+                logger.warning("In graph mode, cell dump does not currently support specifying file_format."
+                               " The file will be stored in npy format.")
+            if self.task == CoreConst.STATISTICS and self.summary_mode == CoreConst.MD5:
+                raise Exception("The L0 level statistics dump mode does not support "
+                            "the calculation of md5 values currently In graph mode.")
+        else:
+            self.rank = []
+            self.scope = []
+            self.list = []
+            self.file_format = []
+            if len(self.data_mode) != 1 or self.data_mode[0] not in Const.GRAPH_CELL_DUMP_DATA_MODE_LIST:
+                self.data_mode = [CoreConst.ALL]
+            if self.task == CoreConst.STATISTICS and self.summary_mode == CoreConst.MD5:
+                self.summary_mode = CoreConst.STATISTICS
+
+        return True
+
+    def set_step(self):
+        if tensordump_flag:
+            _tensordump_set_step(self.step)
+        else:
+            raise Exception(
+                "Importing _tensordump_set_step failed, "
+                "please use the latest version package of MindSpore."
+            )
+
+    def handle(self):
+        os.environ['MS_JIT_MODULES'] = 'msprobe'
+
+        if Runtime.run_mode == Const.PYNATIVE_GRAPH_MODE:
+            dump_path = os.path.join(self.dump_path, Const.GRAPH_MODE)
+        else:
+            dump_path = self.dump_path
+
+        cell_dumper = cellDumperWithDumpGradient
+
+        if self.execution_mode == Const.PYNATIVE_MODE:
+            enable_dump_gradient = hasattr(ops, 'DumpGradient')
+            if hasattr(ops, 'DumpGradient'):
+                try:
+                    ops.DumpGradient()('grad.npy', Tensor([0], dtype=ms.float32), 'in')
+                except Exception:
+                    enable_dump_gradient = False
+                    logger.warning('the DumpGradient operator failed to execute.')
+            if not enable_dump_gradient:
+                cell_dumper = cellDumperWithInsertGradient
+        
+        dump_config = cell_dumper.CellDumpConfig(
+            net=self.net,
+            dump_path=dump_path,
+            data_mode=self.data_mode[0],
+            task=self.task,
+            summary_mode=self.summary_mode,
+            step=self.step
+        )
+
+        cell_dumper.start(
+            dump_config
+        )

msprobe/mindspore/dump/graph_tensor_dump.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2024-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 collections import OrderedDict
+import mindspore as ms
+
+
+def _iterate_items(data):
+    if isinstance(data, (dict, OrderedDict)):
+        return data.items()
+    elif isinstance(data, (list, tuple)):
+        return enumerate(data)
+    else:
+        raise TypeError("Unsupported data type")
+
+
+class _SaveBase:
+    def __init__(self, save_dir):
+        super(_SaveBase, self).__init__()
+        self.path = save_dir
+        self.save_func = _npy_save
+
+    def get_save_func(self):
+        return self.save_func
+
+
+@ms.jit_class
+class _SaveCell(_SaveBase):
+    def __call__(self, name, data):
+        return self.get_save_func()(self.path, name, data)
+
+
+class _SaveGradBase:
+    def __init__(self, save_dir, name):
+        super(_SaveGradBase, self).__init__()
+        self.file = save_dir + name
+
+
+@ms.jit_class
+class _SaveGradCell(_SaveGradBase):
+    def __init__(self, save_dir, name):
+        super(_SaveGradCell, self).__init__(save_dir, name)
+        self.ms_save_grad = ms.ops.InsertGradientOf(
+            _wrapper_save_grad_func(self.file))
+
+    def __call__(self, x):
+        if isinstance(x, ms.Tensor):
+            return self.ms_save_grad(x)
+        else:
+            raise TypeError(f"For 'save_grad', the type of argument 'data' must be mindspore.Tensor or torch.tensor, "
+                            f"but got {type(x)}")
+
+
+def _npy_save_ops(file, data):
+    if isinstance(data, ms.Tensor):
+        if data.dtype == ms.bfloat16:
+            data = data.float()
+        ms.ops.TensorDump()(file, data)
+    else:
+        raise TypeError(f"For 'save', the type of argument 'data' must be mindspore.Tensor or torch.tensor, "
+                        f"but got {type(data)}")
+
+
+def _wrapper_save_grad_func(file):
+    def _save_grad_func(grad):
+        data = grad
+        if data.dtype == ms.bfloat16:
+            data = data.float()
+        ms.ops.TensorDump()(file, data)
+        return grad
+    return _save_grad_func
+
+
+def _npy_save(save_dir, item_name, data):
+    if isinstance(data, (list, tuple, dict, OrderedDict)):
+        for key, val in _iterate_items(data):
+            _npy_save(save_dir, f"{item_name}.{key}", val)
+    else:
+        if data is None:
+            return
+        _npy_save_ops(f"{save_dir}{item_name}", data)
+
+
+def generate_dump_dir(save_dir, sep=os.sep):
+    """
+    usage: generate dump directory path str in mindspore graph mode
+    """
+    full_suffix = '{step}' + sep + '{rank}' + sep
+    if save_dir and save_dir[-1] != sep:
+        result_dir = save_dir + sep + full_suffix
+    else:
+        result_dir = save_dir + full_suffix
+    return result_dir
+
+
+def save(save_dir, name, data):
+    """
+    save tensor.
+    """
+    dump_dir = generate_dump_dir(save_dir)
+    _SaveCell(dump_dir)(name, data)
+
+
+def save_grad(save_dir, name, data):
+    """
+    save grad.
+    """
+    dump_dir = generate_dump_dir(save_dir)
+    suffix_name = name + '_grad'
+    return _SaveGradCell(dump_dir, suffix_name)(data)

msprobe/mindspore/dump/hook_cell/api_register.py

@@ -0,0 +1,176 @@
+# 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
+import inspect
+
+from mindspore import Tensor, ops, mint
+from mindspore.mint import distributed
+from mindspore.mint.nn import functional
+from mindspore.communication import comm_func
+
+from msprobe.core.common.file_utils import load_yaml
+from msprobe.core.common.utils import Const
+from msprobe.core.data_dump.api_registry import ApiRegistry
+from msprobe.mindspore.common.log import logger
+from msprobe.mindspore.common.const import Const as MsConst
+from msprobe.mindspore.common.utils import is_mindtorch
+from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
+
+
+stub_tensor_existed = True
+try:
+    from mindspore.common._stub_tensor import StubTensor
+except ImportError:
+    stub_tensor_existed = False
+
+cur_path = os.path.dirname(os.path.realpath(__file__))
+if not is_mindtorch():
+    _api_types = {
+        Const.MS_FRAMEWORK: {
+            Const.MS_API_TYPE_OPS: (ops, (ops,)),
+            Const.MS_API_TYPE_TENSOR: (Tensor, (Tensor,)),
+            Const.MS_API_TYPE_MINT: (mint, (mint,)),
+            Const.MS_API_TYPE_MINT_FUNC: (functional, (functional,)),
+            Const.MS_API_TYPE_COM: (comm_func, (comm_func,)),
+            Const.MS_API_TYPE_MINT_DIST: (distributed, (distributed,))
+        }
+    }
+    if stub_tensor_existed:
+        _api_types.get(Const.MS_FRAMEWORK).update(
+            {Const.MS_API_TYPE_STUB_TENSOR: (StubTensor, (StubTensor,))}
+        )
+
+    _supported_api_list_path = (os.path.join(cur_path, MsConst.SUPPORTED_API_LIST_FILE),)
+    _backlist = []
+else:
+    import torch
+    import torch_npu
+    _api_types = {
+        Const.MT_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_NPU: (torch_npu, (torch_npu,)),
+            Const.PT_API_TYPE_DIST: (torch.distributed, (torch.distributed, torch.distributed.distributed_c10d))
+        }
+    }
+    _supported_api_list_path = (os.path.join(cur_path, '../../../pytorch/hook_module',
+                                             MsConst.SUPPORTED_API_LIST_FILE),)
+    _backlist = [f'{Const.PT_API_TYPE_TENSOR}.__setitem__']
+
+_inner_used_api = {
+    Const.MS_FRAMEWORK + Const.SEP + Const.MS_API_TYPE_OPS: (
+        ops, "norm", "square", "sqrt", "is_complex", "stack", "is_floating_point"
+    ),
+    Const.MS_FRAMEWORK + Const.SEP + Const.MS_API_TYPE_TENSOR: (
+        Tensor, "to", "numel", 'sum'
+    ),
+    Const.MS_FRAMEWORK + Const.SEP + Const.MS_API_TYPE_MINT: (
+        mint, "max", "min", "mean", "norm"
+    )
+}
+
+
+class ApiTemplate(HOOKCell):
+    def __init__(self, api_name, api_func, prefix, hook_build_func):
+        self.api_name = api_name
+        self.api_func = api_func
+        self.prefix_api_name = prefix + Const.SEP + str(api_name.split(Const.SEP)[-1]) + Const.SEP
+        super().__init__(hook_build_func)
+        distributed_prefix = Const.DIST_API_TYPE_PREFIX if is_mindtorch() else Const.MINT_DIST_API_TYPE_PREFIX
+        if prefix == distributed_prefix:
+            self.op_is_distributed = True
+
+    @staticmethod
+    def async_to_sync(output):
+        # Fake handle, used to return after the CommHandle executes the wait method
+        fake_handle = type("FakeHandle", (), {"wait": lambda self: None})()
+        if isinstance(output, tuple) and len(output) == 2 and hasattr(output[1], "wait"):
+            output[1].wait()
+            output = (output[0], fake_handle)
+        elif hasattr(output, "wait"):
+            output.wait()
+            output = fake_handle
+        return output
+
+    def construct(self, *args, **kwargs):
+        if self.api_name.startswith(MsConst.DROPOUT_API_NAME_PREFIX):
+            return args[0] if args else kwargs.get(Const.INPUT)
+
+        output = self.api_func(*args, **kwargs)
+
+        if self.prefix_api_name.startswith(
+            (MsConst.DISTRIBUTED_DATA_PREFIX, Const.MINT_DIST_API_TYPE_PREFIX)
+        ):
+            try:
+                bound = inspect.signature(self.api_func).bind(*args, **kwargs)
+                bound.apply_defaults()
+                use_asyn_op_flag = bound.arguments.get("asyn_op", False)
+            except Exception as e:
+                use_asyn_op_flag = False
+                logger.warning(f"fail to get dist api's func signature because {e}, no wait")
+
+            if use_asyn_op_flag or self.api_name in ["isend", "irecv"]:
+                output = self.async_to_sync(output)
+            if self.api_name == "batch_isend_irecv" and isinstance(output, list):
+                output = [self.async_to_sync(handle) for handle in output]
+
+        return output
+
+    def forward(self, *args, **kwargs):
+        if self.api_name.startswith(MsConst.DROPOUT_API_NAME_PREFIX):
+            return args[0] if args else kwargs.get(Const.INPUT)
+        return self.api_func(*args, **kwargs)
+
+
+api_register = None
+stub_tensor_set = False
+
+
+def get_api_register(return_new=False):
+    global stub_tensor_set
+
+    def stub_method(method):
+        def wrapped_method(*args, **kwargs):
+            return method(*args, **kwargs)
+        return wrapped_method
+    if not is_mindtorch() and stub_tensor_existed and not stub_tensor_set:
+        api_names = load_yaml(_supported_api_list_path[0]).get(Const.MS_API_TYPE_TENSOR, [])
+        for attr_name in dir(StubTensor):
+            attr = getattr(StubTensor, attr_name)
+            if attr_name in api_names and callable(attr):
+                setattr(StubTensor, attr_name, stub_method(attr))
+        stub_tensor_set = True
+
+    if return_new:
+        return ApiRegistry(
+            _api_types,
+            _inner_used_api,
+            _supported_api_list_path,
+            ApiTemplate,
+            _backlist
+        )
+
+    global api_register
+    if api_register is None:
+        api_register = ApiRegistry(
+            _api_types,
+            _inner_used_api,
+            _supported_api_list_path,
+            ApiTemplate,
+            _backlist
+        )
+    return api_register

msprobe/mindspore/dump/hook_cell/hook_cell.py

@@ -15,11 +15,16 @@
 
 from collections import defaultdict
 
+import mindspore as ms
 from mindspore import nn
 
+from msprobe.core.common.runtime import Runtime
 from msprobe.mindspore.common.utils import is_mindtorch, register_backward_hook_functions
 
 
+ms_version = ms.__version__
+
+
 def add_cell_count(name):
     HOOKCell.cell_count[name] += 1
 
@@ -28,29 +33,34 @@ def get_cell_count(name):
     return HOOKCell.cell_count[name]
 
 
-def __init__(self, build_hook) -> None:
+def __init__(self, hook_build_func) -> None:
     super(HOOKCell, self).__init__()
     self.changed_status = False
-    self.input_kwargs = {}
-    self.prefix = ""
+    self.msprobe_input_kwargs = {}
     if not HOOKCell.g_stop_hook:
         HOOKCell.g_stop_hook = True
         self.changed_status = True
-        if hasattr(self, "prefix_api_name"):
-            self.prefix = self.prefix_api_name
-
         self.forward_data_collected = False
-        forward_pre_hook, forward_hook, backward_hook, backward_pre_hook = build_hook(self.prefix)
-        self.register_forward_pre_hook(forward_pre_hook)
-        self.register_forward_hook(forward_hook)
-        register_backward_hook_functions["full"](self, backward_hook)
-        register_backward_hook_functions["pre"](self, backward_pre_hook)
+
+        if not Runtime.is_running:
+            return
+        prefix = self.prefix_api_name if hasattr(self, "prefix_api_name") else ""
+        if callable(hook_build_func):
+            hook_set = hook_build_func(prefix)
+            if ms_version < "2.6.0" and not is_mindtorch():
+                getattr(self, "_forward_pre_hook", {})[id(self)] = hook_set.forward_pre_hook
+                getattr(self, "_forward_hook", {})[id(self)] = hook_set.forward_hook
+            else:
+                self.register_forward_pre_hook(hook_set.forward_pre_hook)
+                self.register_forward_hook(hook_set.forward_hook)
+            register_backward_hook_functions["full"](self, hook_set.backward_hook)
+            register_backward_hook_functions["pre"](self, hook_set.backward_pre_hook)
 
 
 # 重载call，加全局标志。
 def __call__(self, *args, **kwargs):
     try:
-        self.input_kwargs = kwargs
+        self.msprobe_input_kwargs = kwargs
         out = super(HOOKCell, self).__call__(*args, **kwargs)
     except Exception as e:
         raise e

msprobe/mindspore/dump/hook_cell/ms_hook_manager.py

@@ -0,0 +1,88 @@
+# 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 mindspore.common.api import _no_grad
+from msprobe.core.common.const import Const
+from msprobe.core.common.utils import replace_last_occurrence
+from msprobe.core.data_dump.data_processor.base import ModuleBackwardInputs
+from msprobe.core.hook_manager import BaseHookManager, HookSet
+from msprobe.mindspore.common.utils import has_kwargs_in_forward_hook
+from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
+
+
+class MindsproeHookManager(BaseHookManager):
+    @property
+    def _is_recompute(self):
+        return None
+
+    @staticmethod
+    def _no_grad_context():
+        return _no_grad()
+
+    @staticmethod
+    def _add_count(name):
+        HOOKCell.add_cell_count(name)
+
+    @staticmethod
+    def _process_kwargs_and_output(module, hook_type, kwargs_or_output, output_or_kwargs):
+        if not has_kwargs_in_forward_hook() or hook_type == Const.API:
+            kwargs = module.msprobe_input_kwargs if hasattr(module, 'msprobe_input_kwargs') else {}
+            output = kwargs_or_output
+        else:
+            kwargs = kwargs_or_output
+            output = output_or_kwargs
+        return kwargs, output
+
+    def build_hook(self, hook_type, name):
+        if hook_type == Const.API:
+            full_forward_name = name + str(HOOKCell.get_cell_count(name)) + Const.SEP + Const.FORWARD
+        else:
+            full_forward_name = name
+        full_backward_name = replace_last_occurrence(full_forward_name, Const.FORWARD, Const.BACKWARD)
+        hookset = HookSet(
+            forward_hook=self._build_forward_hook(hook_type, full_forward_name),
+            forward_pre_hook=self._build_forward_pre_hook(hook_type, full_forward_name, name),
+            backward_hook=self._build_backward_hook(hook_type, full_backward_name),
+            backward_pre_hook=self._build_backward_pre_hook(hook_type, full_backward_name)   
+        )
+        return hookset
+
+    def _need_exchange(self, module):
+        if not hasattr(module, 'has_pre_hook_called') or not module.has_pre_hook_called:
+            return False
+        else:
+            return True
+
+    def _get_params_dict(self, module):
+        params_dict = {}
+        if self.config.task != Const.STRUCTURE:
+            params_dict = {
+                    key.split(Const.SEP)[-1]: value
+                    for key, value in module.parameters_dict(recurse=False).items()
+                    }
+        return params_dict
+
+    def _build_backward_pre_hook(self, hook_type, name):
+        def backward_pre_hook(module, grad_input):
+            if self.config.level != Const.LEVEL_L2:
+                return
+            if not self._should_execute_hook(hook_type, module, False):
+                return
+            BaseHookManager.inner_switch = True
+            module_input = ModuleBackwardInputs(grad_input=grad_input)
+            self.data_collector.update_api_or_module_name(name)
+            self.data_collector.backward_input_data_collect(name, module, self._pid, module_input)
+            BaseHookManager.inner_switch = False
+        return backward_pre_hook