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

msprobe/mindspore/debugger/precision_debugger.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,20 @@
 # limitations under the License.
 
 import os
-from collections import defaultdict
+from collections import defaultdict, namedtuple
 
 import mindspore as ms
 from mindspore._c_expression import MSContext
 
-from msprobe.core.common.const import Const, MsgConst
+from msprobe.core.common.const import Const, FileCheckConst, MsgConst
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.common.file_utils import FileChecker
+from msprobe.core.common.utils import get_real_step_or_rank
 from msprobe.mindspore.cell_processor import CellProcessor
 from msprobe.mindspore.common.const import Const as MsConst
+from msprobe.mindspore.common.utils import set_register_backward_hook_functions
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
+from msprobe.mindspore.dump.hook_cell.api_registry import api_register
 from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 from msprobe.mindspore.grad_probe.grad_monitor import GradientMonitor
 from msprobe.mindspore.ms_config import parse_json_config
@@ -30,12 +35,21 @@ from msprobe.mindspore.runtime import Runtime
 from msprobe.mindspore.service import Service
 from msprobe.mindspore.task_handler_factory import TaskHandlerFactory
 
+try:
+    from msprobe.lib import _msprobe_c
+except ImportError:
+    _msprobe_c = None
+
+
+ConfigParameters = namedtuple("ConfigParameters", ["config_path", "task", "dump_path", "level"])
+
 
 class PrecisionDebugger:
     _instance = None
     task_not_need_service = [Const.GRAD_PROBE]
 
-    def __new__(cls, config_path=None, opt=None):
+    def __new__(cls, config_path=None, task=None, dump_path=None,
+                level=None, step=None, opt=None):
         if not cls._instance:
             cls._instance = super().__new__(cls)
             cls._instance.initialized = False
@@ -44,22 +58,65 @@ class PrecisionDebugger:
             cls.first_start = False
         return cls._instance
 
-    def __init__(self, config_path=None):
+    def __init__(self, config_path=None, task=None, dump_path=None,
+                 level=None, step=None):
         if self.initialized:
             return
         self.initialized = True
+
+        set_register_backward_hook_functions()
+
         if not config_path:
             config_path = os.path.join(os.path.dirname(__file__), "../../config.json")
+
+        config_params = ConfigParameters(config_path, task, dump_path, level)
+        self.check_input_params(config_params)
+
         common_config, task_config = parse_json_config(config_path)
+        common_config.task = task if task else common_config.task
         self.task = common_config.task
         if self.task == Const.GRAD_PROBE:
             self.gm = GradientMonitor(common_config, task_config)
             return
+        common_config.step = get_real_step_or_rank(
+            step, Const.STEP) if step is not None else common_config.step
+        common_config.level = level if level else common_config.level
+        common_config.dump_path = dump_path if dump_path else common_config.dump_path
         self.config = DebuggerConfig(common_config, task_config)
 
+        if _msprobe_c:
+            _msprobe_c._PrecisionDebugger(framework="MindSpore", config_path=config_path)
+
+        self.config.execution_mode = self._get_execution_mode()
+        if self._need_service():
+            self.service = Service(self.config)
+
         Runtime.step_count = 0
         Runtime.is_running = False
 
+    @staticmethod
+    def check_input_params(args):
+        if args.config_path is not None:
+            if not isinstance(args.config_path, str):
+                raise MsprobeException(
+                    MsprobeException.INVALID_PARAM_ERROR, f"config_path must be a string")
+            file_checker = FileChecker(
+                file_path=args.config_path, path_type=FileCheckConst.FILE, file_type=FileCheckConst.JSON_SUFFIX)
+            file_checker.common_check()
+
+        if args.task is not None and args.task not in Const.TASK_LIST:
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"task must be one of {Const.TASK_LIST}")
+
+        if args.dump_path is not None:
+            if not isinstance(args.dump_path, str):
+                raise MsprobeException(
+                    MsprobeException.INVALID_PARAM_ERROR, f"dump_path must be a string")
+
+        if args.level is not None and args.level not in Const.LEVEL_LIST:
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"level must be one of {Const.LEVEL_LIST}")
+
     @staticmethod
     def _get_execution_mode():
         jit_level = ms.context.get_jit_config().get(MsConst.JIT_LEVEL)
@@ -78,11 +135,23 @@ class PrecisionDebugger:
         else:
             return MsConst.PYNATIVE_MODE
 
+    @staticmethod
+    def _is_graph_dump(config):
+        if config.level != MsConst.KERNEL:
+            return False
+        if not config.list or len(config.list) > 1:
+            return True
+        if '-' in config.list[0] or '/' in config.list[0]:
+            return True
+        return False
+
     @classmethod
     def start(cls, model=None):
         instance = cls._instance
         if not instance:
             raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if _msprobe_c:
+            _msprobe_c._PrecisionDebugger().start()
         if instance.task in PrecisionDebugger.task_not_need_service:
             return
 
@@ -93,6 +162,7 @@ class PrecisionDebugger:
             instance.service.start(model)
         else:
             if not instance.first_start:
+                api_register.api_set_ori_func()
                 handler = TaskHandlerFactory.create(instance.config)
                 handler.handle()
 
@@ -102,18 +172,15 @@ class PrecisionDebugger:
     @classmethod
     def forward_backward_dump_end(cls):
         instance = cls._instance
-        if not instance:
-            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
-        if instance.task in PrecisionDebugger.task_not_need_service:
-            return
-        if instance.service:
-            instance.service.forward_backward_dump_end()
+        instance.stop()
 
     @classmethod
     def stop(cls):
         instance = cls._instance
         if not instance:
             raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if _msprobe_c:
+            _msprobe_c._PrecisionDebugger().stop()
         if instance.task == Const.GRAD_PROBE:
             instance.gm.stop()
         if instance.task in PrecisionDebugger.task_not_need_service:
@@ -127,6 +194,8 @@ class PrecisionDebugger:
         instance = cls._instance
         if not instance:
             raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if _msprobe_c:
+            _msprobe_c._PrecisionDebugger().step()
         if instance.task in PrecisionDebugger.task_not_need_service:
             return
         if instance.service:
@@ -153,4 +222,4 @@ class PrecisionDebugger:
         if instance.config.execution_mode != MsConst.PYNATIVE_MODE:
             return False
         else:
-            return instance.config.task != Const.FREE_BENCHMARK and instance.config.level != MsConst.KERNEL
+            return instance.config.task != Const.FREE_BENCHMARK and not instance._is_graph_dump(instance.config)

msprobe/mindspore/dump/hook_cell/api_registry.py

@@ -1,4 +1,5 @@
-# Copyright 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");
 # you may not use this file except in compliance with the License.
@@ -11,7 +12,6 @@
 # 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 import Tensor, ops, mint
 from mindspore.mint.nn import functional
@@ -20,8 +20,15 @@ from mindspore.communication import comm_func
 
 from msprobe.mindspore.dump.hook_cell.wrap_api import (HOOKTensor, HOOKStubTensor, HOOKFunctionalOP,
                                                        HOOKMintOP, HOOKMintNNFunctionalOP, HOOKDistributedOP,
-                                                       get_wrap_api_list, setup_hooks)
+                                                       HOOKTorchOP, HOOKTorchTensor, HOOKTorchFunctionalOP,
+                                                       HOOKTorchDistributedOP, HOOKTorchNpuOP,
+                                                       get_wrap_api_list, get_wrap_torch_api_list, setup_hooks)
 from msprobe.core.common.utils import Const
+from msprobe.mindspore.common.utils import is_mindtorch
+
+if is_mindtorch():
+    import torch
+    import torch_npu
 
 
 def stub_method(method):
@@ -40,6 +47,12 @@ class ApiRegistry:
         self.distributed_ori_attr = {}
         self.norm_inner_ops_ori_attr = {}
 
+        self.torch_ori_attr = {}
+        self.torch_tensor_ori_attr = {}
+        self.torch_functional_ori_attr = {}
+        self.torch_distributed_ori_attr = {}
+        self.torch_npu_ori_attr = {}
+
         self.tensor_hook_attr = {}
         self.stub_tensor_hook_attr = {}
         self.functional_hook_attr = {}
@@ -48,6 +61,12 @@ class ApiRegistry:
         self.distibuted_hook_attr = {}
         self.norm_inner_ops_hook_attr = {}
 
+        self.torch_hook_attr = {}
+        self.torch_tensor_hook_attr = {}
+        self.torch_functional_hook_attr = {}
+        self.torch_distributed_hook_attr = {}
+        self.torch_npu_hook_attr = {}
+
         self.norm_inner_ops = ["norm", "square", "sqrt", "is_complex"]
 
     @staticmethod
@@ -82,22 +101,71 @@ class ApiRegistry:
         self.set_api_attr(ops, self.norm_inner_ops_ori_attr)
 
     def api_set_hook_func(self):
-        self.set_api_attr(Tensor, self.tensor_hook_attr)
-        self.set_api_attr(StubTensor, self.stub_tensor_hook_attr)
-        self.set_api_attr(ops, self.functional_hook_attr)
-        self.set_api_attr(mint, self.mint_ops_hook_attr)
-        self.set_api_attr(functional, self.mint_func_ops_hook_attr)
-        self.set_api_attr(comm_func, self.distibuted_hook_attr)
+        if is_mindtorch():
+            self.set_api_attr(torch, self.torch_hook_attr)
+            self.set_api_attr(torch.Tensor, self.torch_tensor_hook_attr)
+            self.set_api_attr(torch.nn.functional, self.torch_functional_hook_attr)
+            self.set_api_attr(torch.distributed, self.torch_distributed_hook_attr)
+            self.set_api_attr(torch_npu, self.torch_npu_hook_attr)
+        else:
+            self.set_api_attr(Tensor, self.tensor_hook_attr)
+            self.set_api_attr(StubTensor, self.stub_tensor_hook_attr)
+            self.set_api_attr(ops, self.functional_hook_attr)
+            self.set_api_attr(mint, self.mint_ops_hook_attr)
+            self.set_api_attr(functional, self.mint_func_ops_hook_attr)
+            self.set_api_attr(comm_func, self.distibuted_hook_attr)
 
     def api_set_ori_func(self):
-        self.set_api_attr(Tensor, self.tensor_ori_attr)
-        self.set_api_attr(StubTensor, self.stub_tensor_ori_attr)
-        self.set_api_attr(ops, self.functional_ori_attr)
-        self.set_api_attr(mint, self.mint_ops_ori_attr)
-        self.set_api_attr(functional, self.mint_func_ops_ori_attr)
-        self.set_api_attr(comm_func, self.distributed_ori_attr)
+        if is_mindtorch():
+            self.set_api_attr(torch, self.torch_ori_attr)
+            self.set_api_attr(torch.Tensor, self.torch_tensor_ori_attr)
+            self.set_api_attr(torch.nn.functional, self.torch_functional_ori_attr)
+            self.set_api_attr(torch.distributed, self.torch_distributed_ori_attr)
+            self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
+        else:
+            self.set_api_attr(Tensor, self.tensor_ori_attr)
+            self.set_api_attr(StubTensor, self.stub_tensor_ori_attr)
+            self.set_api_attr(ops, self.functional_ori_attr)
+            self.set_api_attr(mint, self.mint_ops_ori_attr)
+            self.set_api_attr(functional, self.mint_func_ops_ori_attr)
+            self.set_api_attr(comm_func, self.distributed_ori_attr)
 
     def initialize_hook(self, hook):
+        setup_hooks(hook)
+        if is_mindtorch():
+            wrap_torch_api_name = get_wrap_torch_api_list()
+            self.store_ori_attr(torch,
+                                wrap_torch_api_name.torch_api_names, self.torch_ori_attr)
+            self.store_ori_attr(torch.Tensor,
+                                wrap_torch_api_name.tensor_api_names, self.torch_tensor_ori_attr)
+            self.store_ori_attr(torch.nn.functional,
+                                wrap_torch_api_name.functional_api_names, self.torch_functional_ori_attr)
+            self.store_ori_attr(torch.distributed,
+                                wrap_torch_api_name.distributed_api_names, self.torch_distributed_ori_attr)
+            self.store_ori_attr(torch_npu,
+                                wrap_torch_api_name.npu_api_names, self.torch_npu_ori_attr)
+            for attr_name in dir(HOOKTorchOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_hook_attr[api_name] = getattr(HOOKTorchOP, attr_name)
+            for attr_name in dir(HOOKTorchTensor):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_tensor_hook_attr[api_name] = getattr(HOOKTorchTensor, attr_name)
+            for attr_name in dir(HOOKTorchFunctionalOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_functional_hook_attr[api_name] = getattr(HOOKTorchFunctionalOP, attr_name)
+            for attr_name in dir(HOOKTorchDistributedOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_distributed_hook_attr[api_name] = getattr(HOOKTorchDistributedOP, attr_name)
+            for attr_name in dir(HOOKTorchNpuOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_npu_hook_attr[api_name] = getattr(HOOKTorchNpuOP, attr_name)
+            return
+
         wrap_api_name = get_wrap_api_list()
         self.store_ori_attr(Tensor, wrap_api_name.tensor_api_names, self.tensor_ori_attr)
         self.store_ori_attr(StubTensor, wrap_api_name.stub_tensor_api_names, self.stub_tensor_ori_attr)
@@ -106,7 +174,6 @@ class ApiRegistry:
         self.store_ori_attr(functional, wrap_api_name.mint_nn_func_api_names, self.mint_func_ops_ori_attr)
         self.store_ori_attr(comm_func, wrap_api_name.distributed_api_names, self.distributed_ori_attr)
         self.store_ori_attr(ops, self.norm_inner_ops, self.norm_inner_ops_ori_attr)
-        setup_hooks(hook)
         for attr_name in dir(HOOKTensor):
             if attr_name.startswith(Const.ATTR_NAME_PREFIX):
                 api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]

msprobe/mindspore/dump/hook_cell/hook_cell.py

@@ -1,4 +1,5 @@
-# Copyright 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");
 # you may not use this file except in compliance with the License.
@@ -11,45 +12,66 @@
 # 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 collections import defaultdict
 
 from mindspore import nn
 
-from msprobe.core.common.const import Const
-
-
-class HOOKCell(nn.Cell):
-    cell_count = defaultdict(int)
-    g_stop_hook = False
-
-    def __init__(self, build_hook) -> None:
-        super(HOOKCell, self).__init__()
-        self.changed_status = False
-        self.input_kwargs = {}
-        self.prefix = ""
-        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
-
-            HOOKCell.cell_count[self.prefix] += 1
-            self.prefix = self.prefix + str(HOOKCell.cell_count[self.prefix] - 1) + Const.SEP
-            forward_hook, backward_hook = build_hook(self.prefix)
-            self.register_forward_hook(forward_hook)
-            self.register_backward_hook(backward_hook)
-
-    # 重载call，加全局标志。
-    def __call__(self, *args, **kwargs):
-        try:
-            self.input_kwargs = kwargs
-            out = super(HOOKCell, self).__call__(*args, **kwargs)
-        except Exception as e:
-            raise e
-        finally:
-            if self.changed_status:
-                self.changed_status = False
-                HOOKCell.g_stop_hook = False
-        return out
+from msprobe.mindspore.common.utils import is_mindtorch, register_backward_hook_functions
+
+
+def add_cell_count(name):
+    HOOKCell.cell_count[name] += 1
+
+
+def get_cell_count(name):
+    return HOOKCell.cell_count[name]
+
+
+def __init__(self, build_hook) -> None:
+    super(HOOKCell, self).__init__()
+    self.changed_status = False
+    self.input_kwargs = {}
+    self.prefix = ""
+    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)
+
+
+# 重载call，加全局标志。
+def __call__(self, *args, **kwargs):
+    try:
+        self.input_kwargs = kwargs
+        out = super(HOOKCell, self).__call__(*args, **kwargs)
+    except Exception as e:
+        raise e
+    finally:
+        if self.changed_status:
+            self.changed_status = False
+            HOOKCell.g_stop_hook = False
+    return out
+
+
+hook_cell_dict = {
+    "cell_count": defaultdict(int),
+    "g_stop_hook": False,
+    "add_cell_count": staticmethod(add_cell_count),
+    "get_cell_count": staticmethod(get_cell_count),
+    "__init__": __init__,
+    "__call__": __call__
+}
+
+if is_mindtorch():
+    import torch
+    HOOKCell = type("HOOKCell", (torch.nn.Module,), hook_cell_dict)
+else:
+    HOOKCell = type("HOOKCell", (nn.Cell,), hook_cell_dict)

msprobe/mindspore/dump/hook_cell/primitive_hooks.py

@@ -135,6 +135,34 @@ class PrimitiveHookService:
                 return tuple(hooked_outputs)
             return out
 
+        def pre_forward_hook(primitive_name, primitive_instance, args, kwargs):
+            module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
+            try:
+                self.service_instance.data_collector.forward_input_data_collect(
+                    primitive_name,
+                    primitive_instance,
+                    os.getpid(),
+                    module_input_output
+                )
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during forward input data collection: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+
+        def post_forward_hook(primitive_name, primitive_instance, args, kwargs, output):
+            module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+            try:
+                self.service_instance.data_collector.forward_output_data_collect(
+                    primitive_name,
+                    primitive_instance,
+                    os.getpid(),
+                    module_input_output
+                )
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during forward output data collection: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+
         def wrapped_primitive_call(instance_self, *args, **kwargs):
             """
             包装后的 primitive 调用函数，添加输入和输出的 hook。
@@ -163,27 +191,17 @@ class PrimitiveHookService:
                              f"primitive_name: {primitive_name}")
                 raise DumpException(DumpException.INPUT_HOOK_ERROR) from exception
 
+            forward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.FORWARD}"
+            self.service_instance.data_collector.update_api_or_module_name(forward_primitive_name)
+
+            pre_forward_hook(forward_primitive_name, instance_self, hooked_inputs, kwargs)
             try:
                 out = origin_func(*hooked_inputs, **kwargs)
             except Exception as exception:
                 logger.error(f"This is a primitive op dump error during function call: {exception}, "
                              f"primitive_name: {primitive_name}")
                 raise DumpException(DumpException.FUNCTION_CALL_ERROR) from exception
-
-            forward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.FORWARD}"
-            self.service_instance.data_collector.update_api_or_module_name(forward_primitive_name)
-            if self.service_instance.data_collector:
-                module_input_output = ModuleForwardInputsOutputs(args=hooked_inputs, kwargs=kwargs, output=out)
-                try:
-                    self.service_instance.data_collector.forward_data_collect(forward_primitive_name, instance_self,
-                                                                              os.getpid(), module_input_output)
-                except Exception as exception:
-                    logger.error(f"This is a primitive op dump error during forward data collection: {exception}, "
-                                 f"primitive_name: {primitive_name}")
-                    raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
-
-                if self.service_instance.data_collector.if_return_forward_new_output():
-                    out = self.service_instance.data_collector.get_forward_new_output()
+            post_forward_hook(forward_primitive_name, instance_self, hooked_inputs, kwargs, out)
 
             try:
                 out = hook_primitive_outputs(out, captured_grads_output, updated_primitive_name)

msprobe/mindspore/dump/hook_cell/support_wrap_ops.yaml

@@ -15,7 +15,7 @@
 
 # List of ops that register hooks
 
-  
+
 ops:
   - adaptive_avg_pool1d
   - adaptive_avg_pool2d
@@ -85,6 +85,7 @@ ops:
   - relu6
   - celu
   - rrelu
+  - rms_norm
   - selu
   - sigmoid
   - silu
@@ -553,6 +554,7 @@ tensor:
   - acos
   - acosh
   - add
+  - add_
   - addbmm
   - addcdiv
   - addcmul
@@ -607,6 +609,7 @@ tensor:
   - diff
   - digamma
   - div
+  - div_
   - divide
   - equal
   - erf
@@ -739,6 +742,8 @@ tensor:
   - square
   - squeeze
   - std
+  - sub
+  - sub_
   - subtract
   - subtract
   - svd
@@ -983,6 +988,7 @@ mint.nn.functional:
   - one_hot_ext
   - pad
   - relu
+  - relu_
   - sigmoid
   - silu
   - softmax
@@ -1017,3 +1023,7 @@ communication.comm_func:
   - broadcast
   - gather_into_tensor
   - scatter_tensor
+  - send
+  - recv
+  - isend
+  - irecv