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

msprobe/core/common/utils.py

@@ -708,3 +708,20 @@ def check_process_num(process_num):
         raise ValueError(f"process_num({process_num}) is not a positive integer")
     if process_num > Const.MAX_PROCESS_NUM:
         raise ValueError(f"The maximum supported process_num is {Const.MAX_PROCESS_NUM}, current value: {process_num}.")
+
+
+def confirm(prompt, default=False):
+    if default is True:
+        prompt_suffix = " [Y/n] "
+    elif default is False:
+        prompt_suffix = " [y/N] "
+    else:
+        prompt_suffix = " [y/n] "
+
+    full_prompt = prompt + prompt_suffix
+
+    user_input = input(full_prompt).strip().lower()
+    if user_input in ['y', 'yes']:
+        return True
+    else:
+        return default

msprobe/core/compare/diff_analyze/first_diff_analyze.py

@@ -26,8 +26,6 @@ from msprobe.core.compare.utils import gen_api_batches
 
 cur_dir = os.path.dirname(os.path.realpath(__file__))
 diff_threshold_yaml_path = os.path.join(cur_dir, 'diff_analyze_threshold.yaml')
-ignore_op_list_yaml_path = os.path.join(cur_dir, 'ignore_op_list.yaml')
-ignore_list = load_yaml(ignore_op_list_yaml_path)
 thresholds = load_yaml(diff_threshold_yaml_path)
 cmp_metrics = thresholds.get('compare_metrics')
 
@@ -53,7 +51,7 @@ class FirstDiffAnalyze:
                 return True
         return False
 
-    def single_api_check(self, result_slice, header, api_name=None):
+    def single_api_check(self, result_slice, header):
         """
         单个api差异检查
 
@@ -67,18 +65,14 @@ class FirstDiffAnalyze:
         }
 
         column_indices = {name: idx for idx, name in enumerate(header)}
-        output_idx = -1
+
         for line in result_slice:
             op_item = {
                 column_name: line[column_indices[column_name]]
                 for column_name in header
             }
             single_check_result['op_items'].append(op_item)
-            if op_item['state'] != 'output':
-                continue
-            output_idx += 1
-            if output_idx in ignore_list.get(api_name, []):
-                continue
+
             # set is_same
             if self.mode_config.dump_mode == Const.MD5:
                 if line[column_indices[CompareConst.RESULT]] == CompareConst.DIFF:
@@ -123,13 +117,7 @@ class FirstDiffAnalyze:
         with tqdm(total=len(api_batches), desc=bar_desc_add_rank, unit="api/module", ncols=100) as progress_bar:
             for api_batch in api_batches:
                 result_slice = result[api_batch.start: api_batch.params_grad_end_index]
-                api_compo = api_batch.api_name.split('.')
-                # suppose name is Tensor.MatMul.0.forward
-                if len(api_compo) < 4:
-                    continue
-                # get MatMul as api_name
-                api_name = api_compo[-3]
-                check_result[api_batch.api_name] = self.single_api_check(result_slice, header, api_name)
+                check_result[api_batch.api_name] = self.single_api_check(result_slice, header)
                 progress_bar.update(1)
 
         return check_result

msprobe/core/compare/find_first/utils.py

@@ -182,7 +182,7 @@ def analyze_diff_in_group(nodes_group):
     input_diff_nodes = list(filter(lambda node: node.is_diff, src_list))
     # 如果有异常回溯计算节点找到异常来源
     # 使用cpu模拟节点进行计算，查看结果是否有问题。需要对所有计算节点录入/映射，暂不实现。
-    get_compute_ops_from_comm_nodes(nodes_group)
+    get_compute_ops_from_comm_nodes(input_diff_nodes)
     # 筛选入参没问题但出参有问题的通信节点
     output_diff_nodes = list(filter(lambda node: node.data.is_diff, nodes_group))
     get_comm_ops(output_diff_nodes)

msprobe/core/config_check/ckpt_compare/ckpt_comparator.py

@@ -19,11 +19,11 @@ from tqdm import tqdm
 from msprobe.core.common.file_utils import save_json, check_path_before_create, check_path_not_exists, \
     check_file_or_directory_path
 from msprobe.core.common.log import logger
+from msprobe.core.common.utils import confirm
 from msprobe.core.config_check.ckpt_compare.megatron_loader import load_megatron_weights
 from msprobe.core.config_check.ckpt_compare.metrics import METRIC_FUNC
 
 
-
 def compare_checkpoints(ckpt_path1, ckpt_path2, output_path) -> Dict:
     """Compare weights between two checkpoints using cosine similarity and L2 distance.
     
@@ -45,6 +45,11 @@ def compare_checkpoints(ckpt_path1, ckpt_path2, output_path) -> Dict:
     """
 
     # Load both checkpoints
+    if not confirm("You are using torch.load with weights_only is False, it may cause arbitrary code "
+                   "execution. Do it only if you get the file from a trusted source. Input yes to continue, "
+                   "otherwise exit", False):
+        logger.error("Insecure risks found and exit!")
+        raise Exception("Insecure risks found and exit!")
     check_file_or_directory_path(ckpt_path1, isdir=True)
     check_file_or_directory_path(ckpt_path2, isdir=True)
     check_path_before_create(output_path)

msprobe/core/hook_manager.py

@@ -63,7 +63,6 @@ class BaseHookManager(ABC):
     def reset_status():
         BaseHookManager.inner_switch = defaultdict(bool)
         BaseHookManager.inner_api_count = defaultdict(int)
-        BaseHookManager.hook_handle_dict.clear()
         BaseHookManager.params_grad_info.clear()
 
     @staticmethod

msprobe/docs/01.installation.md

@@ -16,6 +16,8 @@ pip install mindstudio-probe
 
 |  版本   |    发布日期    |支持 PyTorch 版本|支持 MindSpore 版本|                                                                下载链接                                                                |校验码|
 |:-----:|:----------:|:--:|:--:|:----------------------------------------------------------------------------------------------------------------------------------:|:--:|
+| 8.3.0 | 2025.10.30 |1.11/2.0/2.1/2.2/2.5/2.6/2.7|2.4.0/2.5.0/2.6.0| [mindstudio_probe-8.3.0-py3-none-any.whl](https://ptdbg.obs.myhuaweicloud.com/msprobe/8.3/mindstudio_probe-8.3.0-py3-none-any.whl) |e933657b8ceb20774f924865d1f47978bd49cd1d9cec5fb60dec4f18802afbcc|
+| 8.2.1 | 2025.9.29  |1.11/2.0/2.1/2.2|2.4.0/2.5.0/2.6.0| [mindstudio_probe-8.2.1-py3-none-any.whl](https://ptdbg.obs.myhuaweicloud.com/msprobe/8.2/mindstudio_probe-8.2.1-py3-none-any.whl) |2152fadefec3d70148a910a54f2cfb6fccfa60b5ccd59f835f20289a7a5f4764|
 | 8.2.0 | 2025.9.03  |1.11/2.0/2.1/2.2|2.4.0/2.5.0/2.6.0| [mindstudio_probe-8.2.0-py3-none-any.whl](https://ptdbg.obs.myhuaweicloud.com/msprobe/8.2/mindstudio_probe-8.2.0-py3-none-any.whl) |bbc1577d76754adf987069308177d3e0a04e36de9c7f22e75c34cf4ad0ce1af2|
 | 8.1.2 | 2025.8.01  |1.11/2.0/2.1/2.2|2.4.0/2.5.0/2.6.0| [mindstudio_probe-8.1.2-py3-none-any.whl](https://ptdbg.obs.myhuaweicloud.com/msprobe/8.1/mindstudio_probe-8.1.2-py3-none-any.whl) |ff07bb81fddd3b8f3096d119ca1481bde8fdb24f10644def5250caad727448ab|
 | 8.1.1 | 2025.6.20  |1.11/2.0/2.1/2.2|2.4.0/2.5.0/2.6.0| [mindstudio_probe-8.1.1-py3-none-any.whl](https://ptdbg.obs.myhuaweicloud.com/msprobe/8.1/mindstudio_probe-8.1.1-py3-none-any.whl) |2aad10a243575544d7feef552caf4d06aa93028488ebd0bbc9aa350379da859d|

msprobe/docs/02.config_introduction.md

@@ -145,7 +145,7 @@ PyTorch、MSAdapter 以及 MindSpore 动态图场景下，"level"须为"L0"或"L
     <tr><td>pert_mode</td><td>无标杆扰动因子，str 类型。可选参数：<br/>  "improve_precision"：对输入做升精度，默认值；<br/>        "add_noise"：对输入增加噪声；<br/>        "no_change"：不加扰动直接二次执行；<br/>        "bit_noise"：输入的末位比特翻转，MindSpore 场景不支持 BF16 类型的向量；<br/>        "change_value"：输入的张量首尾值调换；<br/>        "to_cpu"：在 CPU 等价执行（仅 PyTorch 场景支持）。<br/><b>配置示例</b>："pert_mode": "improve_precision"。</td><td>否</td></tr>
     <tr><td>handler_type</td><td>处理类型，可选参数：<br/> "check"：进行无标杆比对检查，默认值；<br/> "fix"：将扰动后的 API 输出结果覆盖原始 API 输出结果，尝试将 Loss 曲线恢复正常，该模式下不支持预热功能与反向过程，且仅支持"improve_precision"、"to_cpu"（ PyTorch 场景）两种扰动因子。<br/> <b>配置示例</b>："handler_type": "check"。</td><td>否</td></tr>
     <tr><td>fuzz_level</td><td>无标杆数据 dump 级别，即选择比对结果文件应输出的表头属性，当前仅支持取值为："L1"。输出结果详见 <a href="#161-无标杆比对数据存盘格式">1.6.1 无标杆比对数据存盘格式</a>。</td><td>否</td></tr>
-    <tr><td>fuzz_stage</td><td>比对过程，选择对 API 前向或反向进行无标杆比对，可选参数：<br/> "forward"：前向，默认值；<br/>  "backward"：反向。当 fuzz_stage 为 "backward" 时，handler_type 只能为 "check"。<br/>  <b>配置示例</b>："fuzz_stage": "backward"。</td><td>否</td></tr>
+    <tr><td>fuzz_stage</td><td>比对过程，选择对 API 前向或反向进行无标杆比对，可选参数：<br/> "forward"：前向，默认值；<br/>  "backward"：反向。当 fuzz_stage 为 "backward" 时，handler_type 只能为 "check"。pytorch场景下，当 fuzz_stage 为 "backward" 时， list 参数不能为空。<br/>  <b>配置示例</b>："fuzz_stage": "backward"。</td><td>否</td></tr>
     <tr><td>if_preheat</td><td>预热功能（仅 PyTorch 场景支持），bool 类型。开启功能后工具可以根据每次迭代的输出调整精度算法的阈值，从而更准确地找出存在精度问题的 API。当"handler_type": "fix"时，不支持预热。可选参数：<br/>  true（开启）或 false（关闭），默认关闭。<br/>  <b>配置示例</b>："if_preheat": "true"。</td><td>否</td></tr>
     <tr><td>preheat_step</td><td>开启预热的迭代数量（仅 PyTorch 场景支持），int 类型，默认值为 15。须配置 "if_preheat": "true"。</td><td>否</td></tr>
     <tr><td>max_sample</td><td>每个算子预热的采样次数的最大阈值（仅 PyTorch 场景支持），int 类型，默认值为 20。须配置 "if_preheat": "true"。</td><td>否</td></tr>

msprobe/docs/14.data_parse_PyTorch.md

@@ -8,6 +8,8 @@
 
 依赖：CANN 包中的 msaccucmp 工具，需要安装 Ascend-CANN-toolkit，详见《[CANN 软件安装指南](https://www.hiascend.com/document/detail/zh/canncommercial/700/envdeployment/instg/instg_0001.html)》。
 
+**安全注意事项**： 工具使用时传入的msaccucmp.py文件路径参数请务必确保为CANN中自带的msaccucmp.py原始路径且文件内容未被篡改，工具会在解析数据时直接执行该文件，用户需自行保证该文件内容安全性。
+
 ## 2 数据解析操作指导
 
 ### 2.1 进入 parse 交互式界面

msprobe/docs/15.free_benchmarking_PyTorch.md

@@ -87,7 +87,7 @@ D-->config.json配置
     <tr><td>scope</td><td>否</td><td>自定义</td><td>需要通过指定算子名来限制算子插桩范围 如：["Torch.matmul.0.forward", "Tensor.pow.4.forward"]。</td></tr>
     <tr><td>list</td><td>否</td><td>自定义</td><td>需要通过指定算子类型来限制算子插桩范围 如：["relu"] 会匹配所有算子名中包含relu的算子。</td></tr>
     <tr><td rowspan="2">fuzz_stage</td><td rowspan="2">否</td><td>"forward"（默认）</td><td>需要进行算子<b>前向</b>计算的精度问题排查或<b>验证可疑算子。</b></td></tr>
-    <tr><td>"backward"</td><td>需要进行算子<b>反向</b>计算的精度问题排查，不支持仅反向验证，前向验证包括反向。</td><td></td></tr>
+    <tr><td>"backward"</td><td>需要进行算子<b>反向</b>计算的精度问题排查，不支持仅反向验证，前向验证包括反向。必须设置list参数指定需要检测的算子（详见3.2 config.json配置章节）, 指定的算子会暂存前向激活值，增加内存的占用。</td><td></td></tr>
 </table>
 
 #### 3.2.2 选择扰动因子

msprobe/docs/21.visualization_PyTorch.md

@@ -2,7 +2,7 @@
 
 分级可视化工具将msprobe工具dump的精度数据进行解析，还原模型图结构，实现模型各个层级的精度数据比对，方便用户理解模型结构、分析精度问题。
 
-工具支持PyTorch版本：2.1/2.2
+工具支持PyTorch版本：2.1/2.2/2.5/2.6/2.7
 
 ## 工具特性
 

msprobe/docs/26.data_dump_PyTorch_baseline.md

@@ -6,9 +6,9 @@
 
 | 采集模式 | 无工具 (耗时) |  加工具但未使能 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倍） |
+| L0       | ≈95.1 ms |   ≈95.5 ms （无膨胀）    | ≈420.0 ms  （膨胀4.5倍）  |  ≈1011.3 s  （膨胀10倍）  |
+| L1       | ≈95.1 ms  | ≈115.8 ms  （膨胀1.2倍） | ≈2469.0 ms  （膨胀26倍）  |  ≈8636.0 s  （膨胀90倍）  |
+| mix      | ≈95.1 ms  | ≈117.8 ms  （膨胀1.2倍） | ≈3635.4 ms  （膨胀38 倍） | ≈10698.3 s  （膨胀112倍） |
 
 
 ## "tensor"模式采集数据量参考基线

msprobe/docs/32.ckpt_compare.md

@@ -19,11 +19,11 @@ Megatron、MindSpeed的ckpt加载依赖megatron，请确保megatron在python环
 msprobe --framework pytorch config_check --compare path1 path2 -o output_path.json
 ```
 
-| 参数名 | 解释 | 是否必选 |
-|--------|-------|--------|
-| -f 或 --framework | 深度学习框架，str类型。支持参数：pytorch,mindspore，注意：msadaptor场景传入mindspore。 | 是 |
-| -c 或 --compare | 2个ckpt的路径 | 是 |
-| -o 或 --output | 比对结果输出路径，默认为 ./ckpt_similarity.json。输出路径存在时将报错终止。 | 否 |
+| 参数名 | 解释                                                                                           | 是否必选 |
+|--------|----------------------------------------------------------------------------------------------|--------|
+| -f 或 --framework | 深度学习框架，str类型。支持参数：pytorch,mindspore，注意：msadaptor场景传入mindspore。                               | 是 |
+| -c 或 --compare | 2个ckpt的路径。在ckpt传给工具加载前，用户需要确保ckpt是安全可信的，若ckpt来源官方有提供SHA256等校验值，用户必须要进行校验，以确保ckpt没有被篡改。 | 是 |
+| -o 或 --output | 比对结果输出路径，默认为 ./ckpt_similarity.json。输出路径存在时将报错终止。                                            | 否 |
 
 Megatron-LM 和 MindSpeed 的 ckpt 目录结构如下：
 

msprobe/mindspore/api_accuracy_checker/generate_op_script/op_generator.py

@@ -45,7 +45,7 @@ API_INFO = 2
 FOUR_SEGMENT = 4
 FIVE_SEGMENT = 5
 DATA_NAME = "data_name"
-API_MAX_LENGTH = 30
+API_MAX_LENGTH = 300
 PROPAGATION_LIST = [Const.FORWARD, Const.BACKWARD]
 DATAMODE_LIST = ["random_data", "real_data"]
 ITER_MAX_TIMES = 1000

msprobe/mindspore/compare/utils.py

@@ -26,8 +26,7 @@ def read_npy_data(dir_path, file_name):
         return None
 
     data_path = os.path.join(dir_path, file_name)
-    path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE,
-                               FileCheckConst.NUMPY_SUFFIX, False)
+    path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE, FileCheckConst.NUMPY_SUFFIX)
     data_path = path_checker.common_check()
     data_value = load_npy(data_path)
     return data_value

msprobe/mindspore/monitor/module_hook.py

@@ -163,15 +163,11 @@ class GradContext:
     def __init__(self) -> None:
         self.pre = {}
         self.post = {}
-        self.acc_metric = {}
-        self.acc = {}
         self.actv = {}
 
     def reset(self):
         self.pre.clear()
         self.post.clear()
-        self.acc_metric.clear()
-        self.acc.clear()
         self.actv.clear()
 
 
@@ -312,7 +308,6 @@ class TrainerMon:
         self.recording_l2_features = self.config.get('recording_l2_features', False)
         self.sa_order = self.config.get('sa_order', "s,b,h,d")
 
-
         if not self.cc_distribution.get('enable', False):
             self.cc_log_only = False
         else:
@@ -403,7 +398,7 @@ class TrainerMon:
             if self.monitoring:
                 module_rank_valid = self.is_target_rank()
                 step_condition = (context.step >= self.start_step and (
-                            context.step - self.start_step) % self.step_interval == 0)
+                        context.step - self.start_step) % self.step_interval == 0)
                 if module_rank_valid and step_condition:
                     self.has_collect_times += 1
 
@@ -447,6 +442,7 @@ class TrainerMon:
                     hook(optimizer, args, kwargs)
                 step_final_hook(optimizer, args, kwargs)
                 return out
+
             return wrapper
 
         if self.is_mindtorch:
@@ -541,7 +537,7 @@ class TrainerMon:
             if self.mv_distribution or self.ur_distribution or self.mg_direction:
                 if self.is_mindtorch:
                     context.param_exp_avg, context.param_exp_avg_sq, context.param_adam_update, \
-                    context.param_adam_ratio = self.optimizer_mon.fetch_mv(self, self.param2name)
+                        context.param_adam_ratio = self.optimizer_mon.fetch_mv(self, self.param2name)
                 else:
                     context.param_exp_avg, context.param_exp_avg_sq = self.get_mv_for_ms(optimizer)
 
@@ -564,7 +560,6 @@ class TrainerMon:
             context = self.optimizer_context[optimizer]
             self.generate_param_metrics(context, MonitorConst.POST_PARAM)
 
-
         if self.optimizer_hooked or not self.is_target_rank():
             return
 
@@ -577,7 +572,6 @@ class TrainerMon:
         if not self.wg_distribution:
             return
 
-        get_metrics(self.ops, self.grad_context.acc, self.eps, self.grad_context.acc_metric)
         get_metrics(self.ops, grad_dict, self.eps, self.grad_context.post)
 
     def generate_param_map(self, tag, param_tensor):
@@ -671,7 +665,7 @@ class TrainerMon:
         if not self.wg_distribution:
             return
 
-        self.summary_writer.write_metrics(self.ops, self.grad_context.acc_metric, step, 'grad_unreduced',
+        self.summary_writer.write_metrics(self.ops, self.grad_context.pre, step, 'grad_unreduced',
                                           use_micro_step=self.monitor_mbs_grad)
         self.summary_writer.write_metrics(self.ops, self.grad_context.post, step, 'grad_reduced')
 
@@ -810,7 +804,7 @@ class TrainerMon:
                     f'{context.module_name}.{Const.INPUT}', f'{MonitorConst.NAME_SEP}{context.micro_step}',
                     MonitorConst.ACTV, module_input))
             module_output = [tensor for tensor in module_output if isinstance(tensor, Tensor)] \
-                            if isinstance(module_output, tuple) else module_output
+                if isinstance(module_output, tuple) else module_output
             tbtag_tensor_map.update(
                 self.build_tbtag_tensor_map(
                     f'{context.module_name}.{Const.OUTPUT}', f'{MonitorConst.NAME_SEP}{context.micro_step}',
@@ -868,11 +862,13 @@ class TrainerMon:
             if version.parse(mindspore.__version__) >= version.parse('2.6.0'):
                 def wrapper(module, args, kwargs, module_output):
                     return fwd_hook_fun(module, args, kwargs, module_output, name)
+
                 return module.register_forward_hook(wrapper, with_kwargs=True)
 
             else:
                 def wrapper(module, args, module_output):
                     return fwd_hook_fun(module, args, None, module_output, name)
+
                 return module.register_forward_hook(wrapper)
 
         def extract_attention_feature_hook(module, args, kwargs, module_output, name):
@@ -880,7 +876,7 @@ class TrainerMon:
             if kwargs:
                 kwargs_tensors = [tensor for tensor in kwargs.values() if isinstance(tensor, Tensor)]
                 module_input.extend(kwargs_tensors)
-            
+
             if module not in self.feature_hook_context_by_module:
                 self.feature_hook_context_by_module[module] = FeatureHookContext(name)
             context: FeatureHookContext = self.feature_hook_context_by_module[module]
@@ -890,7 +886,7 @@ class TrainerMon:
                 logger.warning(
                     "Calculate attention feature failed, the length of module_input in attention hook's module should "
                     "be greater than or equal to 2.")
-            
+
             q_h = module_input[0]
             k_h = module_input[1]
             qkt = cal_qkt(q_h, k_h, order=self.sa_order)
@@ -985,25 +981,26 @@ class TrainerMon:
         self._hook_weights()
 
     def _hook_weights(self):
-        context = self.grad_context
 
         @_no_grad()
-        def param_hook(grad, context_dict, param, name):
+        def param_hook(grad, param, name):
             key = name
             if self.monitor_mbs_grad:
                 key += f'{MonitorConst.NAME_SEP}{param.micro_step}'
             key = get_summary_writer_tag_name(key, 'acc_grad', self.rank)
             self.register_param_call_id("param_hook", key)
             param.micro_step += 1
-
+            grad_dict = {}
             if self.monitor_mbs_grad or (param.micro_step == self.micro_batch_number):
-                context_dict[key] = grad
+                grad_dict[key] = grad
+            get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
+
             if param.micro_step == self.micro_batch_number:
                 param.micro_step = 0
 
-        def param_hook_wrapper(param_hook, context_dict, param, name):
+        def param_hook_wrapper(param_hook, param, name):
             def wrapper(grad):
-                return param_hook(grad, context_dict, param, name)
+                return param_hook(grad, param, name)
 
             return wrapper
 
@@ -1011,7 +1008,7 @@ class TrainerMon:
         for param, name in self.param2name.items():
             setattr(param, 'micro_step', 0)
             handle = param.register_hook(
-                param_hook_wrapper(param_hook, context_dict=context.acc, param=param, name=name))
+                param_hook_wrapper(param_hook, param=param, name=name))
             self.handles['wgrads'].append(handle)
         self.weight_hooked = True
 

msprobe/msprobe.py

@@ -14,16 +14,16 @@
 # limitations under the License.
 
 import argparse
-import sys
 import importlib.util
+import sys
 
 from msprobe.core.common.const import Const
+from msprobe.core.common.file_utils import root_privilege_warning
 from msprobe.core.common.log import logger
-from msprobe.core.compare.utils import _compare_parser
 from msprobe.core.compare.compare_cli import compare_cli
 from msprobe.core.compare.merge_result.merge_result_cli import _merge_result_parser, merge_result_cli
-from msprobe.core.config_check.config_check_cli import _config_checking_parser, \
-            _run_config_checking_command
+from msprobe.core.compare.utils import _compare_parser
+from msprobe.core.config_check.config_check_cli import _config_checking_parser, _run_config_checking_command
 
 
 def is_module_available(module_name):
@@ -64,6 +64,8 @@ def main():
     if len(sys.argv) < 4:
         parser.print_help()
         sys.exit(0)
+
+    root_privilege_warning()
     framework_args = parser.parse_args(sys.argv[1:3])
     if framework_args.framework == Const.PT_FRAMEWORK:
         from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut import _run_ut_parser, run_ut_command

msprobe/pytorch/api_accuracy_checker/generate_op_script/op_generator.py

@@ -47,7 +47,7 @@ API_INFO = 2
 FOUR_SEGMENT = 4
 FIVE_SEGMENT = 5
 DATA_NAME = "data_name"
-API_MAX_LENGTH = 30
+API_MAX_LENGTH = 300
 PROPAGATION_LIST = [Const.FORWARD, Const.BACKWARD]
 DATAMODE_LIST = ["random_data", "real_data"]
 ITER_MAX_TIMES = 1000

msprobe/pytorch/api_accuracy_checker/run_ut/multi_run_ut.py

@@ -39,7 +39,12 @@ from msprobe.core.common.const import FileCheckConst, Const
 from msprobe.core.common.utils import CompareException
 
 
-def split_json_file(input_file, num_splits, filter_api):
+def split_json_file(input_file, num_splits, filter_api, device_id):
+    max_processes = len(device_id) * 8
+    if num_splits > max_processes:
+        logger.warning(f"A device supports a maximum of 8 processes. "
+                       f"The total number of processes exceeds the limit, and it is set to {max_processes}.")
+        num_splits = max_processes
     forward_data, backward_data, real_data_path = parse_json_info_forward_backward(input_file)
     input_dir = os.path.dirname(os.path.abspath(input_file))
     if filter_api:
@@ -88,7 +93,7 @@ def split_json_file(input_file, num_splits, filter_api):
                     logger.error(f"File not found or could not be deleted: {file}")
             msg = 'ERROR: Split json file failed, please check the input file and try again.'
             raise CompareException(CompareException.PARSE_FILE_ERROR, msg) from e
-    return split_files, total_items
+    return split_files, total_items, num_splits
 
 
 def signal_handler(signum, frame):
@@ -127,7 +132,8 @@ def run_parallel_ut(config):
     def read_process_output(process):
         try:
             while True:
-                if process.poll() is not None:
+                # 子进程标准输出流与进程本身状态是分开的，因此增加判断。子进程返回值非None表示子进程结束，标准输出为None表示结束。
+                if process.poll() is not None or process.stdout is None:
                     break
                 output = process.stdout.readline()
                 if output == '':
@@ -175,12 +181,17 @@ def run_parallel_ut(config):
 
     try:
         for process in processes:
-            process.communicate(timeout=None)
+            process.wait()  # wait仅阻塞，不捕获标准输出和标准错误，原communicate不仅阻塞，而且捕获标准输出和标准错误
     except KeyboardInterrupt:
         logger.warning("Interrupted by user, terminating processes and cleaning up...")
     except Exception as e:
         logger.error(f"An unexpected error occurred: {e}")
     finally:
+        # 最后再更新一次进度条，避免因缓存写入等原因子进程结束而进度未刷新的问题
+        if wait_for_file_write_complete(config.result_csv_path):
+            result_file = read_csv(config.result_csv_path)
+            completed_items = len(result_file)
+            progress_bar.update(completed_items - progress_bar.n)
         if progress_bar.n < config.total_items:
             logger.warning("The UT task has not been completed. The parameter '-csv_path' along with the path to " \
                            "the result CSV file will be utilized to resume the UT task.")
@@ -195,6 +206,22 @@ def run_parallel_ut(config):
         logger.error(f"An unexpected error occurred: {e}")
 
 
+def wait_for_file_write_complete(file_path, timeout=3600):
+    last_size = 0
+    start_time = time.time()  # 记录开始时间
+    while True:
+        current_size = os.path.getsize(file_path)
+        # 检查是否文件大小未变化
+        if current_size == last_size:
+            return True  # 文件写入完成，返回 True
+        last_size = current_size
+        # 检查是否超时
+        if time.time() - start_time > timeout:
+            logger.error("write the result csv file timeout.")
+            return False  # 超时，返回 False
+        time.sleep(0.1)  # 适当的延时
+
+
 def prepare_config(args):
     api_info_file_checker = FileChecker(file_path=args.api_info_file, path_type=FileCheckConst.FILE,
                                         ability=FileCheckConst.READ_ABLE, file_type=FileCheckConst.JSON_SUFFIX)
@@ -203,7 +230,9 @@ def prepare_config(args):
     create_directory(out_path)
     out_path_checker = FileChecker(out_path, FileCheckConst.DIR, ability=FileCheckConst.WRITE_ABLE)
     out_path = out_path_checker.common_check()
-    split_files, total_items = split_json_file(api_info, args.num_splits, args.filter_api)
+    split_files, total_items, modified_num_splits = split_json_file(api_info, args.num_splits,
+                                                                    args.filter_api, args.device_id)
+    args.num_splits = modified_num_splits
     config_path = args.config_path if args.config_path else None
     if config_path:
         config_path_checker = FileChecker(config_path, FileCheckConst.FILE,

msprobe/pytorch/common/utils.py

@@ -39,7 +39,6 @@ except ImportError:
 else:
     is_gpu = False
 
-
 torch_without_guard_version = torch.__version__ >= '2.1'
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
@@ -338,56 +337,6 @@ def save_pt(tensor, filepath):
     change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
-class TypeCheckingUnpickler(pickle.Unpickler):
-    """
-    This class is a subclass of pickle.Unpickler, which is used to unpickle pickled objects.
-    It overrides the find_class method to add type checking functionality.
-    """
-    allowed_types = [
-        "str",
-        "ApiData",
-        "OrderedDict",
-        "_rebuild_tensor_v2",  # from torch.utils
-        "_load_from_bytes"  # from torch.storage
-    ]
-
-    def find_class(self, module, name):
-        """
-        Method to find the class of the object to be unpickled.
-        Throws pickle.UnpicklingError If the object type is not in the allowed types list.
-        """
-        if name in self.allowed_types:
-            return super().find_class(module, name)
-        raise pickle.UnpicklingError("Unsupported object type: {}.{}".format(module, name))
-
-
-def save_pkl(tensor, filepath):
-    """Save ApiData or str objection by pickle"""
-    check_path_before_create(filepath)
-    filepath = os.path.realpath(filepath)
-    try:
-        with FileOpen(filepath, 'wb') as f:
-            pickle.dump(tensor, f)
-    except Exception as e:
-        logger.error("Save pt file failed, please check according possible error causes: "
-                     "1. out of disk space or disk error, "
-                     "2. no permission to write files, etc.")
-        raise RuntimeError(f"save pt file {filepath} failed") from e
-    change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
-
-
-def load_pkl(pt_path):
-    """Load ApiData or str objection by pickle for accuracy_checker_online"""
-    check_file_or_directory_path(pt_path)
-    pt_path = os.path.realpath(pt_path)
-    try:
-        with FileOpen(pt_path, 'rb') as f:
-            pt = TypeCheckingUnpickler(f).load()
-    except Exception as e:
-        raise RuntimeError(f"load pt file {pt_path} failed: {e}") from e
-    return pt
-
-
 def is_recomputation():
     """Check if the current operation is in the re-computation phase.
 
@@ -416,7 +365,8 @@ def is_recomputation():
 
     # Identify indices in the call stack where the specific function is being executed
     for idx, frame_info in enumerate(call_stack):
-        if frame_info.function == Const.BACKWARD or frame_info.function == 'checkpoint_function_backward':
+        if (frame_info.function == Const.BACKWARD or frame_info.function == 'checkpoint_function_backward' and
+                "megatron" in frame_info.filename):
             backward_function_indices.append(idx)
 
     # Check if the execution is within 'torch/autograd/function.py' file

msprobe/pytorch/compare/utils.py

@@ -27,8 +27,7 @@ def read_pt_data(dir_path, file_name):
         return None
 
     data_path = os.path.join(dir_path, file_name)
-    path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE,
-                               FileCheckConst.PT_SUFFIX, False)
+    path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE, FileCheckConst.PT_SUFFIX)
     data_path = path_checker.common_check()
     try:
         # detach because numpy can not process gradient information

msprobe/pytorch/dump/module_dump/hook_wrapper.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 from functools import wraps
+from typing import Any, Callable
 
 import torch
 from torch.utils.hooks import BackwardHook
@@ -21,6 +22,9 @@ from torch.utils.hooks import BackwardHook
 from msprobe.core.common.const import Const
 from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.hook_module.api_register import get_api_register
+
+torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
 
 def wrap_setup_backward_hook(func):
@@ -92,3 +96,23 @@ def wrap_setup_backward_hook(func):
 def wrap_setup_input_output_hook():
     BackwardHook.setup_input_hook = wrap_setup_backward_hook(BackwardHook.setup_input_hook)
     BackwardHook.setup_output_hook = wrap_setup_backward_hook(BackwardHook.setup_output_hook)
+
+
+def get_apply_func_wrapper(original_func: Callable) -> Callable:
+    @wraps(original_func)
+    def wrapped_apply(*args, **kwargs) -> Any:
+        api_register = get_api_register()
+        if api_register:
+            api_register.restore_inner_used_api()
+        result = original_func(*args, **kwargs)
+        if api_register:
+            api_register.register_inner_used_api()
+        return result
+
+    return wrapped_apply
+
+
+def wrap_backward_hook_function_apply():
+    if torch_version_above_or_equal_2:
+        original_apply = torch.nn.modules._functions.BackwardHookFunction.apply
+        torch.nn.modules._functions.BackwardHookFunction.apply = get_apply_func_wrapper(original_apply)