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

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,11 +43,12 @@ class MindsporeDataProcessor(BaseDataProcessor):
         self.mindspore_object_key = {
             "dtype": self.analyze_dtype_in_kwargs
         }
+        self._async_dump_cache = {}
 
     @staticmethod
     def get_md5_for_tensor(x):
         x = convert_bf16_to_fp32(x)
-        tensor_bytes = x.contiguous().asnumpy().tobytes()
+        tensor_bytes = x.asnumpy().tobytes()
         crc32_hash = zlib.crc32(tensor_bytes)
         return f"{crc32_hash:08x}"
 
@@ -49,22 +56,17 @@ 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_:
-            data_np = data.contiguous().asnumpy()
+        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()
         elif not data.shape:
             tensor_stat.max = tensor_stat.min = tensor_stat.mean = tensor_stat.norm = data.item()
         elif data.dtype == ms.complex64 or data.dtype == ms.complex128:
-            data_abs = np.abs(data.contiguous().asnumpy())
+            data_abs = np.abs(data.asnumpy())
             tensor_stat.max = np.max(data_abs).item()
             tensor_stat.min = np.min(data_abs).item()
             tensor_stat.mean = np.mean(data_abs).item()
@@ -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

msprobe/core/data_dump/data_processor/pytorch_processor.py

@@ -54,6 +54,7 @@ class PytorchDataProcessor(BaseDataProcessor):
             "device": self.analyze_device_in_kwargs,
             "dtype": self.analyze_dtype_in_kwargs
         }
+        self._async_dump_cache = {}
 
     @staticmethod
     def get_md5_for_tensor(x):
@@ -82,49 +83,80 @@ class PytorchDataProcessor(BaseDataProcessor):
         return {"type": "torch.dtype", "value": str(element)}
 
     @staticmethod
-    def get_stat_info(data):
+    def get_stat_info_async(data):
         tensor_stat = TensorStatInfo()
-        if data.is_meta:
-            return tensor_stat
-        data_clone = data.detach()
-        if data_clone.numel() == 0:
+        if torch.is_complex(data):
+            logger.warning("Async dump do not support complex data!")
             return tensor_stat
-        elif data_clone.dtype == torch.bool:
-            tensor_stat.max = torch._C._VariableFunctionsClass.any(data_clone).item()
-            tensor_stat.min = torch._C._VariableFunctionsClass.all(data_clone).item()
-        elif not data_clone.shape:
-            tensor_stat.max = tensor_stat.min = tensor_stat.mean = tensor_stat.norm = data_clone.item()
-        elif torch.is_complex(data_clone):
-            data_np = data_clone.cpu().numpy()
+        elif data.dtype == torch.bool:
+            tensor_stat.stack_tensor_stat = (["Max", "Min"], torch.stack(
+                [torch.any(data), torch.all(data)]))
+        elif not data.shape:
+            tensor_stat.stack_tensor_stat = (["Max", "Min", "Mean", "Norm"], torch.stack([data, data, data, data]))
+        else:
+            if not data.is_floating_point() or data.dtype == torch.float64:
+                data = data.float()
+            tensor_stat.stack_tensor_stat = (["Max", "Min", "Mean", "Norm"], torch.stack([
+                torch.max(data),
+                torch.min(data),
+                torch.mean(data),
+                torch.norm(data)
+            ]))
+        return tensor_stat
+
+    @staticmethod
+    def get_stat_info_sync(data):
+        tensor_stat = TensorStatInfo()
+        if torch.is_complex(data):
+            data_np = data.cpu().numpy()
             data_abs = np.abs(data_np)
             tensor_stat.max = np.max(data_abs).item()
             tensor_stat.min = np.min(data_abs).item()
             tensor_stat.mean = np.mean(data_abs).item()
+        elif data.dtype == torch.bool:
+            tensor_stat.max = torch.any(data).item()
+            tensor_stat.min = torch.all(data).item()
+        elif not data.shape:
+            tensor_stat.max = tensor_stat.min = tensor_stat.mean = tensor_stat.norm = data.item()
         else:
-            if not data_clone.is_floating_point() or data_clone.dtype == torch.float64:
-                data_clone = data_clone.float()
-            tensor_stat.max = torch._C._VariableFunctionsClass.max(data_clone).item()
-            tensor_stat.min = torch._C._VariableFunctionsClass.min(data_clone).item()
-            tensor_stat.mean = torch._C._VariableFunctionsClass.mean(data_clone).item()
-            tensor_stat.norm = torch._C._VariableFunctionsClass.norm(data_clone).item()
+            if not data.is_floating_point() or data.dtype == torch.float64:
+                data = data.float()
+            tensor_stat.max = torch.max(data).item()
+            tensor_stat.min = torch.min(data).item()
+            tensor_stat.mean = torch.mean(data).item()
+            tensor_stat.norm = torch.norm(data).item()
         return tensor_stat
 
+    @staticmethod
+    def get_stat_info(data, async_dump=False):
+        tensor_stat = TensorStatInfo()
+        if data.is_meta:
+            return tensor_stat
+        data_clone = data.detach()
+        if data_clone.numel() == 0:
+            return tensor_stat
+        else:
+            if data_clone.device.type == Const.CPU_LOWERCASE or not async_dump:
+                return PytorchDataProcessor.get_stat_info_sync(data_clone)
+            else:
+                return PytorchDataProcessor.get_stat_info_async(data_clone)
+
     @staticmethod
     def handle_tensor_extremum_nan_inf(tensor, operator):
         data_clone = tensor.detach()
-        data_nan = torch._C._VariableFunctionsClass.isnan(data_clone)
-        if int(torch._C._VariableFunctionsClass.sum(data_nan)) == data_clone.numel():
+        data_nan = torch.isnan(data_clone)
+        if int(torch.sum(data_nan)) == data_clone.numel():
             return float('nan')
 
-        finite_mask = torch._C._VariableFunctionsClass.isfinite(data_clone)
-        if int(torch._C._VariableFunctionsClass.sum(finite_mask)) > 0:
-            finite_values = getattr(torch._C._TensorBase, "__getitem__")(data_clone, finite_mask)
-            return torch._C._VariableFunctionsClass.max(finite_values).item() if operator == 'max' else \
-                torch._C._VariableFunctionsClass.min(finite_values).item()
+        finite_mask = torch.isfinite(data_clone)
+        if int(torch.sum(finite_mask)) > 0:
+            finite_values = data_clone[finite_mask]
+            return torch.max(finite_values).item() if operator == 'max' else \
+                torch.min(finite_values).item()
         else:
-            data_no_nan = getattr(torch._C._TensorBase, "__getitem__")(data_clone, ~data_nan)
-            return torch._C._VariableFunctionsClass.max(data_no_nan).item() if operator == 'max' else \
-                torch._C._VariableFunctionsClass.min(data_no_nan).item()
+            data_no_nan = data_clone[~data_nan]
+            return torch.max(data_no_nan).item() if operator == 'max' else \
+                torch.min(data_no_nan).item()
 
     @staticmethod
     def process_group_hash(arg):
@@ -132,6 +164,10 @@ class PytorchDataProcessor(BaseDataProcessor):
         group_ranks_hash = hashlib.md5(str(group_ranks).encode('utf-8')).hexdigest()
         return group_ranks_hash
 
+    @staticmethod
+    def is_distributed_op(module):
+        return getattr(module, "op_is_distributed", False)
+
     @staticmethod
     def _analyze_torch_size(arg):
         return {"type": "torch.Size", "value": list(arg)}
@@ -177,26 +213,35 @@ class PytorchDataProcessor(BaseDataProcessor):
             return self._analyze_builtin(element)
         return {}
 
+    def analyze_forward_output(self, name, module, module_input_output: ModuleForwardInputsOutputs):
+        if self.is_distributed_op(module):
+            module_input_output.update_output_with_args_and_kwargs()
+        return super().analyze_forward_output(name, module, module_input_output)
+
     def _analyze_tensor(self, tensor, suffix):
-        tensor_stat = self.get_stat_info(tensor)
+        tensor_stat = self.get_stat_info(tensor, self.config.async_dump)
         tensor_json = {}
         tensor_json.update({'type': 'torch.Tensor'})
         tensor_json.update({'dtype': str(tensor.dtype)})
         tensor_json.update({"shape": tensor.shape})
-        tensor_json.update({"Max": tensor_stat.max})
-        tensor_json.update({"Min": tensor_stat.min})
-        tensor_json.update({"Mean": tensor_stat.mean})
-        tensor_json.update({"Norm": tensor_stat.norm})
-        tensor_json.update({"requires_grad": tensor.requires_grad})
-
-        if tensor_stat.max is not None:
-            if np.isinf(tensor_stat.max) or np.isnan(tensor_stat.max):
-                tensor_json['Max_except_inf_nan'] = self.handle_tensor_extremum_nan_inf(tensor, "max")
-        if tensor_stat.min is not None:
-            if np.isinf(tensor_stat.min) or np.isnan(tensor_stat.min):
-                tensor_json['Min_except_inf_nan'] = self.handle_tensor_extremum_nan_inf(tensor, "min")
-
-        if self.config.summary_mode == Const.MD5:
+        if tensor_stat.stack_tensor_stat is None:
+            tensor_json.update({"Max": tensor_stat.max})
+            tensor_json.update({"Min": tensor_stat.min})
+            tensor_json.update({"Mean": tensor_stat.mean})
+            tensor_json.update({"Norm": tensor_stat.norm})
+            tensor_json.update({"requires_grad": tensor.requires_grad})
+            if tensor_stat.max is not None:
+                if np.isinf(tensor_stat.max) or np.isnan(tensor_stat.max):
+                    tensor_json['Max_except_inf_nan'] = self.handle_tensor_extremum_nan_inf(tensor, "max")
+            if tensor_stat.min is not None:
+                if np.isinf(tensor_stat.min) or np.isnan(tensor_stat.min):
+                    tensor_json['Min_except_inf_nan'] = self.handle_tensor_extremum_nan_inf(tensor, "min")
+
+        else:
+            tensor_json.update({"requires_grad": tensor.requires_grad})
+            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
@@ -207,12 +252,20 @@ class StatisticsDataProcessor(PytorchDataProcessor):
 
 
 class TensorDataProcessor(PytorchDataProcessor):
+    def dump_async_data(self):
+        for file_path, tensor in self._async_dump_cache.items():
+            save_pt(tensor.contiguous(), file_path)
+        self._async_dump_cache.clear()
+
     def _analyze_tensor(self, tensor, suffix):
         dump_data_name, file_path = self.get_save_file_path(suffix)
-        saved_tensor = tensor.clone().contiguous().detach()
-        save_pt(saved_tensor, file_path)
         single_arg = super()._analyze_tensor(tensor, suffix)
         single_arg.update({"data_name": dump_data_name})
+        if self.config.async_dump:
+            self._async_dump_cache[file_path] = tensor.clone().detach()
+        else:
+            saved_tensor = tensor.clone().contiguous().detach()
+            save_pt(saved_tensor, file_path)
         return single_arg
 
 
@@ -223,7 +276,7 @@ class OverflowCheckDataProcessor(PytorchDataProcessor):
         super().__init__(config, data_writer)
         self.has_overflow = False
         self.support_inf_nan = None
-        self.cached_inplace_api_info = {}
+        self.cached_api_info = {}
         self.cached_tensors_and_file_paths = {}
         self.bits_for_overflow = 8
         self.real_overflow_nums = 0
@@ -237,21 +290,21 @@ class OverflowCheckDataProcessor(PytorchDataProcessor):
             return True
         return False
 
-    def analyze_pre_forward_inplace(self, name, module_input_output: ModuleForwardInputsOutputs):
+    def analyze_forward_input(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         self.has_overflow = False
         self._is_support_inf_nan()
-        self.cached_inplace_api_info = super().analyze_pre_forward_inplace(name, module_input_output)
+        self.cached_api_info = super().analyze_forward_input(name, module, module_input_output)
         return None
 
-    def analyze_forward_inplace(self, name, module_input_output: ModuleForwardInputsOutputs):
+    def analyze_forward_output(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         self._is_support_inf_nan()
-        api_info_struct = super().analyze_forward_inplace(name, module_input_output)
-        if name in self.cached_inplace_api_info and name in api_info_struct:
-            self.cached_inplace_api_info[name].update(api_info_struct[name])
+        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_inplace_api_info = api_info_struct
+            self.cached_api_info = api_info_struct
         self.handle_overflow()
-        return self.cached_inplace_api_info if self.has_overflow else None
+        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
@@ -266,6 +319,13 @@ class OverflowCheckDataProcessor(PytorchDataProcessor):
         api_info_struct = super().analyze_backward(name, module, module_input_output)
         self.handle_overflow()
         return api_info_struct if self.has_overflow else None
+    
+    def analyze_params(self, name, param_name, grad):
+        self.has_overflow = False
+        self._is_support_inf_nan()
+        api_info_struct = super().analyze_params(name, param_name, grad)
+        self.handle_overflow()
+        return api_info_struct if self.has_overflow else None
 
     def handle_overflow(self):
         if not self.support_inf_nan:
@@ -340,10 +400,10 @@ class FreeBenchmarkDataProcessor(PytorchDataProcessor):
             )
         return
 
-    def analyze_pre_forward(self, name, module, module_input_output: ModuleForwardInputsOutputs):
+    def analyze_forward_input(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         self.checker.pre_forward(name, module, self, module_input_output.args, module_input_output.kwargs)
 
-    def analyze_forward(self, name, module, module_input_output: ModuleForwardInputsOutputs):
+    def analyze_forward_output(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         new_output, unequal_rows = self.checker.forward(
             name,
             module,
@@ -388,7 +448,7 @@ class KernelDumpDataProcessor(PytorchDataProcessor):
     def _print_unsupported_log(api_name):
         logger.warning(f"The kernel dump does not support the {api_name} API.")
 
-    def analyze_pre_forward(self, name, module, module_input_output):
+    def analyze_forward_input(self, name, module, module_input_output):
         if not self.enable_kernel_dump:
             return
         if is_gpu:
@@ -413,7 +473,7 @@ class KernelDumpDataProcessor(PytorchDataProcessor):
             return
         self.start_kernel_dump(self.config.kernel_config_path)
 
-    def analyze_forward(self, name, module, module_input_output):
+    def analyze_forward_output(self, name, module, module_input_output):
         if not self.enable_kernel_dump:
             return
         if self.config.is_backward_kernel_dump:

msprobe/core/data_dump/json_writer.py

@@ -15,10 +15,12 @@
 
 import csv
 import os
+import numpy as np
 
 from msprobe.core.common.const import Const, FileCheckConst
-from msprobe.core.common.file_utils import change_mode, FileOpen, save_json
+from msprobe.core.common.file_utils import change_mode, FileOpen, save_json, load_json
 from msprobe.core.common.log import logger
+from msprobe.core.common.exceptions import MsprobeException
 
 
 class DataWriter:
@@ -115,3 +117,29 @@ class DataWriter:
             self.write_stack_info_json(self.stack_file_path)
         if self.cache_construct:
             self.write_construct_info_json(self.construct_file_path)
+
+    def fill_stack_tensor_data(self):
+        self.process_stat_data_recursive(self.cache_data)
+
+    def process_stat_data_recursive(self, data, depth=0):
+        if depth > Const.MAX_DEPTH:
+            logger.error(f"The maximum depth of recursive process stat data, {Const.MAX_DEPTH} is reached.")
+            raise MsprobeException(MsprobeException.RECURSION_LIMIT_ERROR)
+        if isinstance(data, dict):
+            if "tensor_stat" in data.keys():
+                tensor_stat = data["tensor_stat"]
+                if len(tensor_stat) != Const.TENSOR_STAT_LEN or len(tensor_stat[0]) != len(tensor_stat[1]):
+                    logger.warning("Some bad data in async dump")
+                else:
+                    tensor_stat_index, tensor_stat_data = tensor_stat[0], tensor_stat[1]
+                    if hasattr(tensor_stat_data, "device") and tensor_stat_data.device != Const.CPU_LOWERCASE:
+                        tensor_stat_data = tensor_stat_data.cpu()
+                    for index, stat in zip(tensor_stat_index, tensor_stat_data):
+                        data.update({index, stat.item()})
+                del data["tensor_stat"]
+            else:
+                for key in data.keys():
+                    self.process_stat_data_recursive(data[key], depth + 1)
+        elif isinstance(data, (list, tuple)):
+            for i in data:
+                self.process_stat_data_recursive(i, depth + 1)