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

msprobe/pytorch/dump/module_dump/module_dump.py

@@ -13,75 +13,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-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.dump.module_dump.module_processer import ModuleProcesser
 from msprobe.pytorch.hook_module.api_register import get_api_register
 
-torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
-
 
 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):
+        if hasattr(module, 'msprobe_hook') and not hasattr(module, 'msprobe_module_dump'):
+            logger.info_on_rank_0("The init dump is enabled, and the module dump function will not be available.")
+            return
+
+        ModuleProcesser.enable_module_dump = True
         self.api_register.restore_all_api()
-        self.register_hook(module, dump_name)
+        if not hasattr(module, 'msprobe_module_dump'):
+            self.service.module_processor.register_module_hook(module, self.service.build_hook,
+                                                               recursive=False, module_names=[dump_name])
+            setattr(module, 'msprobe_module_dump', True)
 
     def stop_module_dump(self):
+        ModuleProcesser.enable_module_dump = False
         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()
-        self.hook_handle_list.clear()
-
-    def register_hook(self, module, dump_name):
-        prefix_name = (
-                BaseScope.Module_Type_Module + Const.SEP +
-                dump_name + Const.SEP +
-                module.__class__.__name__ + Const.SEP
-        )
-        module_processor = self.service.module_processor
-        _, forward_hook, backward_hook, forward_hook_torch_version_below_2 = self.service.build_hook(
-            BaseScope.Module_Type_Module,
-            prefix_name
-        )
-
-        if module_processor.has_register_backward_hook(module):
-            logger.warning(
-                f"The {dump_name} module has registered deprecated register_backward_hook,"
-                f"which may cause abnormal data dump. The backward data dump for this module will be skipped."
-            )
-        if torch_version_above_or_equal_2:
-            forward_hook_handle = module.register_forward_hook(forward_hook, with_kwargs=True)
-        else:
-            if not module_processor.has_register_backward_hook(module):
-                backward_hook_handle = module.register_full_backward_hook(
-                    module_processor.node_hook(prefix_name + Const.BACKWARD, Const.STOP)
-                )
-                self.hook_handle_list.append(backward_hook_handle)
-            forward_hook_handle = module.register_forward_hook(forward_hook_torch_version_below_2)
-        self.hook_handle_list.append(forward_hook_handle)
-        if not module_processor.has_register_backward_hook(module):
-            backward_hook_handle = module.register_full_backward_hook(backward_hook)
-            self.hook_handle_list.append(backward_hook_handle)
-
-        forward_pre_hook_handle = module.register_forward_pre_hook(
-            module_processor.node_hook(prefix_name + Const.FORWARD, Const.START)
-        )
-        forward_hook_handle = module.register_forward_hook(
-            module_processor.node_hook(prefix_name + Const.FORWARD, Const.STOP)
-        )
-        self.hook_handle_list.extend([forward_pre_hook_handle, forward_hook_handle])
-        if torch_version_above_or_equal_2 and not module_processor.has_register_backward_hook(module):
-            backward_pre_hook_handle = module.register_full_backward_pre_hook(
-                module_processor.node_hook(prefix_name + Const.BACKWARD, Const.START)
-            )
-            backward_hook_handle = module.register_full_backward_hook(
-                module_processor.node_hook(prefix_name + Const.BACKWARD, Const.STOP)
-            )
-            self.hook_handle_list.extend([backward_pre_hook_handle, backward_hook_handle])

msprobe/pytorch/dump/module_dump/module_processer.py

@@ -13,16 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from functools import wraps
+from collections import OrderedDict
 
 import torch
-from torch.utils.hooks import BackwardHook
+from torch.utils.hooks import BackwardHook, RemovableHandle
 
 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 msprobe.pytorch.common.utils import replace_last_occurrence, is_float8_tensor
+from msprobe.pytorch.common.utils import is_torch_nn_module, register_forward_pre_hook
+from msprobe.pytorch.dump.module_dump.hook_wrapper import wrap_setup_input_output_hook
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 if torch_version_above_or_equal_2:
@@ -39,43 +39,40 @@ def replace_checkpoint():
         torch.utils.checkpoint.checkpoint = checkpoint_without_early_stop
 
 
+def wrap_megatron_deallocate(func):
+    def wrapper_func(out, deallocate_pipeline_outputs=False):
+        if deallocate_pipeline_outputs and isinstance(out, torch.Tensor) and getattr(out, "_base") is not None:
+            out_clone = out.clone()
+            out.data = torch.empty((1,), device=out.device, dtype=out.dtype, )
+            return func(out_clone, deallocate_pipeline_outputs)
+        return func(out, deallocate_pipeline_outputs)
+    return wrapper_func
+
+
 class ModuleProcesser:
     module_count = {}
     module_stack = []
     api_parent_node = ""
     module_node = {}
+    module_bw_hook_kernels = {}
+    module_with_backward_hook = {}
+    enable_module_dump = False
 
     def __init__(self, scope):
         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)
+        wrap_setup_input_output_hook()
         replace_checkpoint()
+        try:
+            from megatron.core.pipeline_parallel import schedules
+            schedules.deallocate_output_tensor = wrap_megatron_deallocate(schedules.deallocate_output_tensor)
+            logger.info_on_rank_0("Patch megatron method success.")
+        except ImportError:
+            logger.info_on_rank_0("No megatron find.")
+        except Exception as e:
+            logger.info_on_rank_0(f"Patch megatron method failed, detail:{str(e)}")
 
     @staticmethod
-    def clone_return_value(func):
-        @wraps(func)
-        def clone_return_value_func(*args, **kwargs):
-            result = func(*args, **kwargs)
-            return ModuleProcesser.clone_if_tensor(result)
-
-        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) and not is_float8_tensor(result):
-            return result.clone()
-        elif type(result) is tuple:
-            return tuple(ModuleProcesser.clone_if_tensor(x) for x in result)
-        elif type(result) is list:
-            return list(ModuleProcesser.clone_if_tensor(x) for x in result)
-        elif type(result) is dict:
-            return {k: ModuleProcesser.clone_if_tensor(v) for k, v in result.items()}
-        else:
-            return result
-
-    @staticmethod
-    def module_count_func(module_name):
+    def set_and_get_calls_number(module_name):
         if module_name not in ModuleProcesser.module_count:
             ModuleProcesser.module_count[module_name] = 0
         else:
@@ -89,13 +86,19 @@ class ModuleProcesser:
             module._is_full_backward_hook is False
 
     @staticmethod
-    def get_modules_and_names(models):
+    def get_modules_and_names(models, recursive, module_names):
         modules_and_names_with_index = {}
         if isinstance(models, (list, tuple)):
+            if not recursive and len(module_names) != len(models):
+                return modules_and_names_with_index
             for index, model in enumerate(models):
-                modules_and_names_with_index[str(index)] = model.named_modules()
+                modules_and_names_with_index[str(index)] = model.named_modules() if recursive else \
+                    [(module_names[index], model)]
         else:
-            modules_and_names_with_index["-1"] = models.named_modules()
+            if not recursive and len(module_names) != 1:
+                return modules_and_names_with_index
+            modules_and_names_with_index["-1"] = models.named_modules() if recursive else \
+                [(module_names[0], models)]
         return modules_and_names_with_index
 
     @classmethod
@@ -104,107 +107,134 @@ class ModuleProcesser:
         cls.module_stack = []
         cls.api_parent_node = ""
         cls.module_node = {}
+        cls.module_bw_hook_kernels = {}
+        cls.enable_module_dump = False
 
-    def register_module_hook(self, models, build_hook):
-        logger.info_on_rank_0("The init dump is enabled, and the module dump function will not be available.")
-        modules_and_names_with_index = self.get_modules_and_names(models)
+    def register_module_hook(self, models, build_hook, recursive=True, module_names=None):
+        if module_names is None:
+            module_names = []
+
+        modules_and_names_with_index = self.get_modules_and_names(models, recursive, module_names)
         for index, modules_and_names in modules_and_names_with_index.items():
             model = models if index == "-1" else models[int(index)]
             for name, module in modules_and_names:
-                if module == model:
+                if recursive and module == model:
+                    continue
+                if not is_torch_nn_module(module):
+                    logger.warning(
+                        f"The module dump does not support {type(module)} type. "
+                        f"The data dump for this module will be skipped."
+                    )
                     continue
                 if module.__class__.__name__ == "FullyShardedDataParallel":
                     continue
+                setattr(module, 'msprobe_hook', True)
                 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)
-                pre_forward_hook, forward_hook, backward_hook, forward_hook_torch_version_below_2 = build_hook(
-                    BaseScope.Module_Type_Module,
-                    prefix_name
-                )
+                prefix_name = f'{BaseScope.Module_Type_Module}{Const.SEP}{module_index}{name}{Const.SEP}' + \
+                              f'{module.__class__.__name__}{Const.SEP}'
+
+                forward_pre_hook = self.build_module_hook(prefix_name, build_hook)
 
                 if self.has_register_backward_hook(module):
                     logger.warning(
                         f"The {prefix_name[:-1]} has registered deprecated register_backward_hook,"
                         f"which may cause abnormal data dump. The backward data dump for this module will be skipped."
                     )
+                    ModuleProcesser.module_with_backward_hook[prefix_name] = True
+                register_forward_pre_hook(module, forward_pre_hook)
+
+    def build_module_hook(self, module_name, build_data_hook):
+        def forward_pre_hook(module, args, kwargs=None):
+            if kwargs is None:
+                kwargs = {}
+
+            if hasattr(module, 'msprobe_module_dump') and not self.enable_module_dump:
+                return (args, kwargs) if torch_version_above_or_equal_2 else args
+
+            index = ModuleProcesser.set_and_get_calls_number(module_name)
+            full_forward_name = f'{module_name}{Const.FORWARD}{Const.SEP}{index}'
+            full_backward_name = f'{module_name}{Const.BACKWARD}{Const.SEP}{index}'
+
+            self.set_construct_info_in_pre_hook(full_forward_name)
+
+            if not hasattr(module, 'msprobe_forward_hook'):
+                forward_hooks_dict = getattr(module, '_forward_hooks', OrderedDict())
+                handle = RemovableHandle(forward_hooks_dict)
+                forward_hooks_dict[handle.id] = forward_hook
+                forward_hooks_dict.move_to_end(handle.id, last=False)
+                if torch_version_above_or_equal_2:
+                    forward_hooks_with_kwargs_dict = getattr(module, '_forward_hooks_with_kwargs', OrderedDict())
+                    forward_hooks_with_kwargs_dict[handle.id] = True
+
+                setattr(module, 'msprobe_forward_hook', True)
+
+            hook_set = build_data_hook(BaseScope.Module_Type_Module, full_forward_name)
+
+            def get_backward_pre_hook(full_backward_name):
+                def backward_pre_hook_fn(module, grad_output):
+                    self.set_construct_info_in_pre_hook(full_backward_name)
+                return backward_pre_hook_fn
+
+            def get_backward_hook(backward_data_hook, full_backward_name):
+                def backward_hook_fn(module, grad_input, grad_output):
+                    new_output = backward_data_hook(module, grad_input, grad_output)
+                    self.set_construct_info_in_hook(full_backward_name, is_forward=False)
+                    return new_output
+                return backward_hook_fn
+
+            if not ModuleProcesser.module_with_backward_hook.get(module_name):
+                backward_pre_hook = get_backward_pre_hook(full_backward_name)
+                backward_hook = get_backward_hook(hook_set.backward_hook, full_backward_name)
                 if torch_version_above_or_equal_2:
-                    module.register_forward_hook(forward_hook, with_kwargs=True)
+                    bw_hook = BackwardHook(module, [backward_hook], [backward_pre_hook])
                 else:
-                    if not self.has_register_backward_hook(module):
-                        module.register_full_backward_hook(self.node_hook(prefix_name + Const.BACKWARD, Const.STOP))
-                    module.register_forward_hook(forward_hook_torch_version_below_2)
-                if not self.has_register_backward_hook(module):
-                    module.register_full_backward_hook(backward_hook)
-
-                module.register_forward_pre_hook(self.node_hook(prefix_name + Const.FORWARD, Const.START))
-                module.register_forward_hook(self.node_hook(prefix_name + Const.FORWARD, Const.STOP))
-                if torch_version_above_or_equal_2 and not self.has_register_backward_hook(module):
-                    module.register_full_backward_pre_hook(self.node_hook(prefix_name + Const.BACKWARD, Const.START))
-                    module.register_full_backward_hook(self.node_hook(prefix_name + Const.BACKWARD, Const.STOP))
-
-    def node_hook(self, name_prefix, start_or_stop, **kwargs):
-
-        def pre_hook(module, input, output=None):
-            try:
-                index = ModuleProcesser.module_count_func(name_prefix)
-            except IndexError as e:
-                index = None
-                pass
-            full_name = name_prefix + Const.SEP + str(index)
-            if not hasattr(module, "mindstudio_reserved_name") or not module.mindstudio_reserved_name:
-                module.mindstudio_reserved_name = []
-            module.mindstudio_reserved_name.append(full_name)
-            if self.module_stack:
-                ModuleProcesser.module_node[full_name] = self.module_stack[-1]
+                    bw_hook = BackwardHook(module, [backward_hook])
+                ModuleProcesser.module_bw_hook_kernels[full_forward_name] = bw_hook
+                args = bw_hook.setup_input_hook(args)
+            return (args, kwargs) if torch_version_above_or_equal_2 else args
+
+        def forward_hook(module, args, kwargs_or_output, output_or_kwargs=None):
+            if hasattr(module, 'msprobe_module_dump') and not self.enable_module_dump:
+                return output_or_kwargs if torch_version_above_or_equal_2 else kwargs_or_output
+
+            index = ModuleProcesser.module_count.get(module_name)
+            full_name = f'{module_name}{Const.FORWARD}{Const.SEP}{index}'
+
+            hook_set = build_data_hook(BaseScope.Module_Type_Module, full_name)
+            hook_result = hook_set.forward_hook(module, args, kwargs_or_output, output_or_kwargs)
+            self.set_construct_info_in_hook(full_name)
+
+            if hook_result is not None:
+                result = hook_result
             else:
-                ModuleProcesser.module_node[full_name] = None
+                result = output_or_kwargs if torch_version_above_or_equal_2 else kwargs_or_output
 
-            ModuleProcesser.module_stack.append(full_name)
-            if self.module_stack:
-                ModuleProcesser.api_parent_node = self.module_stack[-1]
-            if self.scope:
-                self.scope.begin_module(full_name)
+            bw_hook = ModuleProcesser.module_bw_hook_kernels.get(full_name)
+            if bw_hook:
+                result = bw_hook.setup_output_hook(result)
 
-        def end_hook(module, input, output=None):
+            return result
+
+        return forward_pre_hook
+
+    def set_construct_info_in_pre_hook(self, full_name):
+        if self.module_stack:
+            ModuleProcesser.module_node[full_name] = self.module_stack[-1]
+        else:
+            ModuleProcesser.module_node[full_name] = None
+        ModuleProcesser.module_stack.append(full_name)
+        ModuleProcesser.api_parent_node = full_name
+        if self.scope:
+            self.scope.begin_module(full_name)
+
+    def set_construct_info_in_hook(self, full_name, is_forward=True):
+        if torch_version_above_or_equal_2 or is_forward:
             if self.module_stack:
                 ModuleProcesser.module_stack.pop()
-            if self.module_stack:
-                ModuleProcesser.api_parent_node = self.module_stack[-1]
-            else:
-                ModuleProcesser.api_parent_node = None
-            if not hasattr(module, "mindstudio_reserved_name") or not module.mindstudio_reserved_name:
-                raise RuntimeError(f"module reserve name is None when pop")
-            current_name = module.mindstudio_reserved_name.pop()
+            ModuleProcesser.api_parent_node = ModuleProcesser.module_stack[-1] if self.module_stack else None
             if self.scope:
-                self.scope.end_module(current_name)
-
-        def backward_hook(module, input, output=None):
-            try:
-                index = ModuleProcesser.module_count_func(name_prefix)
-            except IndexError as e:
-                index = None
-                pass
-            full_name = name_prefix + Const.SEP + str(index)
-            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 = 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
+                self.scope.end_module(full_name)
+        else:
             if self.scope:
                 self.scope.begin_module(full_name)
-
-        if torch_version_above_or_equal_2:
-            if Const.START in start_or_stop:
-                return pre_hook
-            else:
-                return end_hook
-        else:
-            if Const.FORWARD in name_prefix and Const.START in start_or_stop:
-                return pre_hook
-            elif Const.BACKWARD in name_prefix:
-                return backward_hook
-            else:
-                return end_hook
+            ModuleProcesser.api_parent_node = full_name

msprobe/pytorch/grad_probe/grad_stat_csv.py

@@ -17,6 +17,7 @@ from abc import ABC, abstractmethod
 from collections import namedtuple
 import hashlib
 from functools import wraps
+import zlib
 import torch
 from msprobe.core.grad_probe.constant import GradConst
 
@@ -74,8 +75,8 @@ class CsvMd5(CsvItem):
     def generate_csv_content(csv_content_input):
         grad = csv_content_input.grad
         tensor_bytes = grad.cpu().detach().float().numpy().tobytes()
-        md5_hash = hashlib.md5(tensor_bytes)
-        return [md5_hash.hexdigest()]
+        md5_hash = f"{zlib.crc32(tensor_bytes):08x}"
+        return [md5_hash]
 
 
 @register_csv_item(GradConst.DISTRIBUTION)

msprobe/pytorch/hook_module/api_register.py

@@ -15,21 +15,36 @@
 
 import functools
 import os
+import inspect
 
 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.log import logger
 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
+from msprobe.pytorch.hook_module.utils import dynamic_import_op
+from msprobe.core.common.file_utils import load_yaml
+
+try:
+    import mindspeed.ops
+except ImportError:
+    mindspeed_enable = False
+else:
+    mindspeed_enable = True
 
 
 torch_version_above_2 = torch.__version__.split('+')[0] > '2.0'
 
+_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"}
+
 _api_types = {
     Const.PT_FRAMEWORK: {
         Const.PT_API_TYPE_FUNCTIONAL: (torch.nn.functional, (torch.nn.functional,)),
@@ -57,10 +72,11 @@ if not is_gpu:
                                                                      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"}
+    if mindspeed_enable:
+        _api_types.get(Const.PT_FRAMEWORK).update({Const.PT_API_TYPE_MINDSPEED: (mindspeed.ops, (mindspeed.ops,))})
+        mindspeed_op_list = load_yaml(_supported_api_list_path[0]).get(Const.PT_API_TYPE_MINDSPEED)
+        mindspeed_op_file_list = [op.split(Const.SEP)[0] + Const.PY_SUFFIX for op in mindspeed_op_list]
+        dynamic_import_op(mindspeed.ops, mindspeed_op_file_list)
 
 
 @parameter_adapter
@@ -70,7 +86,15 @@ def tensor_module_forward(module, *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"]:
+    try:
+        bound = inspect.signature(module.api_func).bind(*args, **kwargs)
+        bound.apply_defaults()
+        use_async_op_flag = bound.arguments.get("async_op", False)
+    except Exception as e:
+        use_async_op_flag = False
+        logger.warning(f"fail to get dist api's func signature because {e}, no wait")
+
+    if use_async_op_flag or module.api_name in ["isend", "irecv"]:
         if handle and hasattr(handle, 'wait'):
             handle.wait()
     if module.api_name == "batch_isend_irecv":

msprobe/pytorch/hook_module/hook_module.py

@@ -21,9 +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'
+from msprobe.core.common.runtime import Runtime
+from msprobe.pytorch.common.utils import is_float8_tensor, register_forward_pre_hook, register_forward_hook
 
 
 class HOOKModule(nn.Module):
@@ -41,16 +40,14 @@ class HOOKModule(nn.Module):
         if not self.stop_hook:
             self.forward_data_collected = False
 
+            if not Runtime.is_running:
+                return
             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)
+                hook_set = hook_build_func(prefix)
+                register_forward_pre_hook(self, hook_set.forward_pre_hook)
+                register_forward_hook(self, hook_set.forward_hook)
+                self.register_backward_hook(hook_set.backward_hook)
 
     def __call__(self, *args, **kwargs):
         changed = False
@@ -79,13 +76,7 @@ class HOOKModule(nn.Module):
         if len(self._backward_hooks) > 0:
             full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()
         for hook in self._forward_pre_hooks.values():
-            result_args, result_kwargs = hook(self, args, kwargs)
-            if result_args is not None:
-                if not isinstance(result_args, tuple):
-                    result_args = (result_args,)
-                args = result_args
-            if result_kwargs is not None:
-                kwargs = result_kwargs
+            hook(self, args, kwargs)
         bw_hook = None
         if len(full_backward_hooks) > 0:
             bw_hook = full_hooks.BackwardHook(self, full_backward_hooks)

msprobe/pytorch/hook_module/jit_script_wrapper.py

@@ -0,0 +1,33 @@
+# 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 torch
+
+from msprobe.pytorch.hook_module.api_register import get_api_register
+
+
+def wrap_jit_script_func():
+    def patched_script(*args, **kwargs):
+        all_api_registered = api_register.all_api_registered
+        if all_api_registered:
+            api_register.restore_all_api()
+        result = original_script(*args, **kwargs)
+        if all_api_registered:
+            api_register.register_all_api()
+        return result
+
+    original_script = torch.jit.script
+    api_register = get_api_register()
+    torch.jit.script = patched_script

msprobe/pytorch/hook_module/pt_hook_manager.py

@@ -0,0 +1,68 @@
+# 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 contextlib import nullcontext
+
+from msprobe.core.common.const import Const
+from msprobe.core.common.utils import replace_last_occurrence
+from msprobe.core.hook_manager import BaseHookManager, HookSet
+from msprobe.pytorch.common.utils import is_recomputation, torch_version_above_or_equal_2
+from msprobe.pytorch.hook_module.hook_module import HOOKModule
+
+
+class PytorchHookManager(BaseHookManager): 
+    @property
+    def _is_recompute(self):
+        return is_recomputation()
+
+    @staticmethod
+    def _no_grad_context():
+        return nullcontext()
+
+    @staticmethod
+    def _add_count(name):
+        HOOKModule.add_module_count(name)
+
+    @staticmethod
+    def _process_kwargs_and_output(module, hook_type, kwargs_or_output, output_or_kwargs):
+        kwargs = kwargs_or_output if torch_version_above_or_equal_2 else {}
+        output = output_or_kwargs if torch_version_above_or_equal_2 else kwargs_or_output
+        return kwargs, output
+    
+    def build_hook(self, hook_type, name):
+        if hook_type == Const.API:
+            full_forward_name = name + str(HOOKModule.get_module_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)        
+        )
+        return hookset
+    
+    def _need_exchange(self, module):
+        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.named_parameters(recurse=False)
+                    }
+        return params_dict