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

msprobe/mindspore/dump/hook_cell/primitive_hooks.py

@@ -0,0 +1,203 @@
+# Copyright (c) 2024-2024, 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 mindspore import ops
+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.mindspore.common.log import logger
+
+
+class PrimitiveHookService:
+    def __init__(self, service_instance):
+        self.primitive_counters = {}
+        self.service_instance = service_instance
+
+    def wrap_primitive(self, origin_func, primitive_name):
+        """
+        包装原始的 primitive 函数，添加输入和输出的 hook 以捕获前向和反向数据。
+
+        Args:
+            origin_func (callable): 原始 的 primitive 函数。
+            primitive_name (str): 原始的 primitive 名称。
+
+        Returns:
+            callable: 包装后的 primitive 函数。
+        """
+
+        def create_backward_hook(captured_grads, num_tensors, updated_primitive_name, hook_type):
+            """
+            创建反向 hook 函数，用于捕获梯度。
+
+            Args:
+                captured_grads (list): 用于保存捕获的梯度。
+                num_tensors (int): 张量数量。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+                hook_type (str): hook 类型 (输入/输出)。
+
+            Returns:
+                callable: 反向 hook 函数。
+            """
+
+            def backward_hook(grad):
+                captured_grads.extend(grad)
+                backward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.BACKWARD}"
+
+                try:
+                    if hook_type == Const.INPUT:
+                        self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
+                        new_module_input_output = ModuleBackwardOutputs(grad_output=tuple(captured_grads))
+                        self.service_instance.data_collector.backward_output_data_collect(
+                            backward_primitive_name, self, os.getpid(), new_module_input_output
+                        )
+                    elif hook_type == Const.OUTPUT:
+                        self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
+                        new_module_input_output = ModuleBackwardInputs(grad_input=tuple(captured_grads))
+                        self.service_instance.data_collector.backward_input_data_collect(
+                            backward_primitive_name, self, os.getpid(), new_module_input_output
+                        )
+
+                except Exception as exception:
+                    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
+
+            return backward_hook
+
+        def hook_primitive_inputs(args, captured_grads_input, updated_primitive_name):
+            """
+            针对前向输入添加 hook。
+
+            Args:
+                args (tuple): primitive 输入参数。
+                captured_grads_input (list): 捕获的输入梯度。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+
+            Returns:
+                list: 添加了 hook 的输入。
+            """
+            hooked_inputs = []
+            num_tensors = sum(isinstance(arg, Tensor) for arg in args)
+            input_backward_hook = create_backward_hook(captured_grads_input, num_tensors, updated_primitive_name,
+                                                       Const.INPUT)
+            for arg in args:
+                if isinstance(arg, Tensor):
+                    arg_hooked = ops.HookBackward(input_backward_hook)(arg)
+                    hooked_inputs.append(arg_hooked)
+                else:
+                    hooked_inputs.append(arg)
+            return tuple(hooked_inputs)
+
+        def hook_primitive_outputs(out, captured_grads_output, updated_primitive_name):
+            """
+            针对前向输出添加 hook。
+
+            Args:
+                out (Tensor/tuple): primitive 输出。
+                captured_grads_output (list): 捕获的输出梯度。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+
+            Returns:
+                Tensor/tuple: 添加了 hook 的输出。
+            """
+            if isinstance(out, tuple):
+                num_output_tensors = sum(isinstance(tensor, Tensor) for tensor in out)
+            else:
+                num_output_tensors = 1
+            output_backward_hook = create_backward_hook(captured_grads_output, num_output_tensors,
+                                                        updated_primitive_name, Const.OUTPUT)
+
+            if isinstance(out, Tensor):
+                return ops.HookBackward(output_backward_hook)(out)
+            elif isinstance(out, tuple):
+                hooked_outputs = []
+                for tensor in out:
+                    if isinstance(tensor, Tensor):
+                        hooked_outputs.append(ops.HookBackward(output_backward_hook)(tensor))
+                    else:
+                        hooked_outputs.append(tensor)
+                return tuple(hooked_outputs)
+            return out
+
+        def wrapped_primitive_call(instance_self, *args, **kwargs):
+            """
+            包装后的 primitive 调用函数，添加输入和输出的 hook。
+
+            Args:
+                instance_self (object): primitive 的实例。
+                *args: primitive 输入参数。
+                **kwargs: primitive 关键字参数。
+
+            Returns:
+                Tensor/tuple: primitive 的返回值。
+            """
+            self.update_primitive_counters(primitive_name)
+            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:
+                return origin_func(*args, **kwargs)
+
+            captured_grads_input, captured_grads_output = [], []
+
+            try:
+                hooked_inputs = hook_primitive_inputs(args, captured_grads_input, updated_primitive_name)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during input hooking: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.INPUT_HOOK_ERROR) from exception
+
+            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()
+
+            try:
+                out = hook_primitive_outputs(out, captured_grads_output, updated_primitive_name)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during output hooking: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.OUTPUT_HOOK_ERROR) from exception
+
+            return out
+
+        return wrapped_primitive_call
+
+    def update_primitive_counters(self, primitive_name):
+        if primitive_name not in self.primitive_counters:
+            self.primitive_counters[primitive_name] = 0
+        else:
+            self.primitive_counters[primitive_name] += 1

msprobe/mindspore/dump/hook_cell/support_wrap_ops.yaml

@@ -185,6 +185,7 @@ ops:
   - float_power
   - fmod
   - frac
+  - flash_attention_score
   - gcd
   - hypot
   - igamma
@@ -489,6 +490,31 @@ ops:
   - scatter_update
   - derivative
   - jet
+  - row_stack
+  - gather
+  - arange
+  - cond
+  - slice_scatter
+  - clip_by_norm
+  - eps
+  - layer_norm
+  - cast
+  - numel
+  - permute
+  - select_scatter
+  - group_norm
+  - eq
+  - embedding
+  - ones_like
+  - zeros
+  - nanmean
+  - shape
+  - zeros_like
+  - ones
+  - diagonal_scatter
+  - vander
+  - is_nonzero
+  - rotary_position_embedding
 
 tensor:
   - __abs__
@@ -876,16 +902,60 @@ mint.ops:
   - zeros
   - zeros_ex
   - zeros_like
-
-mint.nn:
-  - Dropout
-  - Embedding
-  - Fold
-  - LayerNorm
-  - Linear
-  - MaxPool2d
-  - Unfold
-  - Upsample
+  - inverse
+  - select
+  - item
+  - unsqueeze
+  - median
+  - floor
+  - histc
+  - special
+  - arctan2
+  - sign
+  - concat
+  - atanh
+  - greater_equal
+  - eye
+  - fix
+  - argmin
+  - asinh
+  - atan
+  - nan_to_num
+  - tan
+  - round
+  - cosh
+  - norm
+  - roll
+  - log1p
+  - reshape
+  - arccos
+  - outer
+  - arcsin
+  - rand_like
+  - acosh
+  - multinomial
+  - logical_xor
+  - acos
+  - linalg
+  - sinc
+  - arcsinh
+  - asin
+  - narrow
+  - arctanh
+  - trace
+  - erfc
+  - bernoulli
+  - expm1
+  - logaddexp
+  - sinh
+  - arccosh
+  - atan2
+  - rand
+  - arange
+  - trunc
+  - arctan
+  - swapaxes
+  - transpose
 
 mint.nn.functional:
   - absolute_import
@@ -920,3 +990,30 @@ mint.nn.functional:
   - softplus
   - tanh
   - unfold
+  - mse_loss
+  - adaptive_avg_pool1d
+  - binary_cross_entropy
+  - adaptive_avg_pool2d
+  - hardsigmoid
+  - selu
+  - softshrink
+  - prelu
+  - logsigmoid
+  - hardswish
+  - mish
+  - log_softmax
+  - hardshrink
+  - l1_loss
+  - elu
+
+communication.comm_func:
+  - all_reduce
+  - all_gather_into_tensor
+  - reduce
+  - reduce_scatter_tensor
+  - all_to_all_single_with_output_shape
+  - all_to_all_with_output_shape
+  - batch_isend_irecv
+  - broadcast
+  - gather_into_tensor
+  - scatter_tensor

msprobe/mindspore/dump/hook_cell/wrap_api.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2024-2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, 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
 #
@@ -13,19 +12,18 @@
 # 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 mindspore import Tensor, ops, mint
-from mindspore.mint.nn import functional
+from mindspore import Tensor, mint, ops
 from mindspore.common._stub_tensor import StubTensor
+from mindspore.communication import comm_func
+from mindspore.mint.nn import functional
 
-from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 from msprobe.core.common.const import Const
-from msprobe.mindspore.common.const import Const as MsConst
 from msprobe.core.common.file_utils import load_yaml
-
+from msprobe.mindspore.common.const import Const as MsConst
+from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, MsConst.SUPPORTED_API_LIST_FILE)
@@ -51,6 +49,10 @@ class HOOKMintNNFunctionalOP(object):
     pass
 
 
+class HOOKDistributedOP(object):
+    pass
+
+
 class ApiTemplate(HOOKCell):
     def __init__(self, api_name, api_dict, prefix, hook):
         self.api_name = api_name
@@ -65,12 +67,14 @@ class ApiTemplate(HOOKCell):
 
 
 class WrapApiName:
-    def __init__(self, tensor_api_names, stub_tensor_api_names, ops_api_names, mint_api_names, mint_nn_func_api_names):
+    def __init__(self, tensor_api_names, stub_tensor_api_names, ops_api_names, mint_api_names, mint_nn_func_api_names,
+                 distributed_api_names):
         self.tensor_api_names = tensor_api_names
         self.stub_tensor_api_names = stub_tensor_api_names
         self.ops_api_names = ops_api_names
         self.mint_api_names = mint_api_names
         self.mint_nn_func_api_names = mint_nn_func_api_names
+        self.distributed_api_names = distributed_api_names
 
 
 def get_wrap_api_list():
@@ -79,11 +83,13 @@ def get_wrap_api_list():
     ops_api = api_list.get(MsConst.SUPPORTED_OPS_LIST_KEY)
     mint_api = api_list.get(MsConst.SUPPORTED_MINT_LIST_KEY)
     mint_nn_func_api = api_list.get(MsConst.SUPPORTED__MINT_NN_FUNC_LIST_KEY)
+    distributed_api = api_list.get(MsConst.SUPPORTED_COMM_LIST_KEY)
     wrap_api_name = WrapApiName(set(tensor_api) & set(dir(Tensor)),
                                 set(tensor_api) & set(dir(StubTensor)),
                                 set(ops_api) & set(dir(ops)),
                                 set(mint_api) & set(dir(mint)),
-                                set(mint_nn_func_api) & set(dir(functional)))
+                                set(mint_nn_func_api) & set(dir(functional)),
+                                set(distributed_api) & set(dir(comm_func)))
     return wrap_api_name
 
 
@@ -111,3 +117,5 @@ def setup_hooks(hook):
                            MsConst.MINT_DATA_PREFIX, hook, HOOKMintOP)
     wrap_api_func_and_bind(wrap_api_name.mint_nn_func_api_names, {f: getattr(functional, f) for f in dir(functional)},
                            MsConst.MINT_NN_FUNC_DATA_PREFIX, hook, HOOKMintNNFunctionalOP)
+    wrap_api_func_and_bind(wrap_api_name.distributed_api_names, {f: getattr(comm_func, f) for f in dir(comm_func)},
+                           MsConst.DISTRIBUTED_DATA_PREFIX, hook, HOOKDistributedOP)

msprobe/mindspore/dump/jit_dump.py

@@ -1,12 +1,30 @@
+# Copyright (c) 2024-2024, 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 defaultdict
 
 from mindspore import Tensor
-from mindspore.common.api import _MindsporeFunctionExecutor
 from mindspore._c_expression import PyNativeExecutor_
+from mindspore.common.api import _MindsporeFunctionExecutor
 
-from msprobe.mindspore.dump.hook_cell.api_registry import api_register
-from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs
+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.mindspore.dump.hook_cell.api_registry import api_register
 
 
 def dump_jit(name, in_feat, out_feat, is_forward):
@@ -15,35 +33,53 @@ def dump_jit(name, in_feat, out_feat, is_forward):
     index = ori_args.find("<")
     if index != 0 and index != -1:
         result = ori_args[0:index]
+    elif name is not None and "<" not in str(name):
+        result = str(name)
     else:
         result = "JitFunction"
-    if is_forward:
-        name_template = "Jit." + result + ".forward"
-    else:
-        name_template = "Jit." + result + ".backward"
     if JitDump.need_dump():
-        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, {}, pid, module_input_output)
+        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
+            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)
+        else:
+            name_template = Const.JIT + Const.SEP + result + Const.SEP + str(JitDump.jit_count[result]) + Const.SEP + \
+                            Const.BACKWARD
+            JitDump.data_collector.update_api_or_module_name(name_template)
+            module_input_output = ModuleBackwardInputsOutputs(grad_input=in_feat, grad_output=out_feat)
+            JitDump.data_collector.backward_data_collect(name_template, None, pid, module_input_output)
 
 
 class JitDump(_MindsporeFunctionExecutor):
     dump_config = None
     jit_enable = False
+    jit_dump_switch = True
+    jit_count = defaultdict(int)
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.name = None
+        if len(args) > 0:
+            self.name = args[0].__name__
         self._executor = PyNativeExecutor_.get_instance()
 
     def __call__(self, *args, **kwargs):
-        api_register.api_set_ori_func()
+        if JitDump.jit_dump_switch:
+            api_register.api_set_ori_func()
         out = super().__call__(*args, **kwargs)
-        if isinstance(args[0], Tensor):
-            dump_jit({}, args, out, True)
-        else:
-            dump_jit(args[0], args[1:], out, True)
-        JitDump.jit_enable = True
-        api_register.api_set_hook_func()
+        if JitDump.jit_dump_switch and len(args) > 0:
+            if self.name and self.name != "construct":
+                dump_jit(self.name, args, out, True)
+            else:
+                dump_jit(args[0], args, out, True)
+            JitDump.jit_enable = True
+        elif len(args) == 0:
+            logger.warning(f"The jit function {self.name} has no input arguments, nothing will be dumped.")
+        if JitDump.jit_dump_switch:
+            api_register.api_set_hook_func()
         return out
 
     @classmethod
@@ -62,11 +98,11 @@ class JitDump(_MindsporeFunctionExecutor):
             return False
         return True
 
-    def grad(self, obj, grad, weights, grad_position, *args,  **kwargs):
-        if JitDump.jit_enable:
+    def grad(self, obj, grad, weights, grad_position, *args, **kwargs):
+        if JitDump.jit_dump_switch and JitDump.jit_enable:
             api_register.api_set_ori_func()
         output = self._executor.grad(grad, obj, weights, grad_position, *args, *(kwargs.values()))
-        if JitDump.jit_enable:
+        if JitDump.jit_dump_switch and JitDump.jit_enable:
             dump_jit(obj, args, None, False)
             api_register.api_set_hook_func()
         return output

msprobe/mindspore/dump/kernel_graph_dump.py

@@ -1,8 +1,23 @@
+# Copyright (c) 2024-2024, 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 json
-from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
+
+from msprobe.core.common.file_utils import create_directory, save_json
 from msprobe.mindspore.common.log import logger
-from msprobe.core.common.file_utils import FileOpen, create_directory
+from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
 
 
 class KernelGraphDump:
@@ -46,8 +61,7 @@ class KernelGraphDump:
         json_path = self.dump_json["common_dump_settings"]["path"]
         create_directory(json_path)
         json_path = os.path.join(json_path, "kernel_graph_dump.json")
-        with FileOpen(json_path, 'w') as f:
-            json.dump(self.dump_json, f)
+        save_json(json_path, self.dump_json, indent=4)
         logger.info(json_path + " has been created.")
         os.environ["MINDSPORE_DUMP_CONFIG"] = json_path
         if self.dump_json["common_dump_settings"]["dump_mode"] == 0:

msprobe/mindspore/dump/kernel_kbyk_dump.py

@@ -1,10 +1,24 @@
+# Copyright (c) 2024-2024, 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 json
 
-from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
-from msprobe.mindspore.common.log import logger
-from msprobe.core.common.file_utils import FileOpen, create_directory
 from msprobe.core.common.const import Const
+from msprobe.core.common.file_utils import create_directory, save_json
+from msprobe.mindspore.common.log import logger
+from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
 
 
 class KernelKbykDump:
@@ -55,8 +69,7 @@ class KernelKbykDump:
         json_path = self.dump_json[KernelKbykDump.COMMON_SETTINGS]["path"]
         create_directory(json_path)
         json_path = os.path.join(json_path, "kernel_kbyk_dump.json")
-        with FileOpen(json_path, 'w') as f:
-            json.dump(self.dump_json, f)
+        save_json(json_path, self.dump_json, indent=4)
         logger.info(json_path + " has been created.")
 
         os.environ["MINDSPORE_DUMP_CONFIG"] = json_path