mindstudio-probe 8.3.2__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/core/dump/dump2db/dump2db.py

@@ -0,0 +1,409 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
+import os
+import sys
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from tqdm import tqdm
+
+from msprobe.core.common.const import Const, Data2DBConst
+from msprobe.core.common.file_utils import (
+    check_file_or_directory_path,
+    create_directory,
+    load_json,
+    recursive_chmod,
+    remove_path,
+)
+from msprobe.core.common.log import logger
+from msprobe.core.dump.dump2db.db_utils import DumpDB
+
+
+def validate_micro_step(micro_step):
+    """
+    校验 micro_step 值是否在有效范围内
+    """
+    if micro_step is not None:
+        if type(micro_step) is not int:
+            raise ValueError("Micro step must be an integer")
+        if micro_step < Data2DBConst.MIN_MICRO_STEP or micro_step > Data2DBConst.MAX_MICRO_STEP:
+            raise ValueError(f"Micro step must be between"
+                             f"{Data2DBConst.MIN_MICRO_STEP} and {Data2DBConst.MAX_MICRO_STEP}")
+
+
+def load_mapping(mapping_path):
+    if mapping_path and isinstance(mapping_path, str):
+        return load_json(mapping_path)
+    else:
+        return {}
+
+
+@dataclass
+class TensorProcessingParams:
+    """Tensor处理函数的参数载体"""
+    tensor_data: Dict[str, Any]
+    target_prefix: str
+    vpp_stage: int
+    micro_step: int
+    step: int
+    rank: int
+    metric_id: int
+    metric_type: Optional[str]
+
+
+class DumpRecordBuilder:
+    def __init__(self, db: DumpDB, data_dir, mapping, micro_step):
+        self.db = db
+        self.data_dir = data_dir
+        self.mapping = mapping
+        self.micro_step = micro_step if micro_step else None
+
+    @staticmethod
+    def extract_target_info(full_key):
+        """
+        从完整的key中提取target信息
+        格式: Module.layerN.operation.M
+        """
+        vpp_stage = 0  # 暂不提取vpp信息
+        micro_step = 0
+        target_prefix = full_key
+
+        parts = full_key.split(Const.SEP)
+        if parts and parts[-1].isdigit():
+            micro_step = int(parts[-1])
+            # 重新构建target_prefix（去掉最后一个数字部分）
+            target_prefix = Const.SEP.join(parts[:-1])
+
+        return target_prefix, vpp_stage, micro_step
+
+    @staticmethod
+    def parse_tensor_target(metric_type, tensor_type, tensor_idx):
+        """根据tensor类型和索引生成target后缀"""
+        if metric_type in [Data2DBConst.FORWARD, Data2DBConst.RECOMPUTE]:
+            if tensor_type == Const.INPUT_ARGS:
+                return f".input.{tensor_idx}"
+            elif tensor_type == Const.OUTPUT:
+                return f".output.{tensor_idx}"
+            elif tensor_type == Const.PARAMS:
+                return f".parameters.{tensor_idx}"
+        elif metric_type == Data2DBConst.BACKWARD:
+            if tensor_type == Const.INPUT:
+                return f".input.{tensor_idx}"
+            elif tensor_type == Const.OUTPUT:
+                return f".output.{tensor_idx}"
+        return ""
+
+    @staticmethod
+    def process_tensor_value(value):
+        """处理统计量值  将inf转换为float极值"""
+        if value is None:
+            return sys.float_info.max
+        elif value == float("inf"):
+            return sys.float_info.max - 1  # 最大float值
+        elif value == float("-inf"):
+            return sys.float_info.min + 1  # 最小float值
+        elif isinstance(value, (int, float)):
+            return float(value)
+        else:
+            return None
+
+    def import_data(self):
+        """导入数据"""
+        logger.info("Scanning data directory...")
+
+        # 扫描所有step和rank
+        max_rank = 0
+        min_step = float('inf')
+        max_step = 0
+        valid_ranks = defaultdict(list)
+
+        # 首先扫描目录结构获取step和rank信息
+        for step_dir in os.listdir(self.data_dir):
+            if not step_dir.startswith('step'):
+                continue
+            try:
+                step = int(step_dir[4:])  # 提取step数字
+            except ValueError:
+                continue
+            min_step = min(min_step, step *
+                           self.micro_step if self.micro_step else step)
+            max_step = max(max_step, (step + 1) *
+                           self.micro_step if self.micro_step else step)
+
+            step_path = os.path.join(self.data_dir, step_dir)
+            for rank_dir in os.listdir(step_path):
+                if not rank_dir.startswith('rank'):
+                    continue
+                if rank_dir == "rank":
+                    rank = 0
+                else:
+                    try:
+                        rank = int(rank_dir[4:])  # 提取rank数字
+                    except ValueError:
+                        continue
+
+                rank_path = os.path.join(step_path, rank_dir)
+                json_path = os.path.join(rank_path, "dump.json")
+                if os.path.exists(json_path):
+                    valid_ranks[step].append((rank, json_path))
+                    max_rank = max(max_rank, rank)
+
+        if min_step == float('inf'):
+            logger.warning(f"No valid step directories found in: {self.data_dir},"
+                           f"looking for directories starting with 'step' (e.g., step0, step1")
+            return
+
+        # 更新全局统计信息
+        global_stats = {
+            "max_rank": max_rank,
+            "min_step": min_step,
+            "max_step": max_step
+        }
+        for metric_name in Data2DBConst.METRICS:
+            global_stats[metric_name] = Data2DBConst.ORDERED_STAT
+        self.db.init_global_stats_data(global_stats)
+
+        for step in tqdm(sorted(valid_ranks.keys()), desc="Processing steps"):
+            step_dir = f"step{step}"
+            step_path = os.path.join(self.data_dir, step_dir)
+
+            check_file_or_directory_path(step_path, isdir=True)
+
+            # 为当前step的所有有效rank创建进度条
+            if not valid_ranks[step]:
+                continue
+            for rank, json_path in tqdm(valid_ranks[step], desc=f"Step {step} ranks", leave=False):
+                self._process_dump_file(
+                    json_path, step, rank)
+        self.db.extract_tags_from_processed_targets()
+
+    def _determine_metric_type(self, full_key, tensor_data):
+        for key, value in self.mapping.items():
+            full_key = full_key.replace(key, value)
+
+        # 确定metric类型
+        metric_type = None
+        if f"{Const.SEP}{Const.FORWARD}{Const.SEP}" in full_key or \
+                full_key.endswith(f"{Const.SEP}{Const.FORWARD}"):
+            if tensor_data.get('is_recompute', False):
+                metric_type = Data2DBConst.RECOMPUTE
+            else:
+                metric_type = Data2DBConst.FORWARD
+            full_key = full_key.replace(
+                f"{Const.SEP}{Const.FORWARD}{Const.SEP}", Const.SEP)  # module
+            full_key = full_key.replace(
+                f"{Const.SEP}{Const.FORWARD}", "")  # api
+        elif f"{Const.SEP}{Const.BACKWARD}{Const.SEP}" in full_key or \
+                full_key.endswith(f"{Const.SEP}{Const.BACKWARD}"):
+            metric_type = Data2DBConst.BACKWARD
+            full_key = full_key.replace(
+                f"{Const.SEP}{Const.BACKWARD}{Const.SEP}", Const.SEP)
+            full_key = full_key.replace(f"{Const.SEP}{Const.BACKWARD}", "")
+        elif full_key.endswith(Const.PARAMS_GRAD):
+            metric_type = Data2DBConst.PARAMETERS_GRAD
+            full_key = full_key.replace(f"{Const.SEP}{Const.PARAMS_GRAD}", "")
+        # fsdp
+        elif len(full_key.split(Const.SEP)) >= 3:
+            parts = full_key.split(Const.SEP)
+            if parts[-2] == Const.PARAMS_GRAD:
+                metric_type = Data2DBConst.PARAMETERS_GRAD
+                full_key = Const.SEP.join(parts[:-2])
+
+        return metric_type, full_key
+
+    def _add_tensor_data(self, tensor, target_name, tensor_params: TensorProcessingParams, batch_data):
+        """添加tensor数据方法"""
+
+        cache_key = (target_name, tensor_params.vpp_stage,
+                     tensor_params.micro_step)
+
+        # 如果缓存中不存在，创建临时ID
+        cache_id_dict = self.db.cache_targets(
+            cache_key, tensor_params.metric_id)
+
+        # 准备数据行, 这里id还是个临时id, 需要更新后读取, 第三个实际为{"id": 0}
+        row_data = [tensor_params.rank, tensor_params.step, cache_id_dict, tensor_params.metric_id]
+        for stat in Data2DBConst.ORDERED_STAT:
+            value = tensor.get(stat.capitalize(), None)
+            row_data.append(DumpRecordBuilder.process_tensor_value(value))
+        batch_data.append((row_data))
+
+    def _process_forward_data(self, tensor_params: TensorProcessingParams, batch_data):
+        """处理forward/recompute数据"""
+        # 处理input_args
+        if Const.INPUT_ARGS in tensor_params.tensor_data:
+            for idx, tensor in enumerate(tensor_params.tensor_data[Const.INPUT_ARGS]):
+                if isinstance(tensor, dict) and tensor.get(Const.TYPE) in Data2DBConst.SUPPORT_TYPE:
+                    if tensor.get(Const.DTYPE) not in Data2DBConst.SUPPORT_DTYPE:
+                        continue
+                    target_suffix = DumpRecordBuilder.parse_tensor_target(
+                        Data2DBConst.FORWARD, Const.INPUT_ARGS, idx)
+                    target_name = tensor_params.target_prefix + target_suffix
+                    self._add_tensor_data(
+                        tensor, target_name, tensor_params, batch_data)
+
+        # 处理output
+        if Const.OUTPUT in tensor_params.tensor_data:
+            for idx, tensor in enumerate(tensor_params.tensor_data[Const.OUTPUT]):
+                if isinstance(tensor, dict) and tensor.get(Const.TYPE) in Data2DBConst.SUPPORT_TYPE:
+                    if tensor.get(Const.DTYPE) not in Data2DBConst.SUPPORT_DTYPE:
+                        continue
+                    target_suffix = DumpRecordBuilder.parse_tensor_target(
+                        Data2DBConst.FORWARD, Const.OUTPUT, idx)
+                    target_name = tensor_params.target_prefix + target_suffix
+                    self._add_tensor_data(
+                        tensor, target_name, tensor_params, batch_data)
+
+        # 处理parameters
+        if Const.PARAMS in tensor_params.tensor_data:
+            for param_name, param_tensors in tensor_params.tensor_data[Const.PARAMS].items():
+                if not (isinstance(param_tensors, list) and param_tensors and isinstance(param_tensors[0], dict)):
+                    continue
+                if param_tensors[0].get(Const.TYPE) in Data2DBConst.SUPPORT_TYPE:
+                    if param_tensors[0].get(Const.DTYPE) not in Data2DBConst.SUPPORT_DTYPE:
+                        continue
+                    target_suffix = DumpRecordBuilder.parse_tensor_target(
+                        Data2DBConst.FORWARD, Const.PARAMS, param_name)
+                    target_name = tensor_params.target_prefix + target_suffix
+                    self._add_tensor_data(
+                        param_tensors[0], target_name, tensor_params, batch_data)
+
+    def _process_parameters_data(self, tensor_params: TensorProcessingParams, batch_data):
+        """处理parameters数据"""
+        for param_name, param_tensors in tensor_params.tensor_data.items():
+            if not (isinstance(param_tensors, list) and param_tensors and isinstance(param_tensors[0], dict)):
+                continue
+            if param_tensors[0].get(Const.TYPE) in Data2DBConst.SUPPORT_TYPE:
+                if param_tensors[0].get(Const.DTYPE) not in Data2DBConst.SUPPORT_DTYPE:
+                    continue
+                target_name = tensor_params.target_prefix + f".{param_name}"
+                self._add_tensor_data(
+                    param_tensors[0], target_name, tensor_params, batch_data)
+
+    def _process_backward_data(self, tensor_params: TensorProcessingParams, batch_data):
+        """处理backward数据"""
+        # 处理input
+        if Const.INPUT in tensor_params.tensor_data:
+            for idx, tensor in enumerate(tensor_params.tensor_data[Const.INPUT]):
+                if isinstance(tensor, dict) and tensor.get(Const.TYPE) in Data2DBConst.SUPPORT_TYPE:
+                    if tensor.get(Const.DTYPE) not in Data2DBConst.SUPPORT_DTYPE:
+                        continue
+                    target_suffix = DumpRecordBuilder.parse_tensor_target(
+                        Data2DBConst.BACKWARD, Const.INPUT, idx)
+                    target_name = tensor_params.target_prefix + target_suffix
+                    self._add_tensor_data(
+                        tensor, target_name, tensor_params, batch_data)
+
+        # 处理output
+        if Const.OUTPUT in tensor_params.tensor_data:
+            for idx, tensor in enumerate(tensor_params.tensor_data[Const.OUTPUT]):
+                if isinstance(tensor, dict) and tensor.get(Const.TYPE) in Data2DBConst.SUPPORT_TYPE:
+                    if tensor.get(Const.DTYPE) not in Data2DBConst.SUPPORT_DTYPE:
+                        continue
+                    target_suffix = DumpRecordBuilder.parse_tensor_target(
+                        Data2DBConst.BACKWARD, Const.OUTPUT, idx)
+                    target_name = tensor_params.target_prefix + target_suffix
+                    self._add_tensor_data(
+                        tensor, target_name, tensor_params, batch_data)
+
+    def _process_dump_file(self, json_path, step, rank):
+        """处理单个dump.json文件"""
+        data = load_json(json_path)
+        if 'data' not in data or not isinstance(data['data'], dict):
+            return
+
+        batch_data = []
+        # 预先计算所有metric_type的table_name
+
+        for i, (ori_key, tensor_data) in enumerate(data['data'].items()):
+            if not isinstance(ori_key, str) or not isinstance(tensor_data, dict):
+                continue
+            metric_type, full_key = self._determine_metric_type(
+                ori_key, tensor_data)
+            if not metric_type:
+                continue
+
+            # 使用缓存获取metric_id和stats
+            metric_id = self.db.get_metric_id(metric_type)
+            target_prefix, vpp_stage, mstep = DumpRecordBuilder.extract_target_info(
+                full_key)
+            if self.micro_step:
+                current_mstep = step * self.micro_step + mstep
+                mstep = 0
+            else:
+                current_mstep = step
+
+            tensor_params = TensorProcessingParams(
+                tensor_data=tensor_data,
+                target_prefix=target_prefix,
+                vpp_stage=vpp_stage,
+                micro_step=mstep,
+                step=current_mstep,
+                rank=rank,
+                metric_id=metric_id,
+                metric_type=metric_type
+            )
+
+            # 处理不同类型的tensor数据
+            if metric_type in [Data2DBConst.FORWARD, Data2DBConst.RECOMPUTE]:
+                self._process_forward_data(tensor_params, batch_data)
+            elif metric_type == Data2DBConst.PARAMETERS_GRAD:
+                self._process_parameters_data(tensor_params, batch_data)
+            elif metric_type == Data2DBConst.BACKWARD:
+                self._process_backward_data(tensor_params, batch_data)
+
+            if len(batch_data) % Data2DBConst.BATCH_SIZE == 0 and len(batch_data) > 0:
+                self.db.batch_insert_targets()
+                self.db.batch_insert_data(batch_data)
+                batch_data = []
+
+        self.db.batch_insert_targets()
+        self.db.batch_insert_data(batch_data)
+
+
+def _data2db_service_parser(parser):
+    parser.add_argument('--db', type=str, required=True,
+                        help='Path to SQLite database file')
+    parser.add_argument('--data', type=str, required=True,
+                        help='Path to dump output directory')
+    parser.add_argument('--mapping', type=str, default=None,
+                        help='Path to optional JSON mapping file')
+    parser.add_argument('--micro_step', type=int, default=None,
+                        help='Specific micro step value to split data in one step (must be between 1 and 10000)')
+
+
+def _data2db_command(args):
+    data_path = args.data
+
+    check_file_or_directory_path(data_path, isdir=True, is_strict=True)
+    create_directory(args.db)
+    db_path = os.path.join(args.db, "monitor_metrics.db")
+
+    if os.path.exists(db_path):
+        logger.warning(f"Existing path {db_path} will be recovered")
+        remove_path(db_path)
+
+    micro_step = args.micro_step
+    validate_micro_step(micro_step)
+    mapping = load_mapping(args.mapping)
+    db = DumpDB(db_path)
+    builder = DumpRecordBuilder(
+        db, data_path, mapping=mapping, micro_step=micro_step)
+    builder.import_data()
+
+    recursive_chmod(args.db)

msprobe/core/{hook_manager.py → dump/hook_manager.py}

@@ -1,17 +1,18 @@
-# Copyright (c) 2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 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.
-# You may obtain a copy of the License at
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 import gc
 import os
@@ -21,7 +22,7 @@ from collections import defaultdict
 
 from msprobe.core.common.runtime import Runtime
 from msprobe.core.common.utils import Const, ThreadSafe
-from msprobe.core.data_dump.data_processor.base import (ModuleBackwardInputsOutputs, ModuleForwardInputsOutputs)
+from msprobe.core.dump.data_dump.data_processor.base import (ModuleBackwardInputsOutputs, ModuleForwardInputsOutputs)
 
 
 class HookSet:
@@ -45,20 +46,17 @@ class BaseHookManager(ABC):
     inner_api_count = defaultdict(int)
     hook_handle_dict = {}
     params_grad_info = {}
+    grad_hook_call = {}
 
     def __init__(self, data_collector, config):
         self.data_collector = data_collector
         self.config = config
+        self.current_pid = self._pid
 
     @property
     def _pid(self):
         return os.getpid()
 
-    @property
-    @abstractmethod
-    def _is_recompute(self):
-        pass
-
     @staticmethod
     def reset_status():
         BaseHookManager.inner_switch = defaultdict(bool)
@@ -82,6 +80,15 @@ class BaseHookManager(ABC):
         if hasattr(module, 'msprobe_input_kwargs') and tid in module.msprobe_input_kwargs:
             del module.msprobe_input_kwargs[tid]
 
+    @staticmethod
+    def _get_grad_hook_call_index(ori_name, param_name):
+        if ori_name not in BaseHookManager.grad_hook_call:
+            BaseHookManager.grad_hook_call[ori_name] = [0, param_name]
+        else:
+            if BaseHookManager.grad_hook_call.get(ori_name)[1] == param_name:
+                BaseHookManager.grad_hook_call[ori_name][0] += 1
+        return BaseHookManager.grad_hook_call.get(ori_name)[0]
+
     @staticmethod
     @abstractmethod
     def _no_grad_context():
@@ -103,7 +110,7 @@ class BaseHookManager(ABC):
         pass
 
     @abstractmethod
-    def build_hook(self):
+    def build_hook(self, hook_type, name):
         pass
 
     @abstractmethod
@@ -115,7 +122,7 @@ class BaseHookManager(ABC):
         pass
 
     @abstractmethod
-    def _register_backward_pre_hook(self, module, full_backward_name, output):
+    def _register_backward_pre_hook(self, module, full_backward_name, args, kwargs, output):
         pass
 
     @abstractmethod
@@ -126,44 +133,17 @@ class BaseHookManager(ABC):
     def _need_exchange(self, module):
         pass
 
+    @abstractmethod
     def _register_param_hook(self, name, module, params_dict):
-        ori_name = name.rsplit(Const.SEP, 2)[0]
-        grad_name = ori_name + Const.SEP + Const.PARAMS_GRAD
-        # 首次执行前向hook时，添加params_grad_name属性，并注册参数hook
-        setattr(module, 'params_grad_name', grad_name)
-        # data_mode为forward时，不注册参数hook
-        if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
-            for param_name, param in params_dict.items():
-                if param.requires_grad:
-                    name = ori_name + Const.SEP + param_name
-                    old_handle = BaseHookManager.hook_handle_dict.get(name)
-                    if old_handle and hasattr(old_handle, "remove"):
-                        old_handle.remove()
-                    handle = param.register_hook(self._build_grad_hook(ori_name, param_name))
-                    BaseHookManager.hook_handle_dict[name] = handle
-
-    def _init_params_grad_info(self, module, params_dict):
-        '''
-        初始化参数梯度信息, 在前向hook结束后, 将参数梯度信息写入cache_data中用于占位
-        '''
-        if not params_dict:
-            return
-        if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
-            grad_name = module.params_grad_name if hasattr(module, 'params_grad_name') else None
-            # 判断是否已经在cache_data中进行了占位, 若没有则先写入cache_data中
-            if not BaseHookManager.params_grad_info.get(grad_name):
-                data_info = {grad_name: {key: [None] for key, value in params_dict.items() if value.requires_grad}}
-                # 当模块中的参数有requires_grad属性为True时，才会进行梯度计算，此时才需要占位
-                if data_info.get(grad_name):
-                    # 将grad_name的data_info先写入cache_data中, 梯度计算后再更新
-                    self.data_collector.handle_data(grad_name, data_info,
-                                                    flush=self.data_collector.data_processor.is_terminated)
-                    self.data_collector.params_grad_record[grad_name] = True
-                # 记录当前模块的参数梯度信息已占位
-                BaseHookManager.params_grad_info[grad_name] = True
+        pass
+
+    def is_child_process(self):
+        return self.current_pid != self._pid
 
     def _should_execute_hook(self, hook_type, tid, is_forward=True):
         is_api_hook = hook_type == Const.API
+        if self.is_child_process():
+            return False
         if BaseHookManager.inner_switch[tid]:
             return False
         if not is_api_hook and not Runtime.is_running:
@@ -174,21 +154,6 @@ class BaseHookManager(ABC):
             return False
         return True
 
-    def _build_grad_hook(self, ori_name, param_name):
-        def hook_fn(grad):
-            tid = threading.get_ident()
-            if not self._should_execute_hook(Const.MODULE, tid):
-                return
-            with ThreadSafe():
-                original_state = self.ensure_gc_enabled()
-                BaseHookManager.inner_switch[tid] = True
-                self.data_collector.params_data_collect(ori_name, param_name, self._pid, grad)
-                BaseHookManager.inner_switch[tid] = False
-                self.restore_gc_state(original_state)
-            return
-
-        return hook_fn
-
     def _build_forward_pre_hook(self, hook_type, api_name):
         def forward_pre_hook(module, args, kwargs=None):
             if hook_type == Const.MODULE:
@@ -220,8 +185,7 @@ class BaseHookManager(ABC):
                         full_forward_name,
                         module,
                         self._pid,
-                        module_input_output,
-                        self._is_recompute
+                        module_input_output
                     )
                 BaseHookManager.inner_switch[tid] = False
                 self.restore_gc_state(original_state)
@@ -257,7 +221,7 @@ class BaseHookManager(ABC):
                 if hook_type == Const.API:
                     full_forward_name = api_name + str(self._get_count(api_name)) + Const.SEP + Const.FORWARD
                     full_backward_name = api_name + str(self._get_count(api_name)) + Const.SEP + Const.BACKWARD
-                    output = self._register_backward_pre_hook(module, full_backward_name, output)
+                    output = self._register_backward_pre_hook(module, full_backward_name, args, kwargs, output)
 
                 with self._no_grad_context():
                     if hook_type == Const.MODULE:
@@ -270,28 +234,19 @@ class BaseHookManager(ABC):
                             api_name,
                             module,
                             self._pid,
-                            module_input_output,
-                            self._is_recompute
+                            module_input_output
                         )
-                        self._init_params_grad_info(module, params_dict)
                     else:
                         self.data_collector.update_api_or_module_name(full_forward_name)
                         self.data_collector.forward_output_data_collect(
                             full_forward_name,
                             module,
                             self._pid,
-                            module_input_output,
-                            self._is_recompute
+                            module_input_output
                         )
                         self._add_count(api_name)
                         BaseHookManager.inner_api_count[tid] -= 1
                     self._clear_input_kwargs(module, tid)
-
-                    if self.data_collector.if_return_forward_new_output():
-                        forward_new_output = self.data_collector.get_forward_new_output()
-                        BaseHookManager.inner_switch[tid] = False
-                        return forward_new_output
-
                 BaseHookManager.inner_switch[tid] = False
                 self.restore_gc_state(original_state)
                 return output
@@ -318,12 +273,8 @@ class BaseHookManager(ABC):
                     full_name,
                     module,
                     self._pid,
-                    module_input_output,
-                    self._is_recompute
+                    module_input_output
                 )
-                if hook_type == Const.MODULE:
-                    params_dict = self._get_params_dict(module)
-                    self.data_collector.params_data_collect_in_bw_hook(params_dict, full_name)
                 BaseHookManager.inner_switch[tid] = False
                 self.restore_gc_state(original_state)
 

msprobe/core/dump/kernel_dump/kernel_config.py

@@ -0,0 +1,34 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
+import os
+
+from msprobe.core.common.file_utils import save_json
+
+
+def create_kernel_config_json(dump_path, cur_rank):
+    kernel_config_name = "kernel_config.json" if cur_rank == '' else f"kernel_config_{cur_rank}.json"
+    kernel_config_path = os.path.join(dump_path, kernel_config_name)
+    config_info = {
+        "dump": {
+            "dump_list": [],
+            "dump_path": dump_path,
+            "dump_mode": "all",
+            "dump_op_switch": "on"
+        }
+    }
+    save_json(kernel_config_path, config_info, indent=4)
+    return kernel_config_path