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

msprobe/core/compare/utils.py

@@ -18,13 +18,14 @@ import re
 import math
 import zlib
 from dataclasses import dataclass
+import multiprocessing
 
 import numpy as np
 import pandas as pd
 
 from msprobe.core.common.const import Const, CompareConst, FileCheckConst
 from msprobe.core.common.utils import CompareException, check_regex_prefix_format_valid, logger, safe_get_value
-from msprobe.core.common.file_utils import check_file_or_directory_path
+from msprobe.core.common.file_utils import check_file_or_directory_path, load_json
 
 json_file_mapping = {
     Const.DUMP_JSON_FILE: "dump.json",
@@ -94,30 +95,39 @@ def check_and_return_dir_contents(dump_dir, prefix):
 
 
 def read_op(op_data, op_name):
+    if not isinstance(op_name, str):
+        logger.error(f"api name error: {op_name} is not a string, please check.")
+        raise CompareException(CompareException.INVALID_API_NAME_ERROR)
     split_name = op_name.split(Const.SEP)
-    if Const.DEBUG in split_name or Const.PARAMS_GRAD in split_name:
-        op_parsed_list = op_item_parse(op_data, op_name)
+    if split_name[-1] == Const.DEBUG:
+        op_parsed_list = op_item_parse(op_data, op_name, Const.DEBUG)
+    elif split_name[-1] == Const.PARAMS_GRAD:
+        op_parsed_list = op_item_parse(op_data, op_name, Const.PARAMS_GRAD)
     else:
         op_parsed_list = []
         for name in CompareConst.IO_NAME_MAPPING:
             if name in op_data:
-                op_parsed_list.extend(op_item_parse(op_data[name], op_name + CompareConst.IO_NAME_MAPPING[name]))
+                op_parsed_list.extend(op_item_parse(op_data[name], op_name + CompareConst.IO_NAME_MAPPING[name], name))
     return op_parsed_list
 
 
-def op_item_parse(op_data, op_name: str, depth: int = 0) -> list:
+def op_item_parse(op_data, op_name: str, state: str, depth: int = 0) -> list:
+    if state == Const.INPUT_ARGS or state == Const.INPUT_KWARGS:
+        state = Const.INPUT
     default_item = {
         'full_op_name': op_name,
-        'type': None,
-        'Max': None,
-        'Min': None,
-        'Mean': None,
-        'Norm': None,
-        'dtype': None,
-        'shape': None,
-        'md5': None,
-        'value': None,
-        'data_name': '-1'
+        Const.TYPE: None,
+        Const.MAX: None,
+        Const.MIN: None,
+        Const.MEAN: None,
+        Const.NORM: None,
+        Const.DTYPE: None,
+        Const.SHAPE: None,
+        Const.MD5: None,
+        Const.VALUE: None,
+        Const.DATA_NAME: '-1',
+        Const.STATE: state,
+        Const.REQ_GRAD: None
     }
 
     if depth > Const.MAX_DEPTH:
@@ -133,33 +143,53 @@ def op_item_parse(op_data, op_name: str, depth: int = 0) -> list:
     if isinstance(op_data, list):
         for i, data in enumerate(op_data):
             if Const.PARAMS_GRAD not in op_name.split(Const.SEP):
-                item_list.extend(op_item_parse(data, op_name + Const.SEP + str(i), depth + 1))
+                item_list.extend(op_item_parse(data, op_name + Const.SEP + str(i), state, depth + 1))
             else:
-                item_list.extend(op_item_parse(data, op_name, depth + 1))
+                item_list.extend(op_item_parse(data, op_name, state, depth + 1))
     elif isinstance(op_data, dict):
+        if is_p2pop_leaf_data(op_data):
+            p2pop_item = {}
+            for key in ['class_type', 'op', 'peer', 'tag', 'group_id']:
+                p2pop_item[key] = op_data.get(key)
+            op_data = op_data.get('tensor')
+            if isinstance(op_data, dict):
+                op_item = gen_op_item(op_data, op_name, state)
+            else:
+                op_item = default_item
+            op_item.update(p2pop_item)
+            return [op_item]
         if is_leaf_data(op_data):
-            return [gen_op_item(op_data, op_name)]
+            return [gen_op_item(op_data, op_name, state)]
         for sub_name, sub_data in op_data.items():
-            item_list.extend(op_item_parse(sub_data, op_name + Const.SEP + str(sub_name), depth + 1))
+            item_list.extend(op_item_parse(sub_data, op_name + Const.SEP + str(sub_name), state, depth + 1))
     return item_list
 
 
+def is_p2pop_leaf_data(op_data):
+    return op_data.get('class_type') == 'torch.distributed.P2POp'
+
+
 def is_leaf_data(op_data):
     return 'type' in op_data and isinstance(op_data['type'], str)
 
 
-def gen_op_item(op_data, op_name):
+def gen_op_item(op_data, op_name, state):
     op_item = {}
-    op_item.update(op_data)
-    data_name = op_data.get('data_name') if op_data.get('data_name') else '-1'  # 如果是""也返回-1
-    op_item['data_name'] = data_name
+    op_item.update({key: str(value) if isinstance(value, bool) else value for key, value in op_data.items()})
+    data_name = op_data.get(Const.DATA_NAME) if op_data.get(Const.DATA_NAME) else '-1'  # 如果是""也返回-1
+    op_item[Const.DATA_NAME] = data_name
     op_item['full_op_name'] = data_name.rsplit(Const.SEP, 1)[0] if data_name != '-1' else op_name
+    op_item[Const.STATE] = state
+    if Const.REQ_GRAD not in op_item:
+        op_item[Const.REQ_GRAD] = None
 
-    params = ['Max', 'Min', 'Mean', 'Norm']
+    # 补齐统计量字段
+    params = [Const.MAX, Const.MIN, Const.MEAN, Const.NORM]
     for i in params:
         if i not in op_item:
             op_item[i] = None
 
+    # special cases
     if not op_item.get('dtype'):
         if op_item.get('type') == 'torch.Size':
             op_item['dtype'] = op_data.get('type')
@@ -172,11 +202,18 @@ def gen_op_item(op_data, op_name):
             op_item['shape'] = '[]'
             for i in params:
                 op_item[i] = op_data.get('value')
+        elif op_name.split(Const.SEP)[-1] in ['src', 'dst', 'group_src', 'group_dst']:
+            op_item['dtype'] = op_data.get('type')
+            op_item['shape'] = '[]'
+            for i in params:
+                op_item[i] = str(op_data.get('value'))
+            op_item['md5'] = str(op_data.get('value'))
         elif op_item.get('type') == 'torch.ProcessGroup':
             op_item['dtype'] = op_data.get('type')
             op_item['shape'] = '[]'
             for i in params:
                 op_item[i] = str(op_data.get('group_ranks'))
+            op_item['md5'] = str(op_data.get('group_ranks'))
         else:
             op_item['dtype'] = str(type(op_data.get('value')))
             op_item['shape'] = '[]'
@@ -205,22 +242,26 @@ def merge_tensor(tensor_list, dump_mode):
         CompareConst.PARAMS_GRAD_STRUCT,
         CompareConst.DEBUG_STRUCT,
         Const.SUMMARY,
-        Const.STACK_INFO
+        Const.STACK_INFO,
+        Const.STATE,
+        Const.REQ_GRAD
     ]
     op_dict = {key: [] for key in keys}
 
     if dump_mode == Const.ALL:
-        op_dict["data_name"] = []
+        op_dict[Const.DATA_NAME] = []
 
     for tensor in tensor_list:
         # A dict(len=2) with 'full_op_name' and 'full_info' is added to the tensor only if self.stack_mode is True
         if len(tensor) == 2:
-            op_dict[Const.STACK_INFO].append(tensor['full_info'])
+            op_dict[Const.STACK_INFO].append(tensor.get('full_info'))
             break
 
-        op_dict[CompareConst.OP_NAME].append(tensor['full_op_name'])
+        op_dict[CompareConst.OP_NAME].append(tensor.get('full_op_name'))
+        state = tensor.get(Const.STATE)
+        op_dict[Const.STATE].append(state)
+        op_dict[Const.REQ_GRAD].append(tensor.get(Const.REQ_GRAD))
 
-        _, state = get_name_and_state(tensor['full_op_name'])
         struct_key = CompareConst.STATE_TO_STRUCT_MAPPING.get(state)
         if not struct_key:
             continue
@@ -228,22 +269,19 @@ def merge_tensor(tensor_list, dump_mode):
             op_dict.get(struct_key).append((tensor[Const.DTYPE], tensor[Const.SHAPE], tensor[Const.MD5]))
         else:
             op_dict.get(struct_key).append((tensor[Const.DTYPE], tensor[Const.SHAPE]))
-        op_dict[Const.SUMMARY].append([tensor[Const.MAX], tensor[Const.MIN], tensor[Const.MEAN], tensor[Const.NORM]])
+
+        # 当统计量为None时，转成字符串None，避免后续操作list放到pd中时None被默认转成NaN
+        op_dict[Const.SUMMARY].append(
+            [str(tensor[key]) if tensor[key] is None else tensor[key] for key in Const.SUMMARY_METRICS_LIST])
 
         if dump_mode == Const.ALL:
-            op_dict["data_name"].append(tensor['data_name'])
+            op_dict[Const.DATA_NAME].append(tensor.get(Const.DATA_NAME))
 
     if not op_dict[CompareConst.KWARGS_STRUCT]:
         del op_dict[CompareConst.KWARGS_STRUCT]
     return op_dict if op_dict[CompareConst.OP_NAME] else {}
 
 
-def check_api_info_len(op_name, info_list, len_require):
-    if len(info_list) < len_require:
-        logger.error(f'Index out of bounds error, please check info of api: {op_name}.')
-        raise CompareException(CompareException.INDEX_OUT_OF_BOUNDS_ERROR)
-
-
 def print_compare_ends_info():
     total_len = len(CompareConst.COMPARE_ENDS_SUCCESSFULLY) + Const.FILL_CHAR_NUMS
     logger.info('*' * total_len)
@@ -263,83 +301,113 @@ def table_value_is_valid(value: str) -> bool:
     return True
 
 
-def get_name_and_state(name):
+class ApiBatch:
+    def __init__(self, api_name: str, start: int):
+        self.api_name = api_name
+        self.start = start
+        self.input_len = 1  # input的数量
+        self.params_end_index = start + 1  # params的结束index
+        self.output_end_index = start + 1  # output的结束index
+        self.params_grad_end_index = start + 1  # params_grad的结束index
+        # 内部state的标志("input", "output", "parameters", "parameters_grad"),
+        # 用于控制计算input_len, output_end_index, params_end_index, self.params_grad_end_index
+        self._state = Const.INPUT  # api_batch初始化为input
+
+    def set_state(self, state: str):
+        """设置当前状态"""
+        if state in {Const.INPUT, Const.OUTPUT, Const.KWARGS, Const.PARAMS, Const.PARAMS_GRAD}:
+            self._state = state
+        else:
+            raise ValueError(f"Invalid state: {state}")
+
+    def increment(self, state: str):
+        self.set_state(state)
+        if self._state == Const.INPUT or self._state == Const.KWARGS:
+            self.input_len += 1
+            self.params_end_index += 1
+            self.output_end_index += 1
+        if self._state == Const.PARAMS:
+            self.params_end_index += 1
+            self.output_end_index += 1
+        if self._state == Const.OUTPUT:
+            self.output_end_index += 1
+        self.params_grad_end_index += 1
+
+
+def api_batches_update(api_batches, api_name, state, index):
     """
-    Get api/module name and state
-    example:
-    name = 'conv2d.forward.1.input.0'
-    return: ('conv2d.forward.1.', 'input')
-
-    name = 'Functional.pad.0.backward.output.0'
-    return: ('Functional.pad.0.backward.', 'output')
+    当一个api的所有item更新完后，input, output的索引范围：
+    input: [start: start+input_len]
+    output: [start+input_len: output_end_index]
+    params: [output_end_index: params_end_index]
+    """
+    if not api_batches:
+        api_batches.append(ApiBatch(api_name, index))
+    else:
+        api_batch = api_batches[-1]
+        if api_batch.api_name == api_name or (
+                not re.search(Const.REGEX_FORWARD_BACKWARD, api_name) and api_name in api_batch.api_name):
+            try:
+                api_batch.increment(state)
+            except ValueError as e:
+                logger.error(f"api_batch: {api_batch} with invalid state, please check! {e}")
+                raise CompareException(CompareException.INVALID_STATE_ERROR) from e
+        else:
+            api_batches.append(ApiBatch(api_name, index))
 
-    name = 'x_tensor.0.debug.{index}'
-    return: ('x_tensor.0.', 'debug')
 
-    state type: input, output, kwargs, parameters, parameters_grad, debug
+def reorder_index(op_parsed_list):
     """
-    if not isinstance(name, str):
-        logger.error(f'Invalid name: {name}, type should be string, please check.')
-        raise CompareException(CompareException.INVALID_API_NAME_ERROR)
+    对单个api解析的op_items的index进行重排，将parameter的index放到output前面，返回新的重排后的index列表，op_parsed_list不变
+    """
+    index_param = []
+    index_output = []
+    index_param_grad = []
+    index_other = []
+    for i, op_item in enumerate(op_parsed_list[:-1]):
+        state = op_item.get(Const.STATE)
+        if state == Const.PARAMS:
+            index_param.append(i)
+        elif state == Const.OUTPUT:
+            index_output.append(i)
+        elif state == Const.PARAMS_GRAD:
+            index_param_grad.append(i)
+        else:
+            index_other.append(i)
+    # 合并others, parameters, 和output，确保parameters排在output前面
+    reordered_index_list = index_other + index_param + index_output + index_param_grad
+    return reordered_index_list
 
-    if Const.DEBUG in name.split(Const.SEP):
-        return name.split(Const.DEBUG)[0], Const.DEBUG
-    if Const.PARAMS_GRAD in name.split(Const.SEP):
-        return name.split(Const.PARAMS_GRAD)[0], Const.PARAMS_GRAD
 
-    split = re.split(Const.REGEX_FORWARD_BACKWARD, name)
-    if len(split) < 3:
-        logger.error(f'Invalid name string: {name}, can not be split by forward/backward, please check.')
-        raise CompareException(CompareException.INVALID_API_NAME_ERROR)
-    api = f'{split[0]}.{split[1]}.'
-    state_str = split[2]
-    match = re.match(r'^(\d+\.)?(input|output|kwargs|parameters)\..+$', state_str)
-    if not match:
-        raise CompareException(f'Invalid name string: {name}')
-    if match.group(1):
-        api = f'{api}{match.group(1)}'
-    state = match.group(2)
-    return api, state
-
-
-def reorder_op_name_list(op_name_list):
+def reorder_op_name_list(op_name_list, state_list):
     if not op_name_list:
-        return op_name_list
+        return op_name_list, state_list
 
     parameters = []
     output = []
     parameters_grad = []
     others = []
-    for x in op_name_list:
-        state = get_name_and_state(x)[1]
+    parameters_s = []
+    output_s = []
+    parameters_grad_s = []
+    others_s = []
+    for op_name, state in zip(op_name_list, state_list):
         if state == Const.PARAMS:
-            parameters.append(x)
+            parameters.append(op_name)
+            parameters_s.append(state)
         elif state == Const.OUTPUT:
-            output.append(x)
+            output.append(op_name)
+            output_s.append(state)
         elif state == Const.PARAMS_GRAD:
-            parameters_grad.append(x)
+            parameters_grad.append(op_name)
+            parameters_grad_s.append(state)
         else:
-            others.append(x)
+            others.append(op_name)
+            others_s.append(state)
     # 合并others, parameters, 和output，确保parameters排在output前面
     op_name_reorder = others + parameters + output + parameters_grad
-    return op_name_reorder
-
-
-def reorder_op_x_list(op_name_list, summary_list, data_name_list):
-    """对op_name, summary, data_name重新排序，把parameters放到input后output前，data_name由于统计量比对时，为None，单独处理"""
-    if not op_name_list or not summary_list:
-        return op_name_list, summary_list, data_name_list
-
-    index_map = {name: index for index, name in enumerate(op_name_list)}
-
-    op_name_reorder = reorder_op_name_list(op_name_list)
-    summary_reorder = [summary_list[index_map.get(name)] for name in op_name_reorder]
-    if data_name_list:
-        data_name_reorder = [data_name_list[index_map.get(name)] for name in op_name_reorder]
-    else:
-        data_name_reorder = data_name_list
-
-    return op_name_reorder, summary_reorder, data_name_reorder
+    state_reorder = others_s + parameters_s + output_s + parameters_grad_s
+    return op_name_reorder, state_reorder
 
 
 def process_summary_data(summary_data):
@@ -393,17 +461,22 @@ def stack_column_process(result_item, has_stack, index, key, npu_stack_info):
     return result_item
 
 
-def result_item_init(n_info, b_info, dump_mode):
+def result_item_init(n_info, b_info, requires_grad_pair, dump_mode):
     n_len = len(n_info.struct)
     b_len = len(b_info.struct)
+    # requires_grad_pair内部创建，固定两个元素
+    n_requires_grad = requires_grad_pair[0]
+    b_requires_grad = requires_grad_pair[1]
+    req_grad_consist = n_requires_grad == b_requires_grad
     struct_long_enough = (n_len > 2 and b_len > 2) if dump_mode == Const.MD5 else (n_len > 1 and b_len > 1)
     if struct_long_enough:
         result_item = [
-            n_info.name, b_info.name, n_info.struct[0], b_info.struct[0], n_info.struct[1], b_info.struct[1]
+            n_info.name, b_info.name, n_info.struct[0], b_info.struct[0], n_info.struct[1], b_info.struct[1],
+            n_requires_grad, b_requires_grad
         ]
         if dump_mode == Const.MD5:
             md5_compare_result = CompareConst.PASS if n_info.struct[2] == b_info.struct[2] else CompareConst.DIFF
-            result_item.extend([n_info.struct[2], b_info.struct[2], md5_compare_result])
+            result_item.extend([n_info.struct[2], b_info.struct[2], req_grad_consist, md5_compare_result])
         elif dump_mode == Const.SUMMARY:
             result_item.extend([" "] * 8)  # 8个统计量数据情况的比对指标
         else:
@@ -449,11 +522,15 @@ def get_accuracy(result, n_dict, b_dict, dump_mode):
             b_name = safe_get_value(b_dict, b_start + index, "b_dict", key="op_name")
             n_struct = safe_get_value(n_dict, index, "n_dict", key=key)
             b_struct = safe_get_value(b_dict, index, "b_dict", key=key)
+            n_requires_grad = safe_get_value(n_dict, n_start + index, "n_dict", key='requires_grad')
+            b_requires_grad = safe_get_value(b_dict, b_start + index, "b_dict", key='requires_grad')
+            requires_grad_pair = [n_requires_grad, b_requires_grad]
+            req_grad_consist = n_requires_grad == b_requires_grad
             err_msg = ""
 
             npu_info = ApiItemInfo(n_name, n_struct, npu_stack_info)
             bench_info = ApiItemInfo(b_name, b_struct, bench_stack_info)
-            result_item = result_item_init(npu_info, bench_info, dump_mode)
+            result_item = result_item_init(npu_info, bench_info, requires_grad_pair, dump_mode)
 
             if dump_mode == Const.MD5:
                 result_item = stack_column_process(result_item, has_stack, index, key, npu_stack_info)
@@ -469,6 +546,8 @@ def get_accuracy(result, n_dict, b_dict, dump_mode):
                 result_item, accuracy_check, err_msg = get_rela_diff_summary_mode(result_item, npu_summary_data,
                                                                                   bench_summary_data, err_msg)
 
+            result_item.append(req_grad_consist)
+            err_msg += "Requires_grad inconsistent." if not req_grad_consist else ""
             result_item.append(accuracy_check if dump_mode == Const.SUMMARY else CompareConst.ACCURACY_CHECK_YES)
             result_item.append(err_msg)
             result_item = stack_column_process(result_item, has_stack, index, key, npu_stack_info)
@@ -482,23 +561,30 @@ def get_accuracy(result, n_dict, b_dict, dump_mode):
         if n_len > b_len:
             for index in range(b_len, n_len):
                 try:
-                    n_name = n_dict['op_name'][n_start + index]
-                    n_struct = n_dict[key][index]
+                    n_name = safe_get_value(n_dict, n_start + index, "n_dict", key="op_name")
+                    n_struct = safe_get_value(n_dict, index, "n_dict", key=key)
+                    n_requires_grad = safe_get_value(n_dict, n_start + index, "n_dict", key='requires_grad')
+
                     if dump_mode == Const.MD5:
                         result_item = [
                             n_name, CompareConst.NAN, n_struct[0], CompareConst.NAN, n_struct[1], CompareConst.NAN,
-                            n_struct[2], CompareConst.NAN, CompareConst.NAN
+                            n_requires_grad, CompareConst.NAN,
+                            n_struct[2], CompareConst.NAN,
+                            False,
+                            CompareConst.NAN
                         ]
                         result.append(result_item)
                         continue
                     result_item = [
                         n_name, CompareConst.NAN, n_struct[0], CompareConst.NAN, n_struct[1], CompareConst.NAN,
+                        n_requires_grad, CompareConst.NAN,
                         " ", " ", " ", " ", " ", " "
                     ]
                     summary_data = n_dict.get(CompareConst.SUMMARY)[n_start + index]
                     result_item.extend(summary_data)
                     summary_data = [CompareConst.NAN for _ in range(len(n_dict.get(CompareConst.SUMMARY)[0]))]
                     result_item.extend(summary_data)
+                    result_item.append(False)
                 except IndexError as e:
                     err_msg = "index out of bounds error occurs, please check!\n" \
                               f"n_dict is {n_dict}"
@@ -546,6 +632,23 @@ def make_result_table(result, dump_mode, stack_mode):
     return result_df
 
 
+def gen_api_batches(result: np.ndarray, header: list):
+    api_name_index = header.index(Const.API_ORIGIN_NAME)
+    state_name_index = header.index(Const.STATE)
+    api_batches = []
+    for i, res_i in enumerate(result):
+        api_name = safe_get_value(res_i, api_name_index, "res_i")
+        state = safe_get_value(res_i, state_name_index, "res_i")
+        api_batches_update(api_batches, api_name, state, i)
+    return api_batches
+
+
+def get_paired_dirs(npu_path, bench_path):
+    npu_dirs = set(os.listdir(npu_path))
+    bench_dirs = set(os.listdir(bench_path))
+    return list(npu_dirs & bench_dirs)
+
+
 def _compare_parser(parser):
     parser.add_argument("-i", "--input_path", dest="input_path", type=str,
                         help="<Required> The compare input path, a dict json.", required=True)
@@ -558,6 +661,8 @@ def _compare_parser(parser):
                         help="<optional> Whether to give advisor.", required=False)
     parser.add_argument("-f", "--fuzzy_match", dest="fuzzy_match", action="store_true",
                         help="<optional> Whether to perform a fuzzy match on the api name.", required=False)
+    parser.add_argument("-hl", "--highlight", dest="highlight", action="store_true",
+                        help="<optional> Whether to set result highlighting.", required=False)
     parser.add_argument("-cm", "--cell_mapping", dest="cell_mapping", type=str, nargs='?', const=True,
                         help="<optional> The cell mapping file path.", required=False)
     parser.add_argument("-am", "--api_mapping", dest="api_mapping", type=str, nargs='?', const=True,
@@ -566,40 +671,143 @@ def _compare_parser(parser):
                         help="<optional> The data mapping file path.", required=False)
     parser.add_argument("-lm", "--layer_mapping", dest="layer_mapping", type=str, nargs='?', const=True,
                         help="<optional> The layer mapping file path.", required=False)
+    parser.add_argument("-da", "--diff_analyze", dest="diff_analyze", action="store_true",
+                        help="<optional> Whether to perform a diff analyze on the api name.", required=False)
 
 
-def compare_distributed_inner(npu_dump_dir, bench_dump_dir, output_path, compare_func, **kwargs):
-    if kwargs.get('suffix'):
-        logger.error("Argument 'suffix' is not supported for compare_distributed.")
-        raise CompareException(CompareException.INVALID_PARAM_ERROR)
-    is_print_compare_log = kwargs.get('is_print_compare_log', True)
-    # get the ranks and match by order
-    npu_ranks = sorted(check_and_return_dir_contents(npu_dump_dir, 'rank'))
-    bench_ranks = sorted(check_and_return_dir_contents(bench_dump_dir, 'rank'))
+def get_sorted_ranks(npu_dump_dir, bench_dump_dir):
+    """
+    get the ranks and match by order
+    """
+    unsorted_npu_ranks = check_and_return_dir_contents(npu_dump_dir, 'rank')
+    unsorted_bench_ranks = check_and_return_dir_contents(bench_dump_dir, 'rank')
+    # 正则匹配已经校验rank后面必是数字，或者无数字的rank
+    npu_ranks = sorted(unsorted_npu_ranks, key=lambda x: int(x[4:]) if len(x) > 4 else -1)  # 前四个字符都是rank，后面是卡号
+    bench_ranks = sorted(unsorted_bench_ranks, key=lambda x: int(x[4:]) if len(x) > 4 else -1)
     if len(npu_ranks) != len(bench_ranks):
         logger.error('The number of ranks in the two runs are different. '
                      'Unable to match the ranks. Please use another folder to compare '
                      'or use compare() api and manually match the ranks.')
         raise CompareException(CompareException.INVALID_PATH_ERROR)
-    for nr, br in zip(npu_ranks, bench_ranks):
+    return npu_ranks, bench_ranks
+
+
+def multi_statistics_compare(func, func_args):
+    def err_call(args):
+        logger.error(f'Multiprocess statistics compare failed! Reason: {args}')
+
+    compare_func, input_param_nr_list, output_path, kwargs = func_args
+
+    param_num = len(input_param_nr_list)
+    process_num = max(int((multiprocessing.cpu_count() + 1) // 4), 1)
+    if param_num <= process_num:
+        process_num = param_num
+        chunks = [[input_param_nr] for input_param_nr in input_param_nr_list]
+    else:
+        chunk_size = param_num // process_num
+        remainder = param_num % process_num
+        chunks = [input_param_nr_list[i:i + chunk_size] for i in range(0, param_num - remainder, chunk_size)]
+        for i in range(remainder):
+            chunks[i].append(input_param_nr_list[param_num - remainder + i])
+
+    pool = multiprocessing.Pool(process_num)
+
+    async_results = []
+    for chunk in chunks:
+        result = pool.apply_async(func, args=(compare_func, chunk, output_path, kwargs), error_callback=err_call)
+        async_results.append(result)
+
+    pool.close()
+
+    for ar in async_results:
+        try:
+            ar.get(timeout=3600)
+        except Exception as e:
+            logger.error(f"Task failed with exception: {e}")
+            pool.terminate()
+            raise CompareException(CompareException.MULTIPROCESS_ERROR) from e
+
+    pool.join()
+
+
+def mp_logger_init(ranks_str):
+    """
+    多进程比对需要对logger进行wrap和patch，在日志前加上卡号信息，从而实现不同进程日志的隔离
+    """
+
+    def wrap_logger(fn):
+        def inner(msg, *args, **kwargs):
+            return fn(ranks_str + msg, *args, **kwargs)
+        return inner
+
+    logger.info = wrap_logger(logger.info)
+    logger.warning = wrap_logger(logger.warning)
+    logger.error = wrap_logger(logger.error)
+
+
+def multi_ranks_compare(compare_func, input_param_nr_list, output_path, kwargs):
+    """
+    将多卡数据分成多进程后，单进程内可能还有多张卡的数据，因此还需要多次比对
+    """
+    rank_list = [input_param_nr[1] for input_param_nr in input_param_nr_list]  # input_param_nr内部数据结构，2元素tuple
+    ranks_str = f"[{' '.join(rank_list)}]"
+    mp_logger_init(ranks_str)
+    for input_param_nr in input_param_nr_list:
+        input_param, nr = input_param_nr
+        compare_entry(compare_func, input_param, output_path, nr, kwargs)
+
+
+def compare_entry(compare_func, input_param, output_path, nr, kwargs):
+    try:
+        compare_func(input_param=input_param, output_path=output_path, suffix=f'_{nr}', **kwargs)
+    except CompareException as e:
+        if e.code == CompareException.INVALID_DATA_ERROR:
+            logger.error(f"Invalid or missing 'data' in dump.json. Skipping {nr} comparison.")
+        if e.code == CompareException.INVALID_TASK_ERROR:
+            logger.error(f"Invalid or missing 'task' in dump.json. Skipping {nr} comparison.")
+
+
+def compare_distributed_inner(npu_dump_dir, bench_dump_dir, output_path, compare_func, **kwargs):
+    def extract_compare_param(_file_type):
         npu_data_dir = os.path.join(npu_dump_dir, nr)
         bench_data_dir = os.path.join(bench_dump_dir, br)
+        npu_path = extract_json(npu_data_dir, _file_type)
+        bench_path = extract_json(bench_data_dir, _file_type)
+        if npu_path == "" or bench_path == "":
+            logger.debug(f'Did not find paired {_file_type} in {nr} and {br}, skip comparing.')
+            return {}, True
+        _input_param = {
+            'npu_json_path': npu_path,
+            'bench_json_path': bench_path,
+            'is_print_compare_log': kwargs.get('is_print_compare_log', True)
+        }
+        return _input_param, False
+
+    if kwargs.get('suffix'):
+        logger.error("Argument 'suffix' is not supported for compare_distributed.")
+        raise CompareException(CompareException.INVALID_PARAM_ERROR)
+
+    npu_ranks, bench_ranks = get_sorted_ranks(npu_dump_dir, bench_dump_dir)
+
+    # 统计量、md5比对
+    pre_check_dump_path = os.path.join(npu_dump_dir, npu_ranks[0], 'dump.json') if npu_ranks else ''
+    if not pre_check_dump_path:
+        return
+    dump_data = load_json(pre_check_dump_path)
+    if dump_data.get('task') == Const.STATISTICS:
+        # dump数据为统计量或md5时，多进程加速比对
+        input_param_nr_list = []
+        for nr, br in zip(npu_ranks, bench_ranks):
+            input_param, skip = extract_compare_param(Const.DUMP_JSON_FILE)
+            if not skip:
+                input_param_nr_list.append((input_param, nr))
+        func_args = (compare_func, input_param_nr_list, output_path, kwargs)
+        multi_statistics_compare(multi_ranks_compare, func_args)
+        return
+
+    # 真实数据比对
+    for nr, br in zip(npu_ranks, bench_ranks):
         for file_type in [Const.DUMP_JSON_FILE, Const.DEBUG_JSON_FILE]:
-            npu_path = extract_json(npu_data_dir, file_type)
-            bench_path = extract_json(bench_data_dir, file_type)
-            if npu_path == "" or bench_path == "":
-                logger.debug(f'Did not find paired {file_type} in {npu_data_dir} and {bench_data_dir},'
-                             ' skip comparing.')
-                continue
-            dump_result_param = {
-                'npu_json_path': npu_path,
-                'bench_json_path': bench_path,
-                'is_print_compare_log': is_print_compare_log
-            }
-            try:
-                compare_func(input_param=dump_result_param, output_path=output_path, suffix=f'_{nr}', **kwargs)
-            except CompareException as e:
-                if e.code == CompareException.INVALID_DATA_ERROR:
-                    logger.error(f"Invalid or missing 'data' in dump.json. Skipping {nr} comparison.")
-                if e.code == CompareException.INVALID_TASK_ERROR:
-                    logger.error(f"Invalid or missing 'task' in dump.json. Skipping {nr} comparison.")
+            input_param, skip = extract_compare_param(file_type)
+            if not skip:
+                compare_entry(compare_func, input_param, output_path, nr, kwargs)