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

msprobe/core/data_dump/data_collector.py

@@ -13,9 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import atexit
 import os
 
-from msprobe.core.data_dump.scope import build_scope, ListScope
+from msprobe.core.data_dump.scope import ScopeFactory
 from msprobe.core.data_dump.json_writer import DataWriter
 from msprobe.core.common.log import logger
 from msprobe.core.common.const import Const
@@ -27,7 +28,6 @@ def build_data_collector(config):
 
 
 class DataCollector:
-    multi_output_apis = ["_sort_", "npu_flash_attention"]
     tasks_need_tensor_data = [Const.OVERFLOW_CHECK, Const.TENSOR, Const.FREE_BENCHMARK]
     level_without_construct = [Const.LEVEL_L1, Const.LEVEL_L2]
 
@@ -37,13 +37,10 @@ class DataCollector:
         self.data_processor = DataProcessorFactory.create_processor(self.config, self.data_writer)
         self.module_processor = DataProcessorFactory.get_module_processor(self.config.framework)
         self.module_count = {}
-        if self.config.task == Const.FREE_BENCHMARK:
-            self.scope = build_scope(ListScope, self.config.scope, self.config.list)
-        else:
-            self.scope = build_scope(None, self.config.scope, self.config.list)
-
-    def __del__(self):
-        self.write_json()
+        self.scope = ScopeFactory(self.config).build_scope()
+        self.backward_module_names = {}
+        self.optimizer_status = ""
+        atexit.register(self.write_json)
 
     @property
     def dump_data_dir(self):
@@ -57,10 +54,6 @@ class DataCollector:
     def check_scope_and_pid(scope, name, pid):
         return (not scope or scope.check(name)) and pid == os.getpid()
 
-    @staticmethod
-    def is_inplace(module):
-        return getattr(module, "op_is_inplace", False)
-
     def if_return_forward_new_output(self):
         return self.data_processor.if_return_forward_new_output()
 
@@ -84,36 +77,54 @@ class DataCollector:
         logger.debug(msg)
         self.data_writer.update_data(data_info)
 
-    def pre_forward_data_collect(self, name, module, pid, module_input_output):
-        backward_name = name.replace(Const.FORWARD, Const.BACKWARD)
-        if self.check_scope_and_pid(self.scope, backward_name, pid):
-            self.data_processor.analyze_pre_forward(backward_name, module, module_input_output)
-        if not self.is_inplace(module) or not self.check_scope_and_pid(self.scope, name, pid):
+    def forward_input_data_collect(self, name, module, pid, module_input_output):
+        if self.config.task == Const.FREE_BENCHMARK:
+            backward_name = name.replace(Const.FORWARD, Const.BACKWARD)
+            if self.check_scope_and_pid(self.scope, backward_name, pid):
+                self.data_processor.analyze_forward_input(backward_name, module, module_input_output)
+            return
+
+        if not self.check_scope_and_pid(self.scope, name, pid):
+            return
+
+        data_info = self.data_processor.analyze_forward_input(name, module, module_input_output)
+        if self.config.level == Const.LEVEL_L2:
             return
-        logger.info(f"API {name} is inplace.")
-        data_info = self.data_processor.analyze_pre_forward_inplace(name, module_input_output)
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
-    def forward_data_collect(self, name, module, pid, module_input_output):
+    def forward_output_data_collect(self, name, module, pid, module_input_output):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
             return
 
-        if not self.is_inplace(module):
-            data_info = self.data_processor.analyze_forward(name, module, module_input_output)
-        else:
-            data_info = self.data_processor.analyze_forward_inplace(name, module_input_output)
-        if self.config.level == "L2":
+        data_info = self.data_processor.analyze_forward_output(name, module, module_input_output)
+        if self.config.level == Const.LEVEL_L2:
             return
         self.data_writer.update_stack(self.data_processor.analyze_api_call_stack(name))
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
+    def forward_data_collect(self, name, module, pid, module_input_output):
+        self.update_construct(name)
+        if not self.check_scope_and_pid(self.scope, name, pid):
+            return
+
+        data_info = self.data_processor.analyze_forward(name, module, module_input_output)
+        self.data_writer.update_stack(self.data_processor.analyze_api_call_stack(name))
+        self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
+
     def backward_data_collect(self, name, module, pid, module_input_output):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
             return
 
         data_info = self.data_processor.analyze_backward(name, module, module_input_output)
+        if self.config.level == Const.LEVEL_L2:
+            return
+        # 获取执行反向的模块名称
+        if data_info and name.split(Const.SEP)[0] in Const.MODULE_PREFIX:
+            module_name = name.rsplit(Const.SEP, 2)[0]
+            # 将模块名称加入到反向模块名称集合中，用于梯度收集时判断是否需要收集梯度
+            self.backward_module_names[module_name] = True
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
     def backward_input_data_collect(self, name, module, pid, module_input_output):
@@ -134,12 +145,17 @@ class DataCollector:
 
     def update_construct(self, name):
         if self.config.level not in DataCollector.level_without_construct:
-            self.data_writer.update_construct({name: self.module_processor.api_parent_node})
+            if self.optimizer_status in [Const.OPTIMIZER, Const.CLIP_GRAD]:
+                self.data_writer.update_construct({name: self.optimizer_status})
+            else:
+                self.data_writer.update_construct({name: self.module_processor.api_parent_node})
             self.data_writer.update_construct(self.module_processor.module_node)
 
     def handle_data(self, name, data_info, flush=False):
         if data_info:
             self.update_data(name, data_info)
+        if self.config.async_dump:
+            return
         if not flush:
             self.data_writer.flush_data_periodically()
         else:
@@ -147,7 +163,23 @@ class DataCollector:
 
     def update_dump_paths(self, *args):
         self.data_writer.update_dump_paths(*args)
-        self.data_writer.initialize_json_file(task=self.config.task, level=self.config.level)
+
+    def initialize_json_file(self, framework=Const.UNKNOWN_FRAMEWORK):
+        self.data_writer.initialize_json_file(task=self.config.task, level=self.config.level, framework=framework)
 
     def update_iter(self, current_iter):
         self.data_processor.update_iter(current_iter)
+
+    def params_data_collect(self, name, param_name, pid, data):
+        grad_name = name + Const.SEP + Const.PARAMS_GRAD
+        # 校验scope和pid，以及当前name是否有过反向计算
+        if not self.check_scope_and_pid(self.scope, name, pid) and not self.backward_module_names.get(name):
+            # 如果没有反向计算，则需要清除之前占位写入的grad数据
+            if self.data_writer.cache_data.get("data"):
+                self.data_writer.cache_data.get("data").pop(grad_name, None)
+            return
+        data_info = self.data_processor.analyze_params(grad_name, param_name, data)
+        self.handle_data(grad_name, data_info, flush=self.data_processor.is_terminated)
+
+    def fill_stack_tensor_data(self):
+        self.data_writer.fill_stack_tensor_data()

msprobe/core/data_dump/data_processor/base.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
-# Licensed under the Apache License, Version 2.0  (the "License");
+# 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
 #
@@ -15,10 +15,11 @@
 
 import inspect
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, is_dataclass
 from typing import Tuple, Dict, Optional, Any
 
 import numpy as np
+
 from msprobe.core.common.const import Const
 from msprobe.core.common.log import logger
 from msprobe.core.common.utils import convert_tuple, CompareException
@@ -38,9 +39,8 @@ class ModuleForwardInputsOutputs:
     def output_tuple(self):
         return convert_tuple(self.output)
 
-    def concat_args_and_kwargs(self):
-        args = self.args + tuple(self.kwargs.values())
-        return args
+    def update_output_with_args_and_kwargs(self):
+        self.output = self.args + tuple(self.kwargs.values())
 
 
 @dataclass
@@ -76,11 +76,12 @@ class ModuleBackwardOutputs:
 
 
 class TensorStatInfo:
-    def __init__(self, max_val=None, min_val=None, mean_val=None, norm_val=None):
+    def __init__(self, max_val=None, min_val=None, mean_val=None, norm_val=None, stack_tensor_stat=None):
         self.max = max_val
         self.min = min_val
         self.mean = mean_val
         self.norm = norm_val
+        self.stack_tensor_stat = stack_tensor_stat
 
 
 class BaseDataProcessor:
@@ -101,6 +102,9 @@ class BaseDataProcessor:
         self.current_iter = 0
         self._return_forward_new_output = False
         self._forward_new_output = None
+        self.save_name = None
+        if hasattr(config, "data_mode"):
+            self.allowed_data_mode = self._get_allowed_data_mode(config.data_mode)
 
     @property
     def data_path(self):
@@ -182,6 +186,18 @@ class BaseDataProcessor:
     def _analyze_numpy(value, numpy_type):
         return {"type": numpy_type, "value": value}
 
+    @staticmethod
+    def _get_allowed_data_mode(data_mode):
+        if Const.ALL in data_mode:
+            allowed_data_mode = [Const.FORWARD, Const.BACKWARD, Const.INPUT, Const.OUTPUT]
+        else:
+            allowed_data_mode = list(set(data_mode))
+            if Const.FORWARD not in allowed_data_mode and Const.BACKWARD not in allowed_data_mode:
+                allowed_data_mode += [Const.FORWARD, Const.BACKWARD]
+            if Const.INPUT not in allowed_data_mode and Const.OUTPUT not in allowed_data_mode:
+                allowed_data_mode += [Const.INPUT, Const.OUTPUT]
+        return allowed_data_mode
+
     @classmethod
     def get_special_types(cls):
         return cls.special_type
@@ -194,25 +210,42 @@ class BaseDataProcessor:
         if isinstance(args, cls.get_special_types()):
             arg_transform = transform(args, cls._recursive_key_stack)
             return arg_transform
+        elif isinstance(args, tuple) and hasattr(args, '_fields'):
+            # namedtuple to dict
+            args_dict = {field: getattr(args, field) for field in args._fields}
+            return cls.apply_transform_dict(args_dict, transform, depth)
+        elif is_dataclass(args):
+            # dataclass to dict
+            args_dict = {field: getattr(args, field) for field in args.__dataclass_fields__}
+            return cls.apply_transform_dict(args_dict, transform, depth)
         elif isinstance(args, (list, tuple)):
-            result_list = []
-            for i, arg in enumerate(args):
-                cls._recursive_key_stack.append(str(i))
-                result_list.append(cls.recursive_apply_transform(arg, transform, depth=depth + 1))
-                cls._recursive_key_stack.pop()
+            result_list = cls.apply_transform_list(args, transform, depth)
             return type(args)(result_list)
         elif isinstance(args, dict):
-            result_dict = {}
-            for k, arg in args.items():
-                cls._recursive_key_stack.append(str(k))
-                result_dict[k] = cls.recursive_apply_transform(arg, transform, depth=depth + 1)
-                cls._recursive_key_stack.pop()
-            return result_dict
+            return cls.apply_transform_dict(args, transform, depth)
         elif args is not None:
-            logger.warning(f"Data type {type(args)} is not supported.")
+            logger.debug(f"Data type {type(args)} is not supported.")
             return None
         else:
             return None
+    
+    @classmethod
+    def apply_transform_dict(cls, args, transform, depth):
+        result_dict = {}
+        for k, arg in args.items():
+            cls._recursive_key_stack.append(str(k))
+            result_dict[k] = cls.recursive_apply_transform(arg, transform, depth=depth + 1)
+            cls._recursive_key_stack.pop()
+        return result_dict
+
+    @classmethod
+    def apply_transform_list(cls, args, transform, depth):
+        result_list = []
+        for i, arg in enumerate(args):
+            cls._recursive_key_stack.append(str(i))
+            result_list.append(cls.recursive_apply_transform(arg, transform, depth=depth + 1))
+            cls._recursive_key_stack.pop()
+        return result_list
 
     def if_return_forward_new_output(self):
         return self._return_forward_new_output
@@ -239,17 +272,12 @@ class BaseDataProcessor:
         Return:
             bool: True if the parameters are in data_mode or data_mode is all, False otherwise.
         """
-        return (Const.ALL in self.config.data_mode or
-                forward_backward in self.config.data_mode or
-                input_output in self.config.data_mode)
-
-    def analyze_pre_forward(self, name, module, module_input_output: ModuleForwardInputsOutputs):
-        pass
+        return forward_backward in self.allowed_data_mode and input_output in self.allowed_data_mode
 
     def analyze_element(self, element):
         return self.recursive_apply_transform(element, self.analyze_single_element)
 
-    def analyze_forward(self, name, module, module_input_output: ModuleForwardInputsOutputs):
+    def analyze_forward_input(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         api_info_struct = {}
         # check whether data_mode contains forward or input
         if self.is_dump_for_data_mode(Const.FORWARD, Const.INPUT):
@@ -261,16 +289,22 @@ class BaseDataProcessor:
             kwargs_info_list = self.analyze_element(module_input_output.kwargs)
             api_info_struct[name][Const.INPUT_KWARGS] = kwargs_info_list
 
-        # check whether data_mode contains forward or output
+        return api_info_struct
+
+    def analyze_forward_output(self, name, module, module_input_output: ModuleForwardInputsOutputs):
+        api_info_struct = {}
+        # check whether data_mode contains forward or input
         if self.is_dump_for_data_mode(Const.FORWARD, Const.OUTPUT):
-            api_info_struct[name] = api_info_struct.get(name, {})
+            api_info_struct[name] = {}
             self.api_data_category = Const.OUTPUT
             output_info_list = self.analyze_element(module_input_output.output_tuple)
             api_info_struct[name][Const.OUTPUT] = output_info_list
+
         return api_info_struct
 
-    def analyze_pre_forward_inplace(self, name, module_input_output: ModuleForwardInputsOutputs):
+    def analyze_forward(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         api_info_struct = {}
+        # check whether data_mode contains forward or input
         if self.is_dump_for_data_mode(Const.FORWARD, Const.INPUT):
             api_info_struct[name] = {}
             self.api_data_category = Const.INPUT
@@ -279,16 +313,18 @@ class BaseDataProcessor:
             self.api_data_category = Const.KWARGS
             kwargs_info_list = self.analyze_element(module_input_output.kwargs)
             api_info_struct[name][Const.INPUT_KWARGS] = kwargs_info_list
-        return api_info_struct
 
-    def analyze_forward_inplace(self, name, module_input_output: ModuleForwardInputsOutputs):
-        concat_args = module_input_output.concat_args_and_kwargs()
-        api_info_struct = {}
+        # check whether data_mode contains forward or output
         if self.is_dump_for_data_mode(Const.FORWARD, Const.OUTPUT):
-            api_info_struct[name] = {}
+            api_info_struct[name] = api_info_struct.get(name, {})
             self.api_data_category = Const.OUTPUT
-            output_info_list = self.analyze_element(concat_args)
+            output_info_list = self.analyze_element(module_input_output.output_tuple)
             api_info_struct[name][Const.OUTPUT] = output_info_list
+
+        if name in api_info_struct and hasattr(module_input_output, Const.PARAMS):
+            self.api_data_category = Const.PARAMS
+            api_info_struct[name][Const.PARAMS] = self.analyze_element(getattr(module_input_output, Const.PARAMS))
+
         return api_info_struct
 
     def analyze_backward(self, name, module, module_input_output: ModuleBackwardInputsOutputs):
@@ -329,9 +365,21 @@ class BaseDataProcessor:
             api_info_struct[name][Const.OUTPUT] = output_info_list
         return api_info_struct
 
+    def analyze_params(self, name, param_name, grad):
+        api_info_struct = {}
+        self.save_name = name + Const.SEP + param_name
+        data_info = self.analyze_element(grad)
+        grad_info_dict = {param_name: [data_info]}
+        api_info_struct[name] = grad_info_dict
+        return api_info_struct
+
     def get_save_file_path(self, suffix):
         file_format = Const.PT_SUFFIX if self.config.framework == Const.PT_FRAMEWORK else Const.NUMPY_SUFFIX
-        dump_data_name = (self.current_api_or_module_name + Const.SEP + self.api_data_category + Const.SEP +
-                          suffix + file_format)
+        if self.save_name is not None:
+            dump_data_name = (self.save_name + file_format)
+            self.save_name = None
+        else:
+            dump_data_name = (self.current_api_or_module_name + Const.SEP + self.api_data_category + Const.SEP +
+                              suffix + file_format)
         file_path = os.path.join(self.data_writer.dump_tensor_data_dir, dump_data_name)
         return dump_data_name, file_path

msprobe/core/data_dump/data_processor/factory.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");
@@ -56,7 +56,7 @@ class DataProcessorFactory:
                 FreeBenchmarkDataProcessor as PytorchFreeBenchmarkDataProcessor,
                 KernelDumpDataProcessor as PytorchKernelDumpDataProcessor
             )
-            from msprobe.pytorch.module_processer import ModuleProcesser
+            from msprobe.pytorch.dump.module_dump.module_processer import ModuleProcesser
             cls.register_processor(Const.PT_FRAMEWORK, Const.STATISTICS, PytorchStatisticsDataProcessor)
             cls.register_processor(Const.PT_FRAMEWORK, Const.TENSOR, PytorchTensorDataProcessor)
             cls.register_processor(Const.PT_FRAMEWORK, Const.OVERFLOW_CHECK, PytorchOverflowCheckDataProcessor)
@@ -67,10 +67,12 @@ class DataProcessorFactory:
             from msprobe.core.data_dump.data_processor.mindspore_processor import (
                 StatisticsDataProcessor as MindsporeStatisticsDataProcessor,
                 TensorDataProcessor as MindsporeTensorDataProcessor,
-                OverflowCheckDataProcessor as MindsporeOverflowCheckDataProcessor
+                OverflowCheckDataProcessor as MindsporeOverflowCheckDataProcessor,
+                KernelDumpDataProcessor as MindsporeKernelDumpDataProcessor
             )
             from msprobe.mindspore.cell_processor import CellProcessor
             cls.register_processor(Const.MS_FRAMEWORK, Const.STATISTICS, MindsporeStatisticsDataProcessor)
             cls.register_processor(Const.MS_FRAMEWORK, Const.TENSOR, MindsporeTensorDataProcessor)
             cls.register_processor(Const.MS_FRAMEWORK, Const.OVERFLOW_CHECK, MindsporeOverflowCheckDataProcessor)
+            cls.register_processor(Const.MS_FRAMEWORK, Const.KERNEL_DUMP, MindsporeKernelDumpDataProcessor)
             cls.register_module_processor(Const.MS_FRAMEWORK, CellProcessor)

msprobe/core/data_dump/data_processor/mindspore_processor.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Huawei Technologies Co., Ltd
+# Copyright 2024-2025 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,7 +16,7 @@
 import zlib
 
 import mindspore as ms
-from mindspore import mint, ops
+from mindspore import mint, ops, hal
 from mindspore._c_expression.typing import Number
 import numpy as np
 
@@ -28,6 +28,12 @@ from msprobe.mindspore.common.utils import convert_bf16_to_fp32, save_tensor_as_
 from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.dump.hook_cell.api_registry import api_register
 
+has_adump = True
+try:
+    from msprobe.lib import _msprobe_c
+except ImportError:
+    has_adump = False
+
 
 class MindsporeDataProcessor(BaseDataProcessor):
     mindspore_special_type = tuple([ms.Tensor, Number])
@@ -37,6 +43,7 @@ class MindsporeDataProcessor(BaseDataProcessor):
         self.mindspore_object_key = {
             "dtype": self.analyze_dtype_in_kwargs
         }
+        self._async_dump_cache = {}
 
     @staticmethod
     def get_md5_for_tensor(x):
@@ -49,15 +56,10 @@ class MindsporeDataProcessor(BaseDataProcessor):
     def analyze_dtype_in_kwargs(element):
         return {"type": "mindspore.dtype", "value": str(element)}
 
-    @classmethod
-    def get_special_types(cls):
-        return super().get_special_types() + cls.mindspore_special_type
-
-    def get_stat_info(self, data):
+    @staticmethod
+    def get_stat_info_sync(data):
         tensor_stat = TensorStatInfo()
-        if data.numel() == 0:
-            return tensor_stat
-        elif data.dtype == ms.bool_:
+        if data.dtype == ms.bool_:
             data_np = data.asnumpy()
             tensor_stat.max = np.max(data_np).item()
             tensor_stat.min = np.min(data_np).item()
@@ -70,7 +72,7 @@ class MindsporeDataProcessor(BaseDataProcessor):
             tensor_stat.mean = np.mean(data_abs).item()
             tensor_stat.norm = np.linalg.norm(data_abs).item()
         else:
-            if not ops.is_floating_point(data):
+            if not ops.is_floating_point(data) or data.dtype == ms.float64:
                 data = data.to(ms.float32)
             api_register.norm_inner_op_set_ori_func()
             get_max_value = api_register.mint_ops_ori_attr.get("max", mint.max)
@@ -87,6 +89,47 @@ class MindsporeDataProcessor(BaseDataProcessor):
             api_register.norm_inner_op_set_hook_func()
         return tensor_stat
 
+    @staticmethod
+    def get_stat_info_async(data):
+        tensor_stat = TensorStatInfo()
+        stack_method = api_register.functional_ori_attr.get("stack", ms.ops.stack)
+        if data.dtype == ms.complex64 or data.dtype == ms.complex128:
+            logger.warning("Async dump do not support complex data!")
+            return tensor_stat
+        elif data.dtype == ms.bool_:
+            tensor_stat.stack_tensor_stat = (["Max", "Min"], stack_method([data.any(), data.all()]))
+        elif not data.shape:
+            tensor_stat.stack_tensor_stat = (["Max", "Min", "Mean", "Norm"], stack_method([data, data, data, data]))
+        else:
+            if not ops.is_floating_point(data) or data.dtype == ms.float64:
+                data = data.to(ms.float32)
+            api_register.norm_inner_op_set_ori_func()
+            get_max_value = api_register.mint_ops_ori_attr.get("max", mint.max)
+            get_min_value = api_register.mint_ops_ori_attr.get("min", mint.min)
+            get_mean_value = api_register.mint_ops_ori_attr.get("mean", mint.mean)
+            if hasattr(mint, "norm"):
+                get_norm_value = api_register.mint_ops_ori_attr.get("norm", mint.norm)
+            else:
+                get_norm_value = api_register.functional_ori_attr.get("norm", ops.norm)
+            tensor_stat.stack_tensor_stat = (["Max", "Min", "Mean", "Norm"], stack_method(
+                [get_max_value(data), get_min_value(data), get_mean_value(data), get_norm_value(data)]))
+            api_register.norm_inner_op_set_hook_func()
+        return tensor_stat
+
+    @classmethod
+    def get_special_types(cls):
+        return super().get_special_types() + cls.mindspore_special_type
+
+    def get_stat_info(self, data):
+        tensor_stat = TensorStatInfo()
+        if data.numel() == 0:
+            return tensor_stat
+        else:
+            if self.config.async_dump:
+                return MindsporeDataProcessor.get_stat_info_async(data)
+            else:
+                return MindsporeDataProcessor.get_stat_info_sync(data)
+
     def analyze_single_element(self, element, suffix_stack):
         if suffix_stack and suffix_stack[-1] in self.mindspore_object_key:
             return self.mindspore_object_key[suffix_stack[-1]](element)
@@ -107,13 +150,17 @@ class MindsporeDataProcessor(BaseDataProcessor):
         tensor_json = {
             'type': 'mindspore.Tensor',
             'dtype': str(tensor.dtype),
-            'shape': tensor.shape,
-            'Max': self.transfer_type(tensor_stat.max),
-            'Min': self.transfer_type(tensor_stat.min),
-            'Mean': self.transfer_type(tensor_stat.mean),
-            'Norm': self.transfer_type(tensor_stat.norm),
+            'shape': tensor.shape
         }
-        if self.config.summary_mode == Const.MD5:
+
+        if tensor_stat.stack_tensor_stat is None:
+            tensor_json.update({'Max': self.transfer_type(tensor_stat.max)})
+            tensor_json.update({'Min': self.transfer_type(tensor_stat.min)})
+            tensor_json.update({'Mean': self.transfer_type(tensor_stat.mean)})
+            tensor_json.update({'Norm': self.transfer_type(tensor_stat.norm)})
+        else:
+            tensor_json.update({'tensor_stat': tensor_stat.stack_tensor_stat})
+        if self.config.summary_mode == Const.MD5 and not self.config.async_dump:
             tensor_md5 = self.get_md5_for_tensor(tensor)
             tensor_json.update({Const.MD5: tensor_md5})
         return tensor_json
@@ -124,11 +171,19 @@ class StatisticsDataProcessor(MindsporeDataProcessor):
 
 
 class TensorDataProcessor(MindsporeDataProcessor):
+    def dump_async_data(self):
+        for file_path, tensor in self._async_dump_cache.items():
+            save_tensor_as_npy(tensor, file_path)
+        self._async_dump_cache.clear()
+
     def _analyze_tensor(self, tensor, suffix):
         dump_data_name, file_path = self.get_save_file_path(suffix)
         single_arg = super()._analyze_tensor(tensor, suffix)
         single_arg.update({"data_name": dump_data_name})
-        save_tensor_as_npy(tensor, file_path)
+        if self.config.async_dump:
+            self._async_dump_cache[file_path] = tensor.copy()
+        else:
+            save_tensor_as_npy(tensor, file_path)
         return single_arg
 
 
@@ -138,6 +193,7 @@ class OverflowCheckDataProcessor(MindsporeDataProcessor):
     def __init__(self, config, data_writer):
         super().__init__(config, data_writer)
         self.has_overflow = False
+        self.cached_api_info = {}
         self.cached_tensors_and_file_paths = {}
         self.real_overflow_nums = 0
         self.overflow_nums = config.overflow_nums
@@ -150,6 +206,20 @@ class OverflowCheckDataProcessor(MindsporeDataProcessor):
             return True
         return False
 
+    def analyze_forward_input(self, name, module, module_input_output: ModuleForwardInputsOutputs):
+        self.has_overflow = False
+        self.cached_api_info = super().analyze_forward_input(name, module, module_input_output)
+        return None
+
+    def analyze_forward_output(self, name, module, module_input_output: ModuleForwardInputsOutputs):
+        api_info_struct = super().analyze_forward_output(name, module, module_input_output)
+        if name in self.cached_api_info and name in api_info_struct:
+            self.cached_api_info[name].update(api_info_struct[name])
+        elif name in api_info_struct:
+            self.cached_api_info = api_info_struct
+        self.maybe_save_overflow_data()
+        return self.cached_api_info if self.has_overflow else None
+
     def analyze_forward(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         self.has_overflow = False
         api_info_struct = super().analyze_forward(name, module, module_input_output)
@@ -161,6 +231,12 @@ class OverflowCheckDataProcessor(MindsporeDataProcessor):
         api_info_struct = super().analyze_backward(name, module, module_input_output)
         self.maybe_save_overflow_data()
         return api_info_struct if self.has_overflow else None
+    
+    def analyze_params(self, name, param_name, grad):
+        self.has_overflow = False
+        api_info_struct = super().analyze_params(name, param_name, grad)
+        self.maybe_save_overflow_data()
+        return api_info_struct if self.has_overflow else None
 
     def maybe_save_overflow_data(self):
         if self.has_overflow:
@@ -190,3 +266,61 @@ class OverflowCheckDataProcessor(MindsporeDataProcessor):
         self._analyze_maybe_overflow_tensor(single_arg)
         single_arg.update({"data_name": dump_data_name})
         return single_arg
+
+
+class KernelDumpDataProcessor(MindsporeDataProcessor):
+    def __init__(self, config, data_writer):
+        super().__init__(config, data_writer)
+        self.enable_kernel_dump = True
+
+    @staticmethod
+    def start_kernel_dump(config_path):
+        hal.synchronize()
+        _msprobe_c.init_dump()
+        _msprobe_c.set_dump(config_path)
+        hal.synchronize()
+
+    @staticmethod
+    def stop_kernel_dump():
+        hal.synchronize()
+        _msprobe_c.finalize_dump()
+        hal.synchronize()
+
+    @staticmethod
+    def _print_unsupported_log(api_name):
+        logger.warning(f"The kernel dump does not support the {api_name} API.")
+
+    def analyze_forward_input(self, name, module, module_input_output):
+        if not self.enable_kernel_dump:
+            return
+        if not has_adump:
+            logger.warning("The current msprobe package does not compile adump, and kernel dump cannot be used.")
+            self.enable_kernel_dump = False
+            return
+        self.start_kernel_dump(self.config.kernel_config_path)
+
+    def analyze_forward_output(self, name, module, module_input_output):
+        if not self.enable_kernel_dump:
+            return
+        self.enable_kernel_dump = False
+        self.stop_kernel_dump()
+        logger.info(f"The kernel data of {name} is dumped successfully.")
+
+    def analyze_backward_input(self, name, module, module_input_output):
+        if not self.enable_kernel_dump:
+            return
+        if not has_adump:
+            logger.warning("The current msprobe package does not compile adump, and kernel dump cannot be used.")
+            self.enable_kernel_dump = False
+            return
+        self.start_kernel_dump(self.config.kernel_config_path)
+
+    def analyze_backward(self, name, module, module_input_output):
+        if not self.enable_kernel_dump:
+            return
+        self.enable_kernel_dump = False
+        self.stop_kernel_dump()
+        logger.info(f"The kernel data of {name} is dumped successfully.")
+
+    def reset_status(self):
+        self.enable_kernel_dump = True