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

msprobe/mindspore/dump/dump_tool_factory.py

@@ -15,15 +15,17 @@
 
 from msprobe.mindspore.common.const import Const
 from msprobe.core.common.log import logger
+from msprobe.mindspore.common.utils import is_graph_mode_cell_dump_allowed
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
 from msprobe.mindspore.dump.kernel_graph_dump import KernelGraphDump
 from msprobe.mindspore.dump.kernel_kbyk_dump import KernelKbykDump
+from msprobe.mindspore.dump.graph_mode_cell_dump import GraphModeCellDump
 
 
 class DumpToolFactory:
     tools = {
         Const.CELL: {
-            Const.GRAPH_KBYK_MODE: None,
+            Const.GRAPH_KBYK_MODE: GraphModeCellDump,
             Const.GRAPH_GE_MODE: None,
             Const.PYNATIVE_MODE: None
         },
@@ -40,9 +42,15 @@ class DumpToolFactory:
     }
 
     @staticmethod
-    def create(config: DebuggerConfig):
-        if len(config.data_mode) != 1 or config.data_mode[0] not in Const.GRAPH_DATA_MODE_LIST:
-            raise Exception("data_mode must be one of all, input, output.")
+    def create(config: DebuggerConfig, model=None):
+        if config.level == Const.CELL:
+            if not is_graph_mode_cell_dump_allowed(config):
+                raise Exception("Cell dump is not supported in graph mode.")
+            if len(config.data_mode) != 1 or config.data_mode[0] not in Const.GRAPH_CELL_DUMP_DATA_MODE_LIST:
+                raise Exception("data_mode must be one of all, forward, backward.")
+        else:
+            if len(config.data_mode) != 1 or config.data_mode[0] not in Const.GRAPH_DATA_MODE_LIST:
+                raise Exception("data_mode must be one of all, input, output.")
         tool = DumpToolFactory.tools.get(config.level)
         if not tool:
             raise Exception("Valid level is needed.")
@@ -51,4 +59,4 @@ class DumpToolFactory:
             logger.error(f"Data dump is not supported in {config.execution_mode} mode "
                          f"when dump level is {config.level}.")
             raise ValueError
-        return tool(config)
+        return tool(config, model) if tool == GraphModeCellDump else tool(config)

msprobe/mindspore/dump/graph_mode_cell_dump.py

@@ -0,0 +1,139 @@
+# 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 os
+
+import mindspore as ms
+from mindspore import hal, ops, Tensor
+from mindspore.ops.primitive import _run_op
+
+from msprobe.core.common.const import Const as CoreConst
+from msprobe.core.common.runtime import Runtime
+from msprobe.mindspore.common.const import Const
+from msprobe.mindspore.common.log import logger
+from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
+import msprobe.mindspore.dump.cell_dump_process as cellDumperWithDumpGradient
+import msprobe.mindspore.dump.cell_dump_with_insert_gradient as cellDumperWithInsertGradient
+
+tensordump_flag = True
+try:
+    from mindspore._c_expression import _tensordump_set_step
+except ImportError:
+    tensordump_flag = False
+
+
+class GraphModeCellDump:
+    task = CoreConst.STATISTICS
+
+    def __init__(self, config: DebuggerConfig, model, strict=True):
+        self.net = model
+        self.white_list = []
+        self.black_list = []
+        self.execution_mode = config.execution_mode
+        self.dump_path = config.dump_path if config.dump_path else "./"
+        self.rank = config.rank
+        self.step = config.step
+        self.scope = config.scope
+        self.list = config.list
+        self.data_mode = config.data_mode
+        self.file_format = config.file_format
+        GraphModeCellDump.task = config.task
+        self.summary_mode = config.summary_mode
+        self.check_config(strict)
+        self.set_step()
+
+    @staticmethod
+    def step():
+        # 更新TensorDump Step
+        if GraphModeCellDump.task == CoreConst.TENSOR:
+            hal.synchronize()
+            temp_tensor = ms.Tensor([1], dtype=ms.float32)
+            step_flag = "<tensordump-update-step>"
+            _run_op(ops.TensorDump(), "TensorDump", (step_flag, temp_tensor))
+            ops.tensordump(step_flag, temp_tensor)
+
+    def check_config(self, strict):
+        if not self.net:
+            raise Exception("The model is empty and cell dump is not enabled.")
+
+        if strict:
+            if self.rank:
+                raise Exception("In graph mode, cell dump does not currently support specifying rank.")
+            if self.scope:
+                raise Exception("In graph mode, cell dump does not currently support specifying scope.")
+            if self.list:
+                raise Exception("In graph mode, cell dump does not currently support specifying list.")
+            if len(self.data_mode) != 1 or self.data_mode[0] not in Const.GRAPH_CELL_DUMP_DATA_MODE_LIST:
+                raise Exception("In graph mode and cell dump, data_mode must be one of all, forword, backword.")
+            if self.file_format != []:
+                logger.warning("In graph mode, cell dump does not currently support specifying file_format."
+                               " The file will be stored in npy format.")
+            if self.task == CoreConst.STATISTICS and self.summary_mode == CoreConst.MD5:
+                raise Exception("The L0 level statistics dump mode does not support "
+                            "the calculation of md5 values currently In graph mode.")
+        else:
+            self.rank = []
+            self.scope = []
+            self.list = []
+            self.file_format = []
+            if len(self.data_mode) != 1 or self.data_mode[0] not in Const.GRAPH_CELL_DUMP_DATA_MODE_LIST:
+                self.data_mode = [CoreConst.ALL]
+            if self.task == CoreConst.STATISTICS and self.summary_mode == CoreConst.MD5:
+                self.summary_mode = CoreConst.STATISTICS
+
+        return True
+
+    def set_step(self):
+        if tensordump_flag:
+            _tensordump_set_step(self.step)
+        else:
+            raise Exception(
+                "Importing _tensordump_set_step failed, "
+                "please use the latest version package of MindSpore."
+            )
+
+    def handle(self):
+        os.environ['MS_JIT_MODULES'] = 'msprobe'
+
+        if Runtime.run_mode == Const.PYNATIVE_GRAPH_MODE:
+            dump_path = os.path.join(self.dump_path, Const.GRAPH_MODE)
+        else:
+            dump_path = self.dump_path
+
+        cell_dumper = cellDumperWithDumpGradient
+
+        if self.execution_mode == Const.PYNATIVE_MODE:
+            enable_dump_gradient = hasattr(ops, 'DumpGradient')
+            if hasattr(ops, 'DumpGradient'):
+                try:
+                    ops.DumpGradient()('grad.npy', Tensor([0], dtype=ms.float32), 'in')
+                except Exception:
+                    enable_dump_gradient = False
+                    logger.warning('the DumpGradient operator failed to execute.')
+            if not enable_dump_gradient:
+                cell_dumper = cellDumperWithInsertGradient
+        
+        dump_config = cell_dumper.CellDumpConfig(
+            net=self.net,
+            dump_path=dump_path,
+            data_mode=self.data_mode[0],
+            task=self.task,
+            summary_mode=self.summary_mode,
+            step=self.step
+        )
+
+        cell_dumper.start(
+            dump_config
+        )

msprobe/mindspore/dump/graph_tensor_dump.py

@@ -0,0 +1,123 @@
+# 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.
+# 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 collections import OrderedDict
+import mindspore as ms
+
+
+def _iterate_items(data):
+    if isinstance(data, (dict, OrderedDict)):
+        return data.items()
+    elif isinstance(data, (list, tuple)):
+        return enumerate(data)
+    else:
+        raise TypeError("Unsupported data type")
+
+
+class _SaveBase:
+    def __init__(self, save_dir):
+        super(_SaveBase, self).__init__()
+        self.path = save_dir
+        self.save_func = _npy_save
+
+    def get_save_func(self):
+        return self.save_func
+
+
+@ms.jit_class
+class _SaveCell(_SaveBase):
+    def __call__(self, name, data):
+        return self.get_save_func()(self.path, name, data)
+
+
+class _SaveGradBase:
+    def __init__(self, save_dir, name):
+        super(_SaveGradBase, self).__init__()
+        self.file = save_dir + name
+
+
+@ms.jit_class
+class _SaveGradCell(_SaveGradBase):
+    def __init__(self, save_dir, name):
+        super(_SaveGradCell, self).__init__(save_dir, name)
+        self.ms_save_grad = ms.ops.InsertGradientOf(
+            _wrapper_save_grad_func(self.file))
+
+    def __call__(self, x):
+        if isinstance(x, ms.Tensor):
+            return self.ms_save_grad(x)
+        else:
+            raise TypeError(f"For 'save_grad', the type of argument 'data' must be mindspore.Tensor or torch.tensor, "
+                            f"but got {type(x)}")
+
+
+def _npy_save_ops(file, data):
+    if isinstance(data, ms.Tensor):
+        if data.dtype == ms.bfloat16:
+            data = data.float()
+        ms.ops.TensorDump()(file, data)
+    else:
+        raise TypeError(f"For 'save', the type of argument 'data' must be mindspore.Tensor or torch.tensor, "
+                        f"but got {type(data)}")
+
+
+def _wrapper_save_grad_func(file):
+    def _save_grad_func(grad):
+        data = grad
+        if data.dtype == ms.bfloat16:
+            data = data.float()
+        ms.ops.TensorDump()(file, data)
+        return grad
+    return _save_grad_func
+
+
+def _npy_save(save_dir, item_name, data):
+    if isinstance(data, (list, tuple, dict, OrderedDict)):
+        for key, val in _iterate_items(data):
+            _npy_save(save_dir, f"{item_name}.{key}", val)
+    else:
+        if data is None:
+            return
+        _npy_save_ops(f"{save_dir}{item_name}", data)
+
+
+def generate_dump_dir(save_dir, sep=os.sep):
+    """
+    usage: generate dump directory path str in mindspore graph mode
+    """
+    full_suffix = '{step}' + sep + '{rank}' + sep
+    if save_dir and save_dir[-1] != sep:
+        result_dir = save_dir + sep + full_suffix
+    else:
+        result_dir = save_dir + full_suffix
+    return result_dir
+
+
+def save(save_dir, name, data):
+    """
+    save tensor.
+    """
+    dump_dir = generate_dump_dir(save_dir)
+    _SaveCell(dump_dir)(name, data)
+
+
+def save_grad(save_dir, name, data):
+    """
+    save grad.
+    """
+    dump_dir = generate_dump_dir(save_dir)
+    suffix_name = name + '_grad'
+    return _SaveGradCell(dump_dir, suffix_name)(data)

msprobe/mindspore/dump/hook_cell/api_register.py

@@ -14,14 +14,17 @@
 # limitations under the License.
 
 import os
+import inspect
 
 from mindspore import Tensor, ops, mint
+from mindspore.mint import distributed
 from mindspore.mint.nn import functional
 from mindspore.communication import comm_func
 
 from msprobe.core.common.file_utils import load_yaml
 from msprobe.core.common.utils import Const
 from msprobe.core.data_dump.api_registry import ApiRegistry
+from msprobe.mindspore.common.log import logger
 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
@@ -41,7 +44,8 @@ if not is_mindtorch():
             Const.MS_API_TYPE_TENSOR: (Tensor, (Tensor,)),
             Const.MS_API_TYPE_MINT: (mint, (mint,)),
             Const.MS_API_TYPE_MINT_FUNC: (functional, (functional,)),
-            Const.MS_API_TYPE_COM: (comm_func, (comm_func,))
+            Const.MS_API_TYPE_COM: (comm_func, (comm_func,)),
+            Const.MS_API_TYPE_MINT_DIST: (distributed, (distributed,))
         }
     }
     if stub_tensor_existed:
@@ -50,6 +54,7 @@ if not is_mindtorch():
         )
 
     _supported_api_list_path = (os.path.join(cur_path, MsConst.SUPPORTED_API_LIST_FILE),)
+    _backlist = []
 else:
     import torch
     import torch_npu
@@ -64,13 +69,14 @@ else:
     }
     _supported_api_list_path = (os.path.join(cur_path, '../../../pytorch/hook_module',
                                              MsConst.SUPPORTED_API_LIST_FILE),)
+    _backlist = [f'{Const.PT_API_TYPE_TENSOR}.__setitem__']
 
 _inner_used_api = {
     Const.MS_FRAMEWORK + Const.SEP + Const.MS_API_TYPE_OPS: (
         ops, "norm", "square", "sqrt", "is_complex", "stack", "is_floating_point"
     ),
     Const.MS_FRAMEWORK + Const.SEP + Const.MS_API_TYPE_TENSOR: (
-        Tensor, "to", "numel"
+        Tensor, "to", "numel", 'sum'
     ),
     Const.MS_FRAMEWORK + Const.SEP + Const.MS_API_TYPE_MINT: (
         mint, "max", "min", "mean", "norm"
@@ -84,6 +90,9 @@ class ApiTemplate(HOOKCell):
         self.api_func = api_func
         self.prefix_api_name = prefix + Const.SEP + str(api_name.split(Const.SEP)[-1]) + Const.SEP
         super().__init__(hook_build_func)
+        distributed_prefix = Const.DIST_API_TYPE_PREFIX if is_mindtorch() else Const.MINT_DIST_API_TYPE_PREFIX
+        if prefix == distributed_prefix:
+            self.op_is_distributed = True
 
     @staticmethod
     def async_to_sync(output):
@@ -103,9 +112,22 @@ class ApiTemplate(HOOKCell):
 
         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"]:
+        if self.prefix_api_name.startswith(
+            (MsConst.DISTRIBUTED_DATA_PREFIX, Const.MINT_DIST_API_TYPE_PREFIX)
+        ):
+            try:
+                bound = inspect.signature(self.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 self.api_name in ["isend", "irecv"]:
                 output = self.async_to_sync(output)
+            if self.api_name == "batch_isend_irecv" and isinstance(output, list):
+                output = [self.async_to_sync(handle) for handle in output]
+
         return output
 
     def forward(self, *args, **kwargs):
@@ -134,9 +156,21 @@ def get_api_register(return_new=False):
         stub_tensor_set = True
 
     if return_new:
-        return ApiRegistry(_api_types, _inner_used_api, _supported_api_list_path, ApiTemplate)
+        return ApiRegistry(
+            _api_types,
+            _inner_used_api,
+            _supported_api_list_path,
+            ApiTemplate,
+            _backlist
+        )
 
     global api_register
     if api_register is None:
-        api_register = ApiRegistry(_api_types, _inner_used_api, _supported_api_list_path, ApiTemplate)
+        api_register = ApiRegistry(
+            _api_types,
+            _inner_used_api,
+            _supported_api_list_path,
+            ApiTemplate,
+            _backlist
+        )
     return api_register

msprobe/mindspore/dump/hook_cell/hook_cell.py

@@ -15,11 +15,16 @@
 
 from collections import defaultdict
 
+import mindspore as ms
 from mindspore import nn
 
+from msprobe.core.common.runtime import Runtime
 from msprobe.mindspore.common.utils import is_mindtorch, register_backward_hook_functions
 
 
+ms_version = ms.__version__
+
+
 def add_cell_count(name):
     HOOKCell.cell_count[name] += 1
 
@@ -31,25 +36,31 @@ def get_cell_count(name):
 def __init__(self, hook_build_func) -> None:
     super(HOOKCell, self).__init__()
     self.changed_status = False
-    self.input_kwargs = {}
+    self.msprobe_input_kwargs = {}
     if not HOOKCell.g_stop_hook:
         HOOKCell.g_stop_hook = True
         self.changed_status = True
         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, backward_pre_hook = hook_build_func(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)
+            hook_set = hook_build_func(prefix)
+            if ms_version < "2.6.0" and not is_mindtorch():
+                getattr(self, "_forward_pre_hook", {})[id(self)] = hook_set.forward_pre_hook
+                getattr(self, "_forward_hook", {})[id(self)] = hook_set.forward_hook
+            else:
+                self.register_forward_pre_hook(hook_set.forward_pre_hook)
+                self.register_forward_hook(hook_set.forward_hook)
+            register_backward_hook_functions["full"](self, hook_set.backward_hook)
+            register_backward_hook_functions["pre"](self, hook_set.backward_pre_hook)
 
 
 # 重载call，加全局标志。
 def __call__(self, *args, **kwargs):
     try:
-        self.input_kwargs = kwargs
+        self.msprobe_input_kwargs = kwargs
         out = super(HOOKCell, self).__call__(*args, **kwargs)
     except Exception as e:
         raise e

msprobe/mindspore/dump/hook_cell/ms_hook_manager.py

@@ -0,0 +1,88 @@
+# 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 mindspore.common.api import _no_grad
+from msprobe.core.common.const import Const
+from msprobe.core.common.utils import replace_last_occurrence
+from msprobe.core.data_dump.data_processor.base import ModuleBackwardInputs
+from msprobe.core.hook_manager import BaseHookManager, HookSet
+from msprobe.mindspore.common.utils import has_kwargs_in_forward_hook
+from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
+
+
+class MindsproeHookManager(BaseHookManager):
+    @property
+    def _is_recompute(self):
+        return None
+
+    @staticmethod
+    def _no_grad_context():
+        return _no_grad()
+
+    @staticmethod
+    def _add_count(name):
+        HOOKCell.add_cell_count(name)
+
+    @staticmethod
+    def _process_kwargs_and_output(module, hook_type, kwargs_or_output, output_or_kwargs):
+        if not has_kwargs_in_forward_hook() or hook_type == Const.API:
+            kwargs = module.msprobe_input_kwargs if hasattr(module, 'msprobe_input_kwargs') else {}
+            output = kwargs_or_output
+        else:
+            kwargs = kwargs_or_output
+            output = output_or_kwargs
+        return kwargs, output
+
+    def build_hook(self, hook_type, name):
+        if hook_type == Const.API:
+            full_forward_name = name + str(HOOKCell.get_cell_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),
+            backward_pre_hook=self._build_backward_pre_hook(hook_type, full_backward_name)   
+        )
+        return hookset
+
+    def _need_exchange(self, module):
+        if not hasattr(module, 'has_pre_hook_called') or not module.has_pre_hook_called:
+            return False
+        else:
+            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.parameters_dict(recurse=False).items()
+                    }
+        return params_dict
+
+    def _build_backward_pre_hook(self, hook_type, name):
+        def backward_pre_hook(module, grad_input):
+            if self.config.level != Const.LEVEL_L2:
+                return
+            if not self._should_execute_hook(hook_type, module, False):
+                return
+            BaseHookManager.inner_switch = True
+            module_input = ModuleBackwardInputs(grad_input=grad_input)
+            self.data_collector.update_api_or_module_name(name)
+            self.data_collector.backward_input_data_collect(name, module, self._pid, module_input)
+            BaseHookManager.inner_switch = False
+        return backward_pre_hook

msprobe/mindspore/dump/hook_cell/primitive_hooks.py

@@ -21,6 +21,7 @@ from mindspore.common.tensor import Tensor
 from msprobe.core.common.utils import Const, DumpException
 from msprobe.core.data_dump.data_processor.base import (ModuleBackwardInputs, ModuleBackwardOutputs,
                                                         ModuleForwardInputsOutputs)
+from msprobe.core.hook_manager import BaseHookManager
 from msprobe.mindspore.common.log import logger
 
 
@@ -58,7 +59,7 @@ class PrimitiveHookService:
             def backward_hook(grad):
                 captured_grads.extend(grad)
                 backward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.BACKWARD}"
-
+                self.service_instance.inner_switch = True
                 try:
                     if hook_type == Const.INPUT:
                         self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
@@ -77,6 +78,7 @@ class PrimitiveHookService:
                     logger.error(f"This is a primitive op {hook_type}_backward dump error: {exception}, "
                                  f"updated_primitive_name: {updated_primitive_name}")
                     raise DumpException(DumpException.BACKWARD_DATA_COLLECTION_ERROR) from exception
+                self.service_instance.inner_switch = False
 
             return backward_hook
 
@@ -137,6 +139,7 @@ class PrimitiveHookService:
 
         def pre_forward_hook(primitive_name, primitive_instance, args, kwargs):
             module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
+            self.service_instance.inner_switch = True
             try:
                 self.service_instance.data_collector.forward_input_data_collect(
                     primitive_name,
@@ -148,9 +151,11 @@ class PrimitiveHookService:
                 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
+            self.service_instance.inner_switch = False
 
         def post_forward_hook(primitive_name, primitive_instance, args, kwargs, output):
             module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+            self.service_instance.inner_switch = True
             try:
                 self.service_instance.data_collector.forward_output_data_collect(
                     primitive_name,
@@ -162,6 +167,7 @@ class PrimitiveHookService:
                 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
+            self.service_instance.inner_switch = False
 
         def wrapped_primitive_call(instance_self, *args, **kwargs):
             """
@@ -179,7 +185,7 @@ class PrimitiveHookService:
             current_count = self.primitive_counters.get(primitive_name, 0)
             updated_primitive_name = f"{Const.PRIMITIVE_PREFIX}{Const.SEP}{primitive_name}{Const.SEP}{current_count}"
 
-            if not self.service_instance.primitive_switch:
+            if not self.service_instance.primitive_switch or BaseHookManager.inner_switch:
                 return origin_func(*args, **kwargs)
 
             captured_grads_input, captured_grads_output = [], []

msprobe/mindspore/dump/hook_cell/support_wrap_ops.yaml

@@ -1025,3 +1025,21 @@ communication.comm_func:
   - recv
   - isend
   - irecv
+
+mint.distributed:
+  - send
+  - recv
+  - broadcast
+  - all_reduce
+  - reduce
+  - all_gather
+  - gather
+  - isend
+  - irecv
+  - scatter
+  - reduce_scatter
+  - all_to_all_single
+  - all_to_all
+  - all_gather_into_tensor
+  - reduce_scatter_tensor
+  - batch_isend_irecv