mindstudio-probe 1.1.1__py3-none-any.whl → 1.2.1__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/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)

msprobe/mindspore/mindtorch/mindtorch_adaptor.py

@@ -0,0 +1,255 @@
+# From PyTorch:
+
+# Copyright (c) 2025      Huawei Technologies Co., Ltd
+# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
+# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
+# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
+# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
+# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
+# Copyright (c) 2011-2013 NYU                      (Clement Farabet)
+# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
+# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
+# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
+
+# From Caffe2:
+
+# Copyright (c) 2016-present, Facebook Inc. All rights reserved.
+
+# All contributions by Facebook:
+# Copyright (c) 2016 Facebook Inc.
+
+# All contributions by Google:
+# Copyright (c) 2015 Google Inc.
+# All rights reserved.
+
+# All contributions by Yangqing Jia:
+# Copyright (c) 2015 Yangqing Jia
+# All rights reserved.
+
+# All contributions by Kakao Brain:
+# Copyright 2019-2020 Kakao Brain
+
+# All contributions by Cruise LLC:
+# Copyright (c) 2022 Cruise LLC.
+# All rights reserved.
+
+# All contributions by Tri Dao:
+# Copyright (c) 2024 Tri Dao.
+# All rights reserved.
+
+# All contributions by Arm:
+# Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates
+
+# All contributions from Caffe:
+# Copyright(c) 2013, 2014, 2015, the respective contributors
+# All rights reserved.
+
+# All other contributions:
+# Copyright(c) 2015, 2016 the respective contributors
+# All rights reserved.
+
+# Caffe2 uses a copyright model similar to Caffe: each contributor holds
+# copyright over their contributions to Caffe2. The project versioning records
+# all such contribution and copyright details. If a contributor wants to further
+# mark their specific copyright on a particular contribution, they should
+# indicate their copyright solely in the commit message of the change when it is
+# committed.
+
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+
+# 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories
+#    America, IDIAP Research Institute and Huawei nor the names of its contributors
+#    may be used to endorse or promote products derived from this software without
+#    specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+import warnings
+
+import mindspore as ms
+from mindspore.ops.operations import _inner_ops as inner
+from torch.nn.modules.module import (_global_backward_pre_hooks, _global_backward_hooks,
+                                     _global_is_full_backward_hook, _global_forward_pre_hooks,
+                                     _global_forward_hooks, _global_forward_hooks_always_called)
+from torch.utils.hooks import RemovableHandle
+
+
+def _call_impl(self, *args, **kwargs):
+    forward_call = self.forward
+    if self.__ms_class__:
+        return forward_call(*args, **kwargs)
+
+    # If we don't have any hooks, we want to skip the rest of the logic in
+    # this function, and just call forward.
+    if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
+            or _global_backward_pre_hooks or _global_backward_hooks
+            or _global_forward_hooks or _global_forward_pre_hooks):
+        return forward_call(*args, **kwargs)
+
+    try:
+        result = None
+        called_always_called_hooks = set()
+
+        if self._backward_pre_hooks or _global_backward_pre_hooks:
+            _get_backward_pre_hooks(self)
+
+        if self._backward_hooks or _global_backward_hooks:
+            _get_backward_hooks(self)
+
+        if _global_forward_pre_hooks or self._forward_pre_hooks:
+            for hook_id, hook in (
+                *_global_forward_pre_hooks.items(),
+                *self._forward_pre_hooks.items(),
+            ):
+                if hook_id in self._forward_pre_hooks_with_kwargs:
+                    args_kwargs_result = hook(self, args, kwargs)  # type: ignore[misc]
+                    if args_kwargs_result is not None:
+                        if isinstance(args_kwargs_result, tuple) and len(args_kwargs_result) == 2:
+                            args, kwargs = args_kwargs_result
+                        else:
+                            raise RuntimeError(
+                                "forward pre-hook must return None or a tuple "
+                                f"of (new_args, new_kwargs), but got {args_kwargs_result}."
+                            )
+                else:
+                    args_result = hook(self, args)
+                    if args_result is not None:
+                        if not isinstance(args_result, tuple):
+                            args_result = (args_result,)
+                        args = args_result
+
+        bw_hook = None
+        if self._backward_hooks:
+            bw_hook = inner.CellBackwardHook(self.__class__.__name__ + "(" + str(id(self)) + ")",
+                                             self, self._backward_hooks)
+            bw_hook.register_backward_hook()
+            args = apply_backward_hook_on_tensors(bw_hook, args)
+
+        result = forward_call(*args, **kwargs)
+        if _global_forward_hooks or self._forward_hooks:
+            for hook_id, hook in (
+                *_global_forward_hooks.items(),
+                *self._forward_hooks.items(),
+            ):
+                # mark that always called hook is run
+                if hook_id in self._forward_hooks_always_called or hook_id in _global_forward_hooks_always_called:
+                    called_always_called_hooks.add(hook_id)
+
+                if hook_id in self._forward_hooks_with_kwargs:
+                    hook_result = hook(self, args, kwargs, result)
+                else:
+                    hook_result = hook(self, args, result)
+
+                if hook_result is not None:
+                    result = hook_result
+
+        if bw_hook:
+            if not isinstance(result, (ms.Tensor, tuple)):
+                warnings.warn("For backward hooks to be called,"
+                              " module output should be a Tensor or a tuple of Tensors"
+                              f" but received {type(result)}")
+            result = apply_backward_hook_on_tensors(bw_hook, result)
+
+        if self._backward_pre_hooks:
+            bw_pre_hook = inner.CellBackwardHook(self.__class__.__name__ + "(" + str(id(self)) + ")",
+                                                 self, self._backward_pre_hooks)
+            bw_pre_hook.register_backward_pre_hook()
+            result = apply_backward_hook_on_tensors(bw_pre_hook, result)
+
+        return result
+    except Exception:
+        # run always called hooks if they have not already been run
+        # For now only forward hooks have the always_call option but perhaps
+        # this functionality should be added to full backward hooks as well.
+        for hook_id, hook in _global_forward_hooks.items():
+            # type: ignore[possibly-undefined]
+            if hook_id in _global_forward_hooks_always_called and hook_id not in called_always_called_hooks:
+                try:
+                    hook_result = hook(self, args, result)  # type: ignore[possibly-undefined]
+                    if hook_result is not None:
+                        result = hook_result
+                except Exception as e:
+                    warnings.warn("global module forward hook with ``always_call=True`` raised an exception "
+                                  f"that was silenced as another error was raised in forward: {str(e)}")
+                    continue
+
+        for hook_id, hook in self._forward_hooks.items():
+            # type: ignore[possibly-undefined]
+            if hook_id in self._forward_hooks_always_called and hook_id not in called_always_called_hooks:
+                try:
+                    if hook_id in self._forward_hooks_with_kwargs:
+                        hook_result = hook(self, args, kwargs, result)  # type: ignore[possibly-undefined]
+                    else:
+                        hook_result = hook(self, args, result)  # type: ignore[possibly-undefined]
+                    if hook_result is not None:
+                        result = hook_result
+                except Exception as e:
+                    warnings.warn("module forward hook with ``always_call=True`` raised an exception "
+                                  f"that was silenced as another error was raised in forward: {str(e)}")
+                    continue
+        # raise exception raised in try block
+        raise
+
+
+def register_full_backward_pre_hook(self, hook, prepend: bool = False) -> RemovableHandle:
+    handle = RemovableHandle(self._backward_pre_hooks)
+    self._backward_pre_hooks[handle.id] = hook
+    if prepend:
+        self._backward_pre_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]
+    return handle
+
+
+def register_full_backward_hook(self, hook, prepend: bool = False) -> RemovableHandle:
+    if self._is_full_backward_hook is False:
+        raise RuntimeError(
+            "Cannot use both regular backward hooks and full backward hooks on a "
+            "single Module. Please use only one of them."
+        )
+
+    self._is_full_backward_hook = True
+
+    handle = RemovableHandle(self._backward_hooks)
+    self._backward_hooks[handle.id] = hook
+    if prepend:
+        self._backward_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]
+    return handle
+
+
+def _get_backward_pre_hooks(self):
+    self._backward_pre_hooks.update(_global_backward_pre_hooks)
+
+
+def _get_backward_hooks(self):
+    if (_global_is_full_backward_hook is True):
+        self._backward_hooks.update(_global_backward_hooks)
+
+
+def apply_backward_hook_on_tensors(cell_backward_hook, args):
+    is_tuple = True
+    if not isinstance(args, tuple):
+        args = (args,)
+        is_tuple = False
+    hooked_args = cell_backward_hook(*args)
+    if is_tuple and len(args) == 1:
+        hooked_args = (hooked_args, )
+    return hooked_args

msprobe/mindspore/ms_config.py

@@ -45,7 +45,11 @@ class StatisticsConfig(BaseConfig):
         self._check_config()
 
     def _check_config(self):
-        if self.summary_mode and self.summary_mode not in ["statistics", "md5"]:
+        single_opt = ["statistics", "md5"]
+        muti_opt = ["md5", "max", "min", "mean", "l2norm"]
+        if isinstance(self.summary_mode, str) and self.summary_mode not in single_opt:
+            raise Exception("summary_mode is invalid")
+        if isinstance(self.summary_mode, list) and not all(opt in muti_opt for opt in self.summary_mode):
             raise Exception("summary_mode is invalid")