mindstudio-probe 8.2.0__py3-none-any.whl → 8.2.1__py3-none-any.whl

msprobe/docs/22.visualization_MindSpore.md

@@ -21,40 +21,37 @@
 
 ### 1.1 安装msprobe工具
 
-[msprobe工具安装](https://gitee.com/ascend/mstt/blob/master/debug/accuracy_tools/msprobe/docs/01.installation.md)
+[msprobe工具安装](./01.installation.md)
 
 ### 1.2 安装tb_graph_ascend
 
 **请安装tb_graph_ascend，否则无法解析构图结果。**
 
-``pip3 install tb-graph-ascend``即可。
-
-如需升级工具，请先``pip3 uninstall tb-graph-ascend``再``pip3 install tb-graph-ascend``即可。
+[tb_graph_ascend安装](../../../../plugins/tensorboard-plugins/tb_graph_ascend#2-安装方式)
 
 ## 2.模型结构数据采集
-[MindSpore场景的精度数据采集](https://gitee.com/ascend/mstt/blob/master/debug/accuracy_tools/msprobe/docs/06.data_dump_MindSpore.md)
+[PyTorch场景的数据采集](./06.data_dump_MindSpore.md)
 
 **仅支持动态图场景，需要选择level为L0（cell信息）或者mix（cell信息+api信息），才能采集到模型结构数据，即采集结果件construct.json内容不为空**。
 
 ## 3.生成图结构文件
 
 ### 3.1 构图命令行说明
-
+  
 **命令示例如下**：
 ```
 msprobe -f mindspore graph -i ./compare.json -o ./output
 ```
 **命令行参数说明**：
 
-| 参数名               | 说明                                                                                                                                                                                                                                                    | 是否必选 |
-|-------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -------- |
-| -f 或 --framework    | 指定训练框架。mindspore。                                                                                                                                                     | 是                                 |
-| -i 或 --input_path   | 指定比对文件，参考[比对文件说明](#313-比对文件说明)                                                                                                                                                                                                                        | 是       |
-| -o 或 --output_path  | 配置比对结果文件存盘目录，str 类型。文件名称基于时间戳自动生成，格式为：`compare_{timestamp}.vis或build_{timestamp}.vis`。                                                                                                                                                                | 是       |
+| 参数名               | 说明                                                                                                                                                                                                                                                                                                                                                                    | 是否必选 |
+|-------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -------- |
+| -f 或 --framework    | 指定训练框架。mindspore。                                                                                                                                                                                                                                                                                                                                                     | 是                                 |
+| -i 或 --input_path   | 指定比对文件，参考[比对文件说明](#313-比对文件说明)                                                                                                                                                                                                                                                                                                                                        | 是       |
+| -o 或 --output_path  | 配置比对结果文件存盘目录，str 类型。文件名称基于时间戳自动生成，格式为：`compare_{timestamp}.vis.db或build_{timestamp}.vis.db`。                                                                                                                                                                                                                                                                          | 是       |
 | -lm 或 --layer_mapping| 跨框架比对，MindSpore和PyTorch的比对场景。配置该参数时表示开启跨框架Layer层的比对功能，指定模型代码中的Layer层后，可以识别对应dump数据中的模块或API。需要指定自定义映射文件*.yaml。自定义映射文件的格式请参见[自定义映射文件（Layer）](#71-自定义映射文件layer), 如何配置自定义映射文件请参考[模型分级可视化如何配置layer mapping映射文件](./visualization/layer_mapping_example.md)。配置该参数后，将仅按节点名称进行比对，忽略节点的 type 和 shape。如果调试侧和标杆侧有名称不同的节点，则需要配置自定义映射文件，-lm参数传入自定义映射文件路径；如果调试侧和标杆侧节点名称相同，则仅指定-lm即可。 | 否    |
-| -oc 或 --overflow_check | 是否开启溢出检测模式，开启后会在输出vis文件中（`compare_{timestamp}.vis或build_{timestamp}.vis`）对每个溢出节点进行标记溢出等级，溢出等级说明参考[溢出等级说明](#312-溢出等级说明)                                                                                                                                | 否    |
-| -f 或 --fuzzy_match     | 是否开启模糊匹配，bool类型。模糊匹配说明参考[匹配说明](#311-匹配说明)                                                                                                                                                                                                             | 否    |
-| -cs 或 --complete_stack | 是否使用完整的堆栈信息，bool类型。默认使用精简的堆栈信息，数据量小有助于增加流畅度。完整堆栈和精简堆栈信息参考[堆栈信息说明](#72-堆栈信息说明)                                                                                                                                                                         | 否    |
+| -oc 或 --overflow_check | 是否开启溢出检测模式，开启后会在输出db文件中（`compare_{timestamp}.vis.db或build_{timestamp}.vis.db`）对每个溢出节点进行标记溢出等级，溢出等级说明参考[溢出等级说明](#312-溢出等级说明)                                                                                                                                                                                                                                           | 否    |
+| -f 或 --fuzzy_match     | 是否开启模糊匹配，bool类型。模糊匹配说明参考[匹配说明](#311-匹配说明)                                                                                                                                                                                                                                                                                                                             | 否    |
 
 #### 3.1.1 匹配说明
 
@@ -140,7 +137,7 @@ msprobe -f mindspore graph -i ./compare.json -o ./output
 
 3.md5：dump了API和Module的输入输出数据统计信息和md5信息。
 
-dump类型如何配置见[数据采集配置文件介绍](https://gitee.com/ascend/mstt/blob/master/debug/accuracy_tools/msprobe/docs/02.config_introduction.md)
+dump类型如何配置见[数据采集配置文件介绍](https://gitcode.com/Ascend/mstt/blob/master/debug/accuracy_tools/msprobe/docs/02.config_introduction.md)
 
 **1. 准备比对文件**：
 
@@ -213,25 +210,15 @@ npu_path或bench_path格式：必须只包含rank+数字格式的文件夹，且
 ```
 msprobe -f mindspore graph -i ./compare.json -o ./output
 ```
-比对完成后将在**output**下生成n个**vis后缀文件**。
+比对完成后将在**output**下生成1个**vis.db后缀文件**。
 
 图构建：
 ```
-├── build_rank0_{timestamp}.vis
-├── build_rank1_{timestamp}.vis
-├── build_rank2_{timestamp}.vis
-├── build_rank3_{timestamp}.vis
-├── ...
-├── build_rankn_{timestamp}.vis
+├── build_{timestamp}.vis.db
 ```
 图比对：
 ```
-├── compare_rank0_{timestamp}.vis
-├── compare_rank1_{timestamp}.vis
-├── compare_rank2_{timestamp}.vis
-├── compare_rank3_{timestamp}.vis
-├── ...
-├── compare_rankn_{timestamp}.vis
+├── compare_{timestamp}.vis.db
 ```
 ##### 3.2.3.2 多step批量构建或比对
 批量构建或比对多个step下的所有rank的数据
@@ -278,33 +265,15 @@ npu_path或bench_path格式：必须只包含step+数字格式的文件夹，且
 ```
 msprobe -f mindspore graph -i ./compare.json -o ./output
 ```
-比对完成后将在**output**下生成若干个**vis后缀文件**。
+比对完成后将在**output**下生成1个**vis.db后缀文件**。
 
 图构建：
 ```
-├── build_step0_rank0_{timestamp}.vis
-├── build_step0_rank1_{timestamp}.vis
-├── build_step0_rank2_{timestamp}.vis
-├── build_step0_rank3_{timestamp}.vis
-├── build_step1_rank0_{timestamp}.vis
-├── build_step1_rank1_{timestamp}.vis
-├── build_step1_rank2_{timestamp}.vis
-├── build_step1_rank3_{timestamp}.vis
-├── ...
-├── build_stepn_rankn_{timestamp}.vis
+├── build_{timestamp}.vis.db
 ```
 图比对：
 ```
-├── compare_step0_rank0_{timestamp}.vis
-├── compare_step0_rank1_{timestamp}.vis
-├── compare_step0_rank2_{timestamp}.vis
-├── compare_step0_rank3_{timestamp}.vis
-├── compare_step1_rank0_{timestamp}.vis
-├── compare_step1_rank1_{timestamp}.vis
-├── compare_step1_rank2_{timestamp}.vis
-├── compare_step1_rank3_{timestamp}.vis
-├── ...
-├── compare_stepn_rankn_{timestamp}.vis
+├── compare_{timestamp}.vis.db
 ```
 
 #### 3.2.4 仅模型结构比对
@@ -413,9 +382,11 @@ tensorboard --logdir out_path
 
 ### 5.1 浏览器打开图
 推荐使用谷歌浏览器，在浏览器中输入机器地址+端口号回车，出现TensorBoard页面，其中/#graph_ascend会自动拼接。
+
 ![vis_browser_1](./img/visualization/vis_browser_1.png)
 
 如果您切换了TensorBoard的其他功能，此时想回到模型分级可视化页面，可以点击左上方的**GRAPH_ASCEND**
+
 ![vis_browser_2](./img/visualization/vis_browser_2.png)
 
 ### 5.2 查看图
@@ -530,61 +501,6 @@ yaml文件中只需配置MindSpore与PyTorch模型代码中功能一致但名称
 
 ![ms_dump](./img/ms_layer.png)
 
-### 7.2 堆栈信息说明
-
-**精简堆栈**
-
-保留一条当前模块或api的调用信息
-
-```json
-{
-    "Cell.model.language_model.embedding.word_embeddings.reduce_scatter_to_sp_region.ReduceScatterToSequenceParallelRegion.forward.0": [
-        "File /home/mindformers/experimental/distri_cores/tensor_parallel/layers.py, line 770, in construct, \n output = self.reduce_scatter_to_sp_region(output_parallel)"
-        ]
-}
-```
-
-**完整堆栈**
-
-当前模块或api完整的调用信息
-
-```json
-{
-    "Cell.model.language_model.embedding.word_embeddings.reduce_scatter_to_sp_region.ReduceScatterToSequenceParallelRegion.forward.0": [
-        "File /home/mindspore/nn/cell.py, line 507, in _run_construct, \n output = self._run_forward_hook(inputs, output)",
-        "File /home/mindspore/nn/cell.py, line 759, in _complex_call, \n output = self._run_construct(*args, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 747, in __call__, \n return self._complex_call(*args, **kwargs)",
-        "File /home/mindformers/experimental/distri_cores/tensor_parallel/layers.py, line 770, in construct, \n output = self.reduce_scatter_to_sp_region(output_parallel)",
-        "File /home/mindspore/nn/cell.py, line 2462, in _backward_hook_construct, \n outputs = self.construct(outputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 498, in _run_construct, \n output = self._backward_hook_construct(*inputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 745, in __call__, \n return self._run_construct(*args, **kwargs)",
-        "File /home/mindformers/experimental/distri_cores/transformer/language_model.py, line 151, in construct, \n embeddings = self.word_embeddings(input_ids)",
-        "File /home/mindspore/nn/cell.py, line 2460, in _backward_hook_construct, \n outputs = self.construct(*outputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 498, in _run_construct, \n output = self._backward_hook_construct(*inputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 745, in __call__, \n return self._run_construct(*args, **kwargs)",
-        "File /home/mindformers/experimental/distri_cores/transformer/language_model.py, line 391, in construct, \n text_embedding_out = self.embedding(enc_input_ids, enc_position_ids,",
-        "File /home/mindspore/nn/cell.py, line 2460, in _backward_hook_construct, \n outputs = self.construct(*outputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 498, in _run_construct, \n output = self._backward_hook_construct(*inputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 745, in __call__, \n return self._run_construct(*args, **kwargs)",
-        "File /home/model/gpt_model.py, line 104, in construct, \n lm_output = self.language_model(tokens,",
-        "File /home/mindspore/nn/cell.py, line 2460, in _backward_hook_construct, \n outputs = self.construct(*outputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 498, in _run_construct, \n output = self._backward_hook_construct(*inputs, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 745, in __call__, \n return self._run_construct(*args, **kwargs)",
-        "File /home/mindformers/experimental/distri_cores/pipeline_parallel/pipeline_cell.py, line 429, in construct, \n return self.model(*inputs)",
-        "File /home/mindspore/nn/cell.py, line 757, in _complex_call, \n output = self.construct(*args, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 747, in __call__, \n return self._complex_call(*args, **kwargs)",
-        "File /home/mindformers/experimental/distri_cores/pipeline_parallel/schedules.py, line 121, in run_forward, \n output_tensor = model(*input_data, recv_data=None)",
-        "File /home/mindformers/experimental/distri_cores/pipeline_parallel/schedules.py, line 735, in forward_backward_pipelining_without_interleaving, \n micro_input_data = run_forward(*micro_input_data,",
-        "File /home/mindformers/experimental/distri_cores/training.py, line 409, in forward_backward_with_pipelining, \n loss, logits, grads = forward_backward_pipelining_without_interleaving(",
-        "File /home/mindformers/experimental/distri_cores/training.py, line 533, in construct, \n (loss, _), grads = self.forward_backward_func(*inputs_tuple, loss_scale=current_step_loss_scale, **inputs_dict)",
-        "File /home/mindspore/nn/cell.py, line 757, in _complex_call, \n output = self.construct(*args, **kwargs)",
-        "File /home/mindspore/nn/cell.py, line 747, in __call__, \n return self._complex_call(*args, **kwargs)",
-        "File /home/mindformers/experimental/distri_cores/training.py, line 655, in train, \n loss, is_finite, loss_scale, learning_rate = train_one_step_cell(**data)",
-        "File /home/model/pretrain_gpt.py, line 303, in main, \n train(",
-        "File /home/model/pretrain_gpt.py, line 316, in <module>, \n main()"
-        ]
-}
-```
 # FAQ
 1. 图比对场景，节点呈现灰色，且没有精度比对数据，怎么处理？
 

msprobe/docs/23.generate_operator_PyTorch.md

@@ -13,7 +13,7 @@ b. 在生成单API脚本时可以选择由工具构造随机数获得 dump 数
 ### 前提
 1. 安装 msprobe。详见[ msprobe 安装](./01.installation.md)章节。
 2. 已完成对训练过程的dump，获得dump.json文件。
-   [PyTorch场景的数据采集](https://gitee.com/ascend/mstt/blob/master/debug/accuracy_tools/msprobe/docs/05.data_dump_PyTorch.md)
+   [PyTorch场景的数据采集](https://gitcode.com/Ascend/mstt/blob/master/debug/accuracy_tools/msprobe/docs/05.data_dump_PyTorch.md)
    
    **目前仅支持复现API级的数据，故dump时level可选择L1（API信息）或者mix（module信息+API信息）。如需复现真实数据场景的API脚本，dump时task应选择tensor，如需复现随机数据场景的API脚本，dump时task选择statistics**。
 3. 发现某个算子疑似存在精度问题，并得知算子名，如Functional.softmax.3、Tensor.add.0、Torch.matmul.5等

msprobe/docs/26.data_dump_PyTorch_baseline.md

@@ -1,14 +1,14 @@
 # PyTorch 场景的精度数据采集基线
 
-## "statistics"模式（未开启md5）采集时间膨胀参考基线
+## "statistics"模式采集时间膨胀参考基线
 
-该基线为PyTorch框架下，使用"statistics"模式采集数据性能膨胀的参考基线。本基线测试了LLAMA2-7B语言大模型在不同采集模式8卡下的时间膨胀。
+该基线为PyTorch框架下，使用"statistics"模式采集数据性能膨胀的参考基线。本基线测试了单层 DeepSeek 大模型在不同采集模式8卡下的时间膨胀。
 
-| 采集模式 | 无工具 (耗时) |  加工具但未使能 Dump (耗时)   |  加工具并使能 Dump (耗时)  |
-|:--------:|:--------:|:--------------------:|:------------------:|
-| L0       | ≈17.4 s  |    ≈17.4 s （无膨胀）     | ≈78.4 s  （膨胀4.5倍）  |
-| L1       | ≈17.4 s  |  ≈20.7 s  （膨胀1.2倍）   |  ≈353 s  （膨胀20倍）   |
-| mix      | ≈17.4 s  | ≈20.7 s  （膨胀1.2倍） | ≈430 s  （膨胀24.7 倍） |
+| 采集模式 | 无工具 (耗时) |  加工具但未使能 Dump (耗时)  |   加工具并使能 Dump (耗时)   | 加工具并使能 Md5 Dump (耗时) |
+|:--------:|:--------:|:-------------------:|:--------------------:|:--------------------:|
+| L0       | ≈95.1 ms |   ≈95.5 ms （无膨胀）    | ≈420.0 ms  （膨胀4.5倍）  |  ≈1011.3 ms  （膨胀10倍）  |
+| L1       | ≈95.1 ms  | ≈115.8 ms  （膨胀1.2倍） | ≈2469.0 ms  （膨胀26倍）  |  ≈8636.0 ms  （膨胀90倍）  |
+| mix      | ≈95.1 ms  | ≈117.8 ms  （膨胀1.2倍） | ≈3635.4 ms  （膨胀38 倍） | ≈10698.3 ms  （膨胀112倍） |
 
 
 ## "tensor"模式采集数据量参考基线

msprobe/mindspore/api_accuracy_checker/api_accuracy_checker.py

@@ -17,7 +17,7 @@ import os
 from dataclasses import dataclass
 from typing import Any, Optional
 from tqdm import tqdm
-import numpy as np
+
 from msprobe.core.common.const import Const, CompareConst
 from msprobe.core.common.file_utils import FileOpen, create_directory, write_csv, load_json, load_yaml
 from msprobe.core.common.utils import add_time_as_suffix

msprobe/mindspore/cell_processor.py

@@ -25,6 +25,7 @@ from msprobe.core.common.exceptions import MsprobeException
 from msprobe.core.common.runtime import Runtime
 from msprobe.core.common.utils import ModuleQueue, ThreadSafe
 from msprobe.core.data_dump.scope import ModuleRangeScope, MixRangeScope, BaseScope
+from msprobe.core.common.megatron_utils import wrap_megatron_step, get_micro_step, is_megatron
 from msprobe.mindspore.common.const import Const as MsConst
 from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.common.utils import (
@@ -47,6 +48,28 @@ def get_cell_construct(construct):
     return _construct
 
 
+def patch_schedules_step():
+    try:
+        from mindspeed.mindspore.core.pipeline_parallel import schedules
+        schedules.forward_step = wrap_megatron_step(schedules.forward_step)
+        schedules.backward_step = wrap_megatron_step(schedules.backward_step, is_forward=False)
+        logger.info_on_rank_0("Patch mindspeed.mindspore method success.")
+    except ImportError:
+        logger.info_on_rank_0("No mindspeed.mindspore find.")
+    except Exception as e:
+        logger.info_on_rank_0(f"Patch mindspeed.mindspore method failed, detail:{str(e)}")
+
+    try:
+        from megatron.core.pipeline_parallel import schedules
+        schedules.forward_step = wrap_megatron_step(schedules.forward_step)
+        schedules.backward_step = wrap_megatron_step(schedules.backward_step, is_forward=False)
+        logger.info_on_rank_0("Patch megatron method success.")
+    except ImportError:
+        logger.info_on_rank_0("No megatron find.")
+    except Exception as e:
+        logger.info_on_rank_0(f"Patch megatron method failed, detail:{str(e)}")
+
+
 class CellProcessor:
     cell_queue = ModuleQueue()
     cell_count = {}
@@ -84,6 +107,8 @@ class CellProcessor:
             raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
                                    'The model cannot be None, when level is "L0" or "mix"')
 
+        patch_schedules_step()
+
         is_registered = False
         model_type = Const.MODULE if is_mindtorch() else Const.CELL
         cells_with_index_in_pynative_mode, cells_with_index_in_graph_mode = get_cells_and_names_with_index(models)
@@ -127,6 +152,7 @@ class CellProcessor:
                 Runtime.run_mode = MsConst.PYNATIVE_GRAPH_MODE
                 GraphModeCellDump(config, cells_and_names_in_graph_mode, strict=False).handle()
 
+
     def build_cell_hook(self, cell_name, build_data_hook):
         @ThreadSafe.synchronized
         def forward_pre_hook(cell, args):
@@ -259,24 +285,26 @@ class CellProcessor:
             CellProcessor.cell_stack[tid] = []
 
         if self.cell_stack[tid]:
-            CellProcessor.module_node[full_name] = self.cell_stack[tid][-1]
+            CellProcessor.module_node[full_name] = self.cell_stack[tid][-1] if not is_megatron() \
+                else [self.cell_stack[tid][-1], get_micro_step()]
         else:
             parent_name = CellProcessor.cell_queue.find_last(full_name)
-            CellProcessor.module_node[full_name] = parent_name
+            CellProcessor.module_node[full_name] = parent_name if not is_megatron() else [parent_name, get_micro_step()]
 
         CellProcessor.cell_queue.add_name(full_name)
         CellProcessor.cell_stack[tid].append(full_name)
-        CellProcessor.api_parent_node[tid] = full_name
+        CellProcessor.api_parent_node[tid] = full_name if not is_megatron() else [full_name, get_micro_step()]
         if self.scope:
             self.scope.begin_module(full_name)
 
     def set_construct_info_in_hook(self, full_name):
         tid = threading.get_ident()
         CellProcessor.cell_queue.remove_name(full_name)
-        CellProcessor.api_parent_node[tid] = None
+        CellProcessor.api_parent_node[tid] = None if not is_megatron() else [None, get_micro_step()]
         if self.cell_stack.get(tid):
             CellProcessor.cell_stack[tid].pop()
         if self.cell_stack.get(tid):
-            CellProcessor.api_parent_node[tid] = CellProcessor.cell_stack[tid][-1]
+            CellProcessor.api_parent_node[tid] = CellProcessor.cell_stack[tid][-1] if not is_megatron() \
+                else [CellProcessor.cell_stack[tid][-1], get_micro_step()]
         if self.scope:
             self.scope.end_module(full_name)

msprobe/mindspore/compare/common_dir_compare.py

@@ -212,14 +212,14 @@ def do_multi_process(func, map_dict):
         df_chunks = [result_df]
         process_num = 1
     logger.info(f"Using {process_num} processes with chunk size {df_chunk_size}")
-    
+
     # 分割字典
     map_chunks = split_dict(map_dict, df_chunk_size)
-    
+
     # 创建结果列表和进程池
     results = []
     pool = multiprocessing.Pool(process_num)
-    
+
     progress_bar = tqdm(total=len(result_df), desc="API/Module Item Compare Process", unit="row", ncols=100)
 
     def update_progress(size, progress_lock, extra_param=None):
@@ -228,34 +228,30 @@ def do_multi_process(func, map_dict):
 
     def err_call(args):
         logger.error('multiprocess compare failed! Reason: {}'.format(args))
-        try:
-            pool.close()
-        except OSError as e:
-            logger.error(f'pool terminate failed: {str(e)}')
+
     results = []
+
+    # 提交任务到进程池
+    for process_idx, (df_chunk, map_chunk) in enumerate(zip(df_chunks, map_chunks)):
+        start_idx = df_chunk_size * process_idx
+        result = pool.apply_async(
+            func,
+            args=(df_chunk, start_idx, map_chunk, lock),
+            error_callback=err_call,
+            callback=partial(update_progress, len(map_chunk), lock)
+        )
+        results.append(result)
+    pool.close()
+
     try:
-        # 提交任务到进程池
-        for process_idx, (df_chunk, map_chunk) in enumerate(zip(df_chunks, map_chunks)):
-            start_idx = df_chunk_size * process_idx
-            result = pool.apply_async(
-                func,
-                args=(df_chunk, start_idx, map_chunk, lock),
-                error_callback=err_call,
-                callback=partial(update_progress, len(map_chunk), lock)
-            )
-            results.append(result)
-
-        final_results = [r.get() for r in results]
-        # 等待所有任务完成
-        pool.close()
-        pool.join()
-        return pd.concat(final_results, ignore_index=True)
+        final_results = [r.get(timeout=3600) for r in results]
     except Exception as e:
-        logger.error(f"\nMain process error: {str(e)}")
+        logger.error(f"Task failed with exception: {e}")
         pool.terminate()
         return pd.DataFrame({})
-    finally:
-        pool.close()
+    # 等待所有任务完成
+    pool.join()
+    return pd.concat(final_results, ignore_index=True)
 
 
 def initialize_result_df(total_size):

msprobe/mindspore/debugger/precision_debugger.py

@@ -182,7 +182,7 @@ class PrecisionDebugger(BasePrecisionDebugger):
             with ThreadSafe():
                 instance.service.step()
         if is_graph_mode_cell_dump_allowed(instance.config):
-            GraphModeCellDump.step()
+            GraphModeCellDump.step(instance.config.dump_path, instance.config.step, instance.config.task)
         if enable_dynamic_kbyk_dump and instance.config.level_ori == Const.LEVEL_L2:
             _dump_step(1)
         if cls._is_kernel_dump() and _msprobe_c:

msprobe/mindspore/dump/cell_dump_process.py

@@ -46,9 +46,11 @@ KEY_FORWARD = CoreConst.FORWARD
 KEY_BACKWARD = CoreConst.BACKWARD
 KEY_INPUT = CoreConst.INPUT
 KEY_OUTPUT = CoreConst.OUTPUT
-KEY_DUMP_TENSOR_DATA = "dump_tensor_data_"
+KEY_DUMP_TENSOR_DATA = "dump_tensor_data/"
 KEY_STATISTIC_CSV = "statistic.csv"
 KEY_TD_FLAG = "td_flag"
+# 设置落盘文件检测超时时间
+TIMEOUT = 600
 td = ops.TensorDump()
 if (ms.__version__ >= "2.5.0"):
     td_in = ops.TensorDump("in")
@@ -574,28 +576,33 @@ def generate_stack_info(path):
     logger.info(f"Stack data saved to {json_path}")
 
 
-def is_download_finished(directory, interval=3):
+def is_download_finished(directory, save_flag):
     """
     判断指定目录在一段时间后是否有数据被下载完成
     :param directory: 指定目录的路径
-    :param interval: 检查的时间间隔（秒），默认为 3 秒
+    :param save_flag: 数据落盘完成后的标志文件
     :return: 如有数据被下载完成返回 True，否则返回 False
     """
+    # 设定一定的延迟间隔，避免频繁进行磁盘的io读取操作
+    time.sleep(0.5)
+    logger.info("Waiting for download...")
     # 检查目录是否存在
     if not os.path.exists(directory):
         logger.warning(f"The specified directory {directory} does not exist.")
         return False
-    initial_modification_time = os.path.getmtime(directory)
-    time.sleep(interval)
-    current_modification_time = os.path.getmtime(directory)
-    # 比较初始和当前修改时间
-    if current_modification_time > initial_modification_time:
-        return False
-    else:
-        return True
+    
+    # 遍历当前目录中的所有条目
+    for entry_path in os.listdir(directory):
+        if entry_path.startswith(save_flag):
+            return True
+
+    return False
+
 
+def process_step(dump_path, flag_path, step, step_list):
+    if step not in step_list:
+        return
 
-def process(dump_path):
     if not os.path.exists(dump_path):
         logger.warning('No grap cell data is dumped.')
         create_directory(dump_path)
@@ -606,32 +613,38 @@ def process(dump_path):
     if rank_id is not None:
         rank_dir = CoreConst.RANK + str(rank_id)
 
-    step_dir_list = os.listdir(dump_path)
-    for step_dir in step_dir_list:
-        step_path = os.path.join(dump_path, step_dir)
-        rank_path = os.path.join(step_path, rank_dir)
-        npy_path = os.path.join(rank_path, CoreConst.DUMP_TENSOR_DATA)
-        while True:
-            is_finished = is_download_finished(npy_path)
-            if not is_finished:
-                logger.info("There is data being downloaded in the specified directory, continue checking...")
-            else:
-                logger.info("There is no data being downloaded in the specified directory, Stop checking.")
-                break
-        logger.info("==========Start processing data that has already been stored on the disk!==========")
-        rename_filename(path=npy_path)
-        generate_construct(npy_path)
-        generate_dump_info(npy_path)
-        generate_stack_info(npy_path)
-        # 单卡场景，rank目录名称为rank
-        if rank_id is None:
-            new_rank_path = os.path.join(step_path, CoreConst.RANK)
-            try:
-                move_directory(rank_path, new_rank_path)
-                logger.info(f"Directory was successfully renamed to: {new_rank_path}")
-            except Exception as e:
-                logger.warning(f"Failed to renamed to {new_rank_path}: {e}")
-        logger.info("==========JSON file generation completed!==========")
+    step_dir = CoreConst.STEP + str(step)
+
+    step_path = os.path.join(dump_path, step_dir)
+    rank_path = os.path.join(step_path, rank_dir)
+    npy_path = os.path.join(rank_path, CoreConst.DUMP_TENSOR_DATA)
+    save_finish_flag = f"step_{step}"
+    start_time = time.time() 
+    while True:
+        is_finished = is_download_finished(flag_path, save_finish_flag)
+        if not is_finished:
+            logger.info("There is data being downloaded in the specified directory, continue checking...")
+        else:
+            logger.info("There is no data being downloaded in the specified directory, Stop checking.")
+            break
+        elapsed_time = time.time() - start_time
+        if elapsed_time > TIMEOUT:
+            logger.error(f"Check timed out after {TIMEOUT} seconds. Exiting.")
+            return
+    logger.info(f"==========Start processing step_{step}'s data that has already been stored on the disk!==========")
+    rename_filename(path=npy_path)
+    generate_construct(npy_path)
+    generate_dump_info(npy_path)
+    generate_stack_info(npy_path)
+    # 单卡场景，rank目录名称为rank
+    if rank_id is None:
+        new_rank_path = os.path.join(step_path, CoreConst.RANK)
+        try:
+            move_directory(rank_path, new_rank_path)
+            logger.info(f"Directory was successfully renamed to: {new_rank_path}")
+        except Exception as e:
+            logger.warning(f"Failed to renamed to {new_rank_path}: {e}")
+    logger.info(f"==========Step_{step}'s JSON file generation completed!==========")
 
 
 # 删除csv文件中每行数据最后面的逗号
@@ -689,7 +702,10 @@ def merge_file(dump_path, rank_dir, file_dict):
                          " and the index is out of bounds.")
 
 
-def process_statistics(dump_path):
+def process_statistics_step(dump_path, step, step_list):
+    if step_list and step not in step_list:
+        return
+
     if not os.path.exists(dump_path):
         logger.warning('No grap cell data is dumped.')
         create_directory(dump_path)
@@ -723,25 +739,24 @@ def process_statistics(dump_path):
 
     rank_dir = rank_dir_kbk.replace(CoreConst.REPLACEMENT_CHARACTER, '')
     dir_list = os.listdir(dump_path)
-    step_dir_list = [d for d in dir_list if d.startswith(CoreConst.STEP)]
-    for step_dir in step_dir_list:
-        step_path = os.path.join(dump_path, step_dir)
-        rank_path = os.path.join(step_path, rank_dir)
-        csv_path = os.path.join(rank_path, KEY_STATISTIC_CSV)
-        logger.info("==========Start processing data csv!==========")
-        generate_construct(csv_path)
-        generate_dump_info(csv_path)
-        generate_stack_info(csv_path)
-        remove_path(rank_path_kbk)
-        # 单卡场景，rank目录名称为rank
-        if rank_id is None:
-            new_rank_path = os.path.join(step_path, CoreConst.RANK)
-            try:
-                move_directory(rank_path, new_rank_path)
-                logger.info(f"Directory was successfully renamed to: {new_rank_path}")
-            except Exception as e:
-                logger.warning(f"Failed to renamed to {new_rank_path}: {e}")
-        logger.info("==========JSON file generation completed!==========")
+    step_dir = CoreConst.STEP + str(step)
+    step_path = os.path.join(dump_path, step_dir)
+    rank_path = os.path.join(step_path, rank_dir)
+    csv_path = os.path.join(rank_path, KEY_STATISTIC_CSV)
+    logger.info("==========Start processing data csv!==========")
+    generate_construct(csv_path)
+    generate_dump_info(csv_path)
+    generate_stack_info(csv_path)
+    remove_path(rank_path_kbk)
+    # 单卡场景，rank目录名称为rank
+    if rank_id is None:
+        new_rank_path = os.path.join(step_path, CoreConst.RANK)
+        try:
+            move_directory(rank_path, new_rank_path)
+            logger.info(f"Directory was successfully renamed to: {new_rank_path}")
+        except Exception as e:
+            logger.warning(f"Failed to renamed to {new_rank_path}: {e}")
+    logger.info("==========JSON file generation completed!==========")
 
 
 def get_yaml_keys(yaml_data):
@@ -922,7 +937,3 @@ def start(config: CellDumpConfig):
             cell.data_mode = data_mode
 
     logger.info("==========The cell_dump_process_start phase is Finished!==========")
-    if dump_task == CoreConst.TENSOR:
-        atexit.register(process, dump_path=dump_path)
-    if dump_task == CoreConst.STATISTICS:
-        atexit.register(process_statistics, dump_path=dump_path)