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

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,35 +21,33 @@ import torch
 import torch.nn as nn
 import torch.utils.hooks as full_hooks
 
-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):
     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)
+
+            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)
+                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
@@ -78,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)
@@ -111,6 +103,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/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

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