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

msprobe/mindspore/dump/hook_cell/wrap_api.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");
@@ -23,10 +23,16 @@ from mindspore.mint.nn import functional
 from msprobe.core.common.const import Const
 from msprobe.core.common.file_utils import load_yaml
 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
 
+if is_mindtorch():
+    import torch
+    import torch_npu
+
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, MsConst.SUPPORTED_API_LIST_FILE)
+torch_yaml_path = os.path.join(cur_path, "../../../pytorch/hook_module", MsConst.SUPPORTED_API_LIST_FILE)
 
 
 class HOOKTensor(object):
@@ -53,6 +59,26 @@ class HOOKDistributedOP(object):
     pass
 
 
+class HOOKTorchOP(object):
+    pass
+
+
+class HOOKTorchTensor(object):
+    pass
+
+
+class HOOKTorchFunctionalOP(object):
+    pass
+
+
+class HOOKTorchDistributedOP(object):
+    pass
+
+
+class HOOKTorchNpuOP(object):
+    pass
+
+
 class ApiTemplate(HOOKCell):
     def __init__(self, api_name, api_dict, prefix, hook):
         self.api_name = api_name
@@ -60,7 +86,30 @@ class ApiTemplate(HOOKCell):
         self.prefix_api_name = prefix + str(api_name.split(Const.SEP)[-1]) + Const.SEP
         super().__init__(hook)
 
+    @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):
+            if kwargs.get("async_op") or self.api_name in ["isend", "irecv"]:
+                output = self.async_to_sync(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)
@@ -77,6 +126,15 @@ class WrapApiName:
         self.distributed_api_names = distributed_api_names
 
 
+class WrapTorchApiName:
+    def __init__(self, torch_api_names, tensor_api_names, functional_api_names, distributed_api_names, npu_api_names):
+        self.torch_api_names = torch_api_names
+        self.tensor_api_names = tensor_api_names
+        self.functional_api_names = functional_api_names
+        self.distributed_api_names = distributed_api_names
+        self.npu_api_names = npu_api_names
+
+
 def get_wrap_api_list():
     api_list = load_yaml(yaml_path)
     tensor_api = api_list.get(MsConst.SUPPORTED_TENSOR_LIST_KEY)
@@ -93,6 +151,21 @@ def get_wrap_api_list():
     return wrap_api_name
 
 
+def get_wrap_torch_api_list():
+    api_list = load_yaml(torch_yaml_path)
+    torch_api = api_list.get("torch")
+    tensor_api = api_list.get("tensor")
+    functional_api = api_list.get("functional")
+    distributed_api = api_list.get("distributed")
+    npu_api = api_list.get("torch_npu")
+    wrap_api_name = WrapTorchApiName(set(torch_api) & set(dir(torch)),
+                                     set(tensor_api) & set(dir(torch.Tensor)),
+                                     set(functional_api) & set(dir(torch.nn.functional)),
+                                     set(distributed_api) & set(dir(torch.distributed)),
+                                     set(npu_api) & set(dir(torch_npu)))
+    return wrap_api_name
+
+
 def wrap_api_func(api_name, api_dict, prefix, hook):
     def api_function(*args, **kwargs):
         return ApiTemplate(api_name, api_dict, prefix, hook)(*args, **kwargs)
@@ -106,6 +179,24 @@ def wrap_api_func_and_bind(api_list, api_dict, prefix, hook, hook_class):
 
 
 def setup_hooks(hook):
+    if is_mindtorch():
+        torch_wrap_api_name = get_wrap_torch_api_list()
+        wrap_api_func_and_bind(torch_wrap_api_name.torch_api_names,
+                               {f: getattr(torch, f) for f in dir(torch)},
+                               MsConst.TORCH_DATA_PREFIX, hook, HOOKTorchOP)
+        wrap_api_func_and_bind(torch_wrap_api_name.tensor_api_names,
+                               {f: getattr(torch.Tensor, f) for f in dir(torch.Tensor)},
+                               MsConst.TENSOR_DATA_PREFIX, hook, HOOKTorchTensor)
+        wrap_api_func_and_bind(torch_wrap_api_name.functional_api_names,
+                               {f: getattr(torch.nn.functional, f) for f in dir(torch.nn.functional)},
+                               MsConst.OPS_DATA_PREFIX, hook, HOOKTorchFunctionalOP)
+        wrap_api_func_and_bind(torch_wrap_api_name.distributed_api_names,
+                               {f: getattr(torch.distributed, f) for f in dir(torch.distributed)},
+                               MsConst.DISTRIBUTED_DATA_PREFIX, hook, HOOKTorchDistributedOP)
+        wrap_api_func_and_bind(torch_wrap_api_name.npu_api_names, {f: getattr(torch_npu, f) for f in dir(torch_npu)},
+                               MsConst.TORCH_NPU_DATA_PREFIX, hook, HOOKTorchNpuOP)
+        return
+
     wrap_api_name = get_wrap_api_list()
     wrap_api_func_and_bind(wrap_api_name.tensor_api_names, {f: getattr(Tensor, f) for f in dir(Tensor)},
                            MsConst.TENSOR_DATA_PREFIX, hook, HOOKTensor)

msprobe/mindspore/dump/jit_dump.py

@@ -16,14 +16,15 @@
 import os
 from collections import defaultdict
 
-from mindspore import Tensor
 from mindspore._c_expression import PyNativeExecutor_
-from mindspore.common.api import _MindsporeFunctionExecutor
+try:
+    from mindspore.common.api import _MindsporeFunctionExecutor
+except ImportError:
+    from mindspore.common.api import _JitExecutor as _MindsporeFunctionExecutor
 
 from msprobe.core.common.log import logger
-from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs, ModuleBackwardInputsOutputs
 from msprobe.core.common.const import Const
-from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs
+from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs, ModuleBackwardInputsOutputs
 from msprobe.mindspore.dump.hook_cell.api_registry import api_register
 
 
@@ -40,8 +41,8 @@ def dump_jit(name, in_feat, out_feat, is_forward):
     if JitDump.need_dump():
         if is_forward:
             JitDump.jit_count[result] += 1
-            name_template = Const.JIT + Const.SEP + result + Const.SEP + str(JitDump.jit_count[result]) + Const.SEP + \
-                            Const.FORWARD
+            name_template = (Const.JIT + Const.SEP + result + Const.SEP +
+                             str(JitDump.jit_count[result]) + Const.SEP + Const.FORWARD)
             JitDump.data_collector.update_api_or_module_name(name_template)
             module_input_output = ModuleForwardInputsOutputs(args=in_feat, kwargs={}, output=out_feat)
             JitDump.data_collector.forward_data_collect(name_template, None, pid, module_input_output)

msprobe/mindspore/dump/kernel_dump/kernel_config.py

@@ -0,0 +1,33 @@
+# 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.
+
+import os
+
+from msprobe.core.common.file_utils import save_json
+
+
+def create_kernel_config_json(dump_path, cur_rank):
+    kernel_config_name = "kernel_config.json" if cur_rank == '' else f"kernel_config_{cur_rank}.json"
+    kernel_config_path = os.path.join(dump_path, kernel_config_name)
+    config_info = {
+        "dump": {
+            "dump_list": [],
+            "dump_path": dump_path,
+            "dump_mode": "all",
+            "dump_op_switch": "on"
+        }
+    }
+    save_json(kernel_config_path, config_info, indent=4)
+    return kernel_config_path

msprobe/mindspore/dump/kernel_graph_dump.py

@@ -56,6 +56,13 @@ class KernelGraphDump:
                 self.dump_json["common_dump_settings"]["input_output"] = 2
 
     def handle(self):
+        try:
+            from msprobe.lib import _msprobe_c
+            return
+        except ImportError:
+            # 如果没有_msprobe_ce_c走MindSpore老流程
+            logger.info("Module _msprobe_c has not been installed, use interface in mindspore instead.")
+
         if os.getenv("GRAPH_OP_RUN") == "1":
             raise Exception("Must run in graph mode, not kbk mode")
         json_path = self.dump_json["common_dump_settings"]["path"]

msprobe/mindspore/free_benchmark/api_pynative_self_check.py

@@ -19,7 +19,6 @@ import os
 import traceback
 
 import mindspore as ms
-
 from msprobe.core.common.const import Const
 from msprobe.core.common.exceptions import DistributedNotInitializedError
 from msprobe.core.common.file_utils import check_path_length, load_yaml
@@ -29,6 +28,7 @@ from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.common.utils import get_rank_if_initialized
 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.free_benchmark.common.config import Config
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
 from msprobe.mindspore.free_benchmark.common.utils import Tools
@@ -63,7 +63,10 @@ class ApiPyNativeSelfCheck:
         api_register.initialize_hook(self.build_hook)
         api_register.api_set_hook_func()
 
-    def build_hook(self, api_name_with_id):
+    def build_hook(self, api_name):
+        def pre_hook(cell, input_data):
+            return None
+
         def forward_hook(api_name_with_id, cell, input_data, output_data):
             ret = None
 
@@ -85,7 +88,10 @@ class ApiPyNativeSelfCheck:
         def backward_hook(cell, grad_input, grad_output):
             pass
 
+        HOOKCell.get_cell_count(api_name)
+        api_name_with_id = api_name + str(HOOKCell.get_cell_count(api_name)) + Const.SEP
         forward_hook = functools.partial(forward_hook, api_name_with_id)
+        HOOKCell.add_cell_count(api_name)
 
         def wrap_forward_hook(cell, input_data, output_data):
             return forward_hook(cell, input_data, output_data)
@@ -93,7 +99,10 @@ class ApiPyNativeSelfCheck:
         def wrap_backward_hook(cell, grad_input, grad_output):
             return backward_hook(cell, grad_input, grad_output)
 
-        return wrap_forward_hook, wrap_backward_hook
+        def pre_backward_hook(cell, grad_input):
+            return None
+
+        return pre_hook, wrap_forward_hook, wrap_backward_hook, pre_backward_hook
 
     def store_original_func(self):
         for api_name in self.api_list:
@@ -138,7 +147,7 @@ def get_module(api_name):
     module_obj = importlib.import_module(func_name_list[0])
     for i, module_name in enumerate(func_name_list[1:-1]):
         if not hasattr(module_obj, module_name):
-            importlib.import_module(f"{Const.SEP.join(func_name_list[:i+2])}")
+            importlib.import_module(f"{Const.SEP.join(func_name_list[:i + 2])}")
         module_obj = getattr(module_obj, module_name)
     orig_func = getattr(module_obj, func_name)
 

msprobe/mindspore/free_benchmark/perturbation/bit_noise.py

@@ -35,12 +35,12 @@ class BitNoisePerturbation(BasePerturbation):
                 noise_type = list(FreeBenchmarkConst.MS_NUMPY_DTYPE_DICT.keys())[
                              list(FreeBenchmarkConst.MS_NUMPY_DTYPE_DICT.values()).index(bit_len_type)]
                 noise = ops.full(inputs.shape, 1, dtype=noise_type)
-                input_np = inputs.contiguous().asnumpy()
+                input_np = inputs.asnumpy()
                 input_np_int = input_np.view(bit_len_type)
                 result = Tensor(input_np_int)
                 result = ops.where(ops.abs(inputs) > sub_normal,
                                    ops.bitwise_xor(result, noise), result)
-                result_np = result.contiguous().asnumpy()
+                result_np = result.asnumpy()
                 result_np_float = result_np.view(FreeBenchmarkConst.MS_NUMPY_DTYPE_DICT.get(inputs.dtype))
                 self.is_fuzzed = True
                 return Tensor(result_np_float)

msprobe/mindspore/grad_probe/grad_analyzer.py

@@ -16,6 +16,7 @@
 import multiprocessing
 import os
 import time
+from dataclasses import dataclass
 from multiprocessing import Process
 from typing import List
 
@@ -23,6 +24,7 @@ import mindspore as ms
 import numpy as np
 from mindspore.common.parameter import Parameter
 from mindspore.communication import get_rank
+
 from msprobe.core.common.file_utils import (create_directory, check_file_or_directory_path,
                                             write_csv, remove_path, move_file, load_npy)
 from msprobe.core.grad_probe.constant import GradConst
@@ -31,6 +33,16 @@ from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.grad_probe.global_context import grad_context, GlobalContext
 
 
+@dataclass
+class GradDumpConfig:
+    dump_dir: str
+    g_name: str
+    dump_step: Parameter
+    grad: ms.Tensor
+    level: str
+    bounds: List
+
+
 def get_rank_id():
     try:
         rank_id = get_rank()
@@ -40,35 +52,35 @@ def get_rank_id():
 
 
 @ms.jit
-def grad_dump(dump_dir: str, g_name: str, dump_step: Parameter, grad: ms.Tensor, level: str, bounds: List):
+def grad_dump(config: GradDumpConfig):
     """
     Dump gradient statistic data.
         level0: [step, max, min, norm, shape_dim, shape]
         level1: [step, max, min, norm, shape_dim, shape] + grad_bool_data
         level2: [step, max, min, norm, shape_dim, shape, dist_dim, dist] + grad_bool_data
     """
-    dump_path = os.path.join(dump_dir, g_name)
+    dump_path = os.path.join(config.dump_dir, config.g_name)
     dump_dir_path = dump_path + "_dir"
     save_op = ms.ops.TensorDump()
 
-    grad_flat = grad.reshape(-1)
+    grad_flat = config.grad.reshape(-1)
     max_val = grad_flat.max(axis=0).float()
     min_val = grad_flat.min(axis=0).float()
     norm_val = grad_flat.norm(ord=2).float()
-    shape = grad.shape
-    extrem_list = [dump_step[0].float(), max_val, min_val, norm_val]
+    shape = config.grad.shape
+    extrem_list = [config.dump_step[0].float(), max_val, min_val, norm_val]
     extrem_stat = ms.ops.stack(extrem_list)
     shape_list = [len(shape)] + list(shape)
     shape_stat = ms.Tensor(shape_list).float()
     level0_stat = ms.ops.concat((extrem_stat, shape_stat), axis=0)
     level_stat = level0_stat
 
-    if level == GradConst.LEVEL2:
-        zero_grad = (grad == 0).sum()
-        dist_dim = ms.Tensor([len(bounds) + 2]).float()
-        bucket_result = ms.ops.bucketize(grad.float(), bounds)
+    if config.level == GradConst.LEVEL2:
+        zero_grad = (config.grad == 0).sum()
+        dist_dim = ms.Tensor([len(config.bounds) + 2]).float()
+        bucket_result = ms.ops.bucketize(config.grad.float(), config.bounds)
         bucket_result = bucket_result.astype(ms.int8)
-        dist_stat = [(bucket_result == i).sum() for i in range(len(bounds) + 1)]
+        dist_stat = [(bucket_result == i).sum() for i in range(len(config.bounds) + 1)]
         dist_stat.append(zero_grad)
         dist_stat.append(ms.Tensor(1, dtype=ms.int64))  # make sure dist_stat is not empty
         dist_stat = ms.ops.stack(dist_stat, axis=0).float()
@@ -76,8 +88,8 @@ def grad_dump(dump_dir: str, g_name: str, dump_step: Parameter, grad: ms.Tensor,
         level_stat = level2_stat
 
     save_op(dump_path, level_stat)
-    if level == GradConst.LEVEL1 or level == GradConst.LEVEL2:
-        grad_direction = grad > 0
+    if config.level == GradConst.LEVEL1 or config.level == GradConst.LEVEL2:
+        grad_direction = config.grad > 0
         save_op(dump_dir_path, grad_direction)
 
 

msprobe/mindspore/grad_probe/hook.py

@@ -26,7 +26,7 @@ from msprobe.core.grad_probe.constant import GradConst
 from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.grad_probe.global_context import grad_context
 from msprobe.mindspore.grad_probe.grad_analyzer import csv_generator
-from msprobe.mindspore.grad_probe.grad_analyzer import grad_dump, get_rank_id
+from msprobe.mindspore.grad_probe.grad_analyzer import grad_dump, get_rank_id, GradDumpConfig
 from msprobe.mindspore.grad_probe.grad_stat_csv import GradStatCsv, CsvInput
 from msprobe.mindspore.grad_probe.utils import save_grad_direction, get_adapted_level
 
@@ -38,7 +38,14 @@ class HookInput:
 
     def __init__(self, opt) -> None:
         self.func = opt.construct
-        self.g_names = [param.name for param in opt._parameters]
+        if hasattr(opt, "_parameters"):
+            parameter_list = opt._parameters
+        elif hasattr(opt, "parameters"):
+            parameter_list = opt.parameters
+        else:
+            logger.error_log_with_exp("Given optimizer has no attributes: '_parameters' or 'parameters'. \
+                                      Please check the type of the given optimizer.", ValueError)
+        self.g_names = [param.name for param in parameter_list]
         self.param_list = grad_context.get_context(GradConst.PARAM_LIST)
         self.rank_id = get_rank_id()
         output_path = grad_context.get_context(GradConst.OUTPUT_PATH)
@@ -59,8 +66,10 @@ def hook_graph_mode_optimizer(opt, hook_input):
         for index, grad_value in enumerate(gradients):
             if hook_input.param_list and hook_input.g_names[index] not in hook_input.param_list:
                 continue
-            grad_dump(hook_input.dump_dir, hook_input.g_names[index], self.dump_step,
-                      grad_value, hook_input.level, hook_input.bounds)
+            conf = GradDumpConfig(dump_dir=hook_input.dump_dir, g_name=hook_input.g_names[index],
+                                  dump_step=self.dump_step, grad=grad_value, level=hook_input.level,
+                                  bounds=hook_input.bounds)
+            grad_dump(conf)
         ms.ops.TensorDump()(hook_input.step_finish_flag, self.dump_step)
         self.assignadd(self.dump_step, self.global_step_increase_tensor)
         out = hook_input.func(gradients)

msprobe/mindspore/mindtorch/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .mindtorch_adaptor import (_call_impl,
+                                register_full_backward_pre_hook,
+                                register_full_backward_hook)