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

msprobe/visualization/graph/distributed_analyzer.py

@@ -0,0 +1,395 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from enum import Enum
+from msprobe.visualization.utils import GraphConst
+from msprobe.core.common.const import Const, CompareConst
+from msprobe.core.common.log import logger
+
+
+class CommunicationType(Enum):
+    """
+    通信类型：发送、接收、发送接收
+    """
+    SEND = 'send'
+    RECEIVE = 'receive'
+    SEND_RECEIVE = 'send_receive'
+
+
+class DistributedType(Enum):
+    """
+    分布式类型：点对点通信、集体通信
+    """
+    P2P = 'p2p'
+    COLLECTIVE = 'collective'
+
+
+CANNOT_MATCH = 'cannot match distributed node in rank'
+
+
+class DistributedAnalyzer:
+
+    def __init__(self, graphs: dict, overflow_check: bool):
+        self.graphs = graphs
+        self.overflow_check = overflow_check
+        self.config = {
+            # 当前通信api名称: 匹配目标通信api名称, 获取rank信息的位置参数或关键字参数, 通信类型, 分布式类型
+            'send': ['recv', GraphConst.DST, CommunicationType.SEND.value, DistributedType.P2P],
+            'isend': ['irecv', GraphConst.DST, CommunicationType.SEND.value, DistributedType.P2P],
+            'recv': ['send', GraphConst.SRC, CommunicationType.RECEIVE.value, DistributedType.P2P],
+            'irecv': ['isend', GraphConst.SRC, CommunicationType.RECEIVE.value, DistributedType.P2P],
+            'broadcast': ['broadcast', '1', CommunicationType.SEND.value, DistributedType.COLLECTIVE],
+            'scatter': ['scatter', GraphConst.SRC, CommunicationType.SEND.value, DistributedType.COLLECTIVE],
+            'gather': ['gather', GraphConst.DST, CommunicationType.RECEIVE.value, DistributedType.COLLECTIVE],
+            'reduce': ['reduce', '1', CommunicationType.RECEIVE.value, DistributedType.COLLECTIVE]
+        }
+        self.group_node_mapping = {}
+        self._make_group_node_mapping()
+
+    @staticmethod
+    def _get_opposite_communication_type(action):
+        if action == CommunicationType.SEND.value:
+            return CommunicationType.RECEIVE.value
+        elif action == CommunicationType.RECEIVE.value:
+            return CommunicationType.SEND.value
+        return action
+
+    @staticmethod
+    def _node_output_all_equal(data: dict, target_data: dict):
+        keys_to_compare = [Const.DTYPE, Const.SHAPE, Const.MAX, Const.MIN, Const.MEAN, Const.NORM]
+        return all(data.get(key) == target_data.get(key) for key in keys_to_compare)
+
+    @staticmethod
+    def _get_target_rank(node, rank, parameter):
+        """
+        点对点通信, 从输出数据参数src或dst, 获取通信目标rank
+        一对多通信和多对一通信, 从输出数据参数src或dst或位置参数, 获取发送或接收的rank源头
+        :param node: 当前节点
+        :param rank: 当前rank
+        :param parameter: 输出数据参数
+        :return: 目标rank
+        """
+        target_rank = node.input_data.get(f'{node.id}{GraphConst.INPUT}{parameter}', {}).get('value')
+        if target_rank is None:
+            logger.warning(f'The parameter {parameter} of node {node.id} does not exist, {CANNOT_MATCH}{rank}')
+        return target_rank
+
+    @staticmethod
+    def _get_group_info(node, rank):
+        """
+        获取当前通信节点的group参数中的group_ranks和group_id
+        :param node: 当前通信节点
+        :param rank: 当前rank
+        :return: group_ranks和group_id
+        """
+        group = node.input_data.get(f'{node.id}{GraphConst.INPUT}group', {})
+        if not group:
+            logger.warning(f'The kwarg group of node {node.id} does not exist, {CANNOT_MATCH}{rank}')
+            return None, None
+        group_ranks = group.get('group_ranks')
+        if not group_ranks:
+            logger.warning(f'The group_ranks of node {node.id} does not exist, {CANNOT_MATCH}{rank}')
+            return None, None
+        group_id = group.get('group_id')
+        if not group_id:
+            logger.warning(f'The group_id of node {node.id} does not exist, {CANNOT_MATCH}{rank}')
+            return None, None
+        return group_ranks, group_id
+
+    @staticmethod
+    def _get_batch_group_info(node, rank):
+        for data in node.input_data.values():
+            group_id = data.get('group_id')
+            if group_id is not None:
+                return group_id
+        logger.warning(f'The group_id of node {node.id} does not exist, {CANNOT_MATCH}{rank}')
+        return None
+
+    def distributed_match(self):
+        for rank, graph in self.graphs.items():
+            nodes = graph.node_map
+            for node_id, node in nodes.items():
+                # 不是通信节点或者已经匹配过了
+                if not node_id.startswith(Const.DISTRIBUTED) or node.matched_distributed:
+                    continue
+                api_name, distributed_type = self._get_distributed_name_and_type(node_id)
+                if api_name == GraphConst.BATCH_P2P:
+                    self._batch_p2p_match(node, rank)
+                elif distributed_type == DistributedType.P2P:
+                    self._p2p_match(node, rank, api_name)
+                else:
+                    self._collective_match(node, rank, api_name)
+
+    def _make_group_node_mapping(self):
+        """
+        建立通信节点的全局唯一标识映射
+        key: rank号, value: unique_group_id与node_id之间的映射
+        {
+            "0": {
+                "unique_group_id1": "node_id1",
+                "unique_group_id2": "node_id2",
+                "node_id1": "unique_group_id1",
+                "node_id2": "unique_group_id2"
+            },
+            "1": {},
+            "2": {}
+        }
+        """
+        for rank, graph in self.graphs.items():
+            group_count = {}
+            group_info = {}
+            batch_p2p_count = {}
+            nodes = graph.node_map
+            for node_id, node in nodes.items():
+                if not node_id.startswith(Const.DISTRIBUTED):
+                    continue
+                api_name, distributed_type = self._get_distributed_name_and_type(node_id)
+                if api_name == GraphConst.BATCH_P2P:
+                    self._make_batch_p2p_mapping(node, rank, batch_p2p_count)
+                    continue
+                elif distributed_type == DistributedType.P2P:
+                    config_info = self.config.get(api_name)
+                    target_rank = self._get_target_rank(node, rank, config_info[1])
+                    if target_rank is None:
+                        continue
+                    # p2p通信节点，api名称+传输目标rank作为group_id
+                    group_id = api_name + Const.RANK + str(target_rank)
+                else:
+                    # 其他通信节点直接获取group_id, 并拼接api名称
+                    _, group_id = self._get_group_info(node, rank)
+                    if not group_id:
+                        continue
+                    group_id += api_name
+                # 同group_id的调用次数累计
+                group_count[group_id] = group_count.get(group_id, 0) + 1
+                # group_id+同group_id的调用次数作为唯一的unique_group_id
+                unique_group_id = group_id + Const.REPLACEMENT_CHARACTER + str(group_count.get(group_id))
+                group_info[unique_group_id] = node_id
+                group_info[node_id] = unique_group_id
+            if rank not in self.group_node_mapping:
+                self.group_node_mapping[rank] = {}
+            self.group_node_mapping[rank].update(group_info)
+
+    def _make_batch_p2p_mapping(self, node, rank, batch_p2p_count):
+        """
+        给batch_isend_irecv接口的每个p2p内容赋予唯一标识
+        """
+        if rank not in self.group_node_mapping:
+            self.group_node_mapping[rank] = {}
+        params = []
+        for info_dict in node.batch_p2p_info:
+            op = info_dict.get(GraphConst.OP)
+            target_rank = info_dict.get(GraphConst.PEER)
+            if op is None or target_rank is None:
+                logger.warning('Cannot get param op or peer.')
+                continue
+            group_id = op + Const.REPLACEMENT_CHARACTER + Const.RANK + str(target_rank) + \
+                       Const.REPLACEMENT_CHARACTER + info_dict.get(GraphConst.GROUP_ID, '')
+            batch_p2p_count[group_id] = batch_p2p_count.get(group_id, 0) + 1
+            # 例如: isend_rank0_5a4d31ad765260ba50eb190f1f9fd163_1
+            unique_group_id = group_id + Const.REPLACEMENT_CHARACTER + str(batch_p2p_count.get(group_id))
+            params.append(unique_group_id)
+            self.group_node_mapping.get(rank)[unique_group_id] = node.id
+        if params:
+            self.group_node_mapping.get(rank)[node.id] = params
+
+    def _get_distributed_name_and_type(self, node_id):
+        if Const.SEP not in node_id:
+            raise ValueError(f'Invalid node id {node_id}.')
+        api_name = node_id.split(Const.SEP)[1]
+        if api_name in self.config:
+            return api_name, self.config.get(api_name)[3]
+        return api_name, DistributedType.COLLECTIVE
+
+    def _get_target_node(self, rank, unique_group_id, api_name, target_rank, target_api_name=None):
+        """
+        获取名称匹配上的目标节点
+        :param rank: 当前rank
+        :param unique_group_id: 当前节点唯一group id
+        :param api_name: 当前节点的api名称, 例如Distributed.isend.0.forward, api名称为isend
+        :param target_rank: 与当前节点产生通信的rank
+        :param target_api_name: 与当前节点产生通信的节点api名称, 仅p2p通信需要配置
+        :return: 目标节点
+        """
+        target_graph = self.graphs.get(target_rank)
+        if not target_graph:
+            logger.warning(f'Graph data does not exist, {CANNOT_MATCH}{target_rank}')
+            return None
+        target_group_mapping = self.group_node_mapping.get(target_rank)
+        # p2p通信，想要获取目标节点，需要替换unique_group_id中的rank和api name,
+        # 例如isend发送到rank1，对应的irecv接收自rank0, isend_rank1与irecv_rank0对应
+        target_unique_group_id = (unique_group_id
+                                  .replace(Const.RANK + str(target_rank), Const.RANK + str(rank))
+                                  .replace(api_name, target_api_name)) if target_api_name else unique_group_id
+        target_node_id = target_group_mapping.get(target_unique_group_id, '')
+        target_node = target_graph.node_map.get(target_node_id)
+        if not target_node:
+            logger.warning(f'Node {target_node_id} does not exist, {CANNOT_MATCH}{target_rank}')
+            return None
+        return target_node
+
+    def _add_node_matched_distributed(self, node, target_node, api_name, target_rank, reversal_type=False):
+        """
+        给当前节点添加matched_distributed字段信息
+        :param node: 当前节点
+        :param target_node: 匹配上的目标节点
+        :param api_name: 当前节点的api名称
+        :param target_rank: 匹配上的目标rank
+        :param reversal_type: 是否需要反转通信类型，例如broadcast在rank0通信类型是发送，但在其他rank通信类型是接收
+        """
+        communications_type = self.config.get(api_name)[2]
+        communications_type = self._get_opposite_communication_type(communications_type) if reversal_type \
+            else communications_type
+        index = target_node.data.get(GraphConst.OVERFLOW_LEVEL, CompareConst.NAN) if self.overflow_check \
+            else target_node.data.get(GraphConst.JSON_INDEX_KEY, CompareConst.NAN)
+        matched_distributed = {
+            'communications_type': communications_type,
+            'nodes_info': {target_rank: [str(index), target_node.id]}
+        }
+        node.matched_distributed = matched_distributed
+
+    def _p2p_match(self, node, rank, api_name):
+        """
+        点对点通信匹配
+
+        根据当前点对点通信节点的输出数据中的src或dst参数, 确定目标rank, 并从目标rank中找到对应的点对点通信节点, 校验输出数据是否一致，
+        校验通过则在两个匹配节点增加匹配信息
+        Args:
+            node: 当前点对点通信节点
+            rank: 当前节点所属rank
+            api_name: 当前节点的api名称
+        Returns:
+        """
+        config_info = self.config.get(api_name)
+        target_api_name = config_info[0]
+        #
+        target_rank = self._get_target_rank(node, rank, config_info[1])
+        if target_rank is None:
+            return
+        unique_group_id = self.group_node_mapping.get(rank, {}).get(node.id, '')
+        target_node = self._get_target_node(rank, unique_group_id, api_name, target_rank, target_api_name)
+        if not target_node:
+            return
+        target_config_info = self.config.get(target_api_name)
+        source_rank = (target_node.input_data.get(f'{target_node.id}{GraphConst.INPUT}{target_config_info[1]}', {})
+                       .get('value'))
+        if source_rank is None:
+            logger.warning(
+                f'The kwarg {target_config_info[1]} of node {target_node.id} does not exist, '
+                f'{CANNOT_MATCH}{source_rank}')
+            return
+        if source_rank != rank:
+            # 点对点通信，待匹配目标节点包含的rank信息要与当前rank一致
+            logger.warning(
+                f'{node.id} of rank{rank} is expected to communicate with {target_node.id} of rank{target_rank}, '
+                f'but the data shows that {target_node.id} communicates with rank{source_rank}.'
+                f'The rank is inconsistent, cannot match distributed node')
+            return
+
+        # 点对点通信，两个匹配节点的输出数据要一致
+        if not DistributedAnalyzer._node_output_all_equal(node.output_data.get(node.id + '.output.0'),
+                                                          target_node.output_data.get(target_node.id + '.output.0')):
+            logger.warning(f'{node.id} output of rank{rank} is different from the {target_node.id} '
+                           f'output of rank{target_rank}, cannot match distributed node')
+            return
+
+        self._add_node_matched_distributed(node, target_node, api_name, target_rank)
+        self._add_node_matched_distributed(target_node, node, target_api_name, rank)
+
+    def _collective_match(self, node, rank, api_name):
+        """
+        集体通信匹配
+
+        一对多通信和多对一通信, 需要先获取节点输出数据中的src或dst或位置参数, 确定发送源或接收源, 多对多通信不需要
+        :param node: 当前集体通信节点
+        :param rank: 当前节点所属rank
+        :param api_name: 当前节点的api名称
+        :return:
+        """
+        communications_type = CommunicationType.SEND_RECEIVE.value
+        config_info = self.config.get(api_name)
+        if config_info:
+            # 此时为一对多通信或多对一通信
+            source_rank = self._get_target_rank(node, rank, config_info[1])
+            if source_rank is None or str(source_rank) != str(rank):
+                return
+            communications_type = config_info[2]
+        group_ranks, group_id = self._get_group_info(node, rank)
+        if not group_ranks or not group_id:
+            return
+        unique_group_id = self.group_node_mapping.get(rank, {}).get(node.id, '')
+        matched_distributed = {'communications_type': communications_type}
+        nodes_info = {}
+        for target_rank in group_ranks:
+            if str(target_rank) == str(rank):
+                continue
+            target_node = self._get_target_node(rank, unique_group_id, api_name, target_rank)
+            if not target_node:
+                continue
+            _, target_group_id = self._get_group_info(target_node, target_rank)
+            if not target_group_id:
+                continue
+            if group_id != target_group_id:
+                logger.warning(
+                    f'{node.id} of rank{rank} is expected to communicate with {target_node.id} of rank{target_rank}'
+                    f', but the data shows that the group id of the two nodes are different, '
+                    f'cannot match distributed node')
+                continue
+            # 给当前通信节点添加matched_distributed字段信息
+            index = target_node.data.get(GraphConst.OVERFLOW_LEVEL, CompareConst.NAN) if self.overflow_check \
+                else target_node.data.get(GraphConst.JSON_INDEX_KEY, CompareConst.NAN)
+            nodes_info[target_rank] = [str(index), target_node.id]
+            if config_info:
+                # 给匹配上的目标节点也添加matched_distributed字段信息
+                self._add_node_matched_distributed(target_node, node, api_name, rank, True)
+        if nodes_info:
+            matched_distributed['nodes_info'] = nodes_info
+            node.matched_distributed = matched_distributed
+
+    def _batch_p2p_match(self, node, rank):
+        """
+        批量点对点匹配
+
+        针对torch.distributed.batch_isend_irecv接口，其入参是一个包含点对点通信信息的集合，需要遍历集合对每个点对点通信信息进行匹配
+        :param node: 当前集体通信节点
+        :param rank: 当前节点所属rank
+        :return:
+        """
+        unique_group_ids = self.group_node_mapping.get(rank, {}).get(node.id)
+        if not unique_group_ids:
+            return
+        matched_distributed = [] if len(unique_group_ids) > 1 else {}
+        for unique_group_id in unique_group_ids:
+            try:
+                id_info = unique_group_id.split(Const.REPLACEMENT_CHARACTER)
+                api_name = id_info[0]
+                target_api_name = self.config.get(api_name)[0]
+                target_rank = int(id_info[1].replace(Const.RANK, ''))
+            except Exception as e:
+                logger.warning(f'Failed to parsing batch p2p parameter with error info: {e}.')
+                continue
+            target_node = self._get_target_node(rank, unique_group_id, api_name, target_rank, target_api_name)
+            if not target_node:
+                continue
+            communications_type = self.config.get(api_name)[2]
+            index = target_node.data.get(GraphConst.OVERFLOW_LEVEL, CompareConst.NAN) if self.overflow_check \
+                else target_node.data.get(GraphConst.JSON_INDEX_KEY, CompareConst.NAN)
+            matched_info = {
+                'communications_type': communications_type,
+                'nodes_info': {target_rank: [str(index), target_node.id]}
+            }
+            matched_distributed.append(matched_info) if isinstance(matched_distributed, list) \
+                else matched_distributed.update(matched_info)
+        if matched_distributed:
+            node.matched_distributed = matched_distributed

msprobe/visualization/graph/graph.py

@@ -67,6 +67,15 @@ class Graph:
         ancestors_b = node_b.get_ancestors()
         return node_b, ancestors_n, ancestors_b
 
+
+    @staticmethod
+    def fuzzy_match(node_n, node_b):
+        if not node_n or not node_b or not node_n.fuzzy_eq(node_b):
+            return None, [], []
+        ancestors_n = node_n.get_ancestors()
+        ancestors_b = node_b.get_ancestors()
+        return node_b, ancestors_n, ancestors_b
+
     @staticmethod
     def dfs(node, result):
         info = node.to_dict()

msprobe/visualization/graph/node_op.py

@@ -16,6 +16,7 @@
 from enum import Enum
 import re
 from msprobe.visualization.builder.msprobe_adapter import op_patterns
+from msprobe.core.common.log import logger
 
 
 class NodeOp(Enum):
@@ -32,8 +33,9 @@ class NodeOp(Enum):
         for op in NodeOp:
             index = op.value
             if index < 0 or index >= len(op_patterns):
-                raise Exception("NodeOp and op_patterns in MsprobeAdapter do not match")
+                continue
             pattern = op_patterns[index]
             if re.match(pattern, node_name):
                 return op
-        raise Exception(f"Cannot parse node_name {node_name} into NodeOp")
+        logger.warning(f"Cannot parsing node_name {node_name} into NodeOp, default parsing as module.")
+        return NodeOp.module

msprobe/visualization/graph_service.py

@@ -16,8 +16,8 @@
 import os
 import time
 import json
-from msprobe.core.common.file_utils import (FileOpen, check_file_type, create_directory, FileChecker,
-                                            check_file_or_directory_path)
+from msprobe.core.common.file_utils import (check_file_type, create_directory, FileChecker,
+                                            check_file_or_directory_path, load_json)
 from msprobe.core.common.const import FileCheckConst, Const
 from msprobe.core.common.utils import CompareException
 from msprobe.core.overflow_check.checker import AnomalyDetector
@@ -28,11 +28,12 @@ from msprobe.core.common.log import logger
 from msprobe.visualization.graph.node_colors import NodeColors
 from msprobe.core.compare.layer_mapping import generate_api_mapping_by_layer_mapping
 from msprobe.core.compare.utils import check_and_return_dir_contents
+from msprobe.visualization.graph.distributed_analyzer import DistributedAnalyzer
 
 current_time = time.strftime("%Y%m%d%H%M%S")
 
 
-def _compare_graph(input_param, args, output_file_name=f'compare_{current_time}.vis'):
+def _compare_graph(input_param, args):
     logger.info('Start building model graphs...')
     # 对两个数据进行构图
     dump_path_n = input_param.get('npu_path')
@@ -49,8 +50,8 @@ def _compare_graph(input_param, args, output_file_name=f'compare_{current_time}.
                                FileCheckConst.READ_ABLE).common_check()
     stack_path_b = FileChecker(os.path.join(dump_path_b, GraphConst.STACK_FILE), FileCheckConst.FILE,
                                FileCheckConst.READ_ABLE).common_check()
-    graph_n = GraphBuilder.build(construct_path_n, data_path_n, stack_path_n)
-    graph_b = GraphBuilder.build(construct_path_b, data_path_b, stack_path_b)
+    graph_n = GraphBuilder.build(construct_path_n, data_path_n, stack_path_n, complete_stack=args.complete_stack)
+    graph_b = GraphBuilder.build(construct_path_b, data_path_b, stack_path_b, complete_stack=args.complete_stack)
     logger.info('Model graphs built successfully, start Comparing graphs...')
     # 基于graph、stack和data进行比较
     dump_path_param = {
@@ -66,8 +67,7 @@ def _compare_graph(input_param, args, output_file_name=f'compare_{current_time}.
             mapping_dict = generate_api_mapping_by_layer_mapping(data_path_n, data_path_b, yaml_path)
         except Exception:
             logger.warning('The layer mapping file parsing failed, please check file format, mapping is not effective.')
-    graph_comparator = GraphComparator([graph_n, graph_b], dump_path_param, args.output_path, args.framework,
-                                       mapping_dict=mapping_dict)
+    graph_comparator = GraphComparator([graph_n, graph_b], dump_path_param, args, mapping_dict=mapping_dict)
     graph_comparator.compare()
     micro_steps = graph_n.paging_by_micro_step(graph_b)
     # 开启溢出检测
@@ -75,16 +75,22 @@ def _compare_graph(input_param, args, output_file_name=f'compare_{current_time}.
         graph_n.overflow_check()
         graph_b.overflow_check()
 
+    return CompareGraphResult(graph_n, graph_b, graph_comparator, micro_steps)
+
+
+def _export_compare_graph_result(args, graphs, graph_comparator, micro_steps,
+                                 output_file_name=f'compare_{current_time}.vis'):
     create_directory(args.output_path)
     output_path = os.path.join(args.output_path, output_file_name)
     task = GraphConst.GRAPHCOMPARE_MODE_TO_DUMP_MODE_TO_MAPPING.get(graph_comparator.ma.compare_mode)
-    export_config = GraphExportConfig(graph_n, graph_b, graph_comparator.ma.get_tool_tip(),
-                                      NodeColors.get_node_colors(graph_comparator.ma.compare_mode), micro_steps, task)
+    export_config = GraphExportConfig(graphs[0], graphs[1], graph_comparator.ma.get_tool_tip(),
+                                      NodeColors.get_node_colors(graph_comparator.ma.compare_mode), micro_steps, task,
+                                      args.overflow_check)
     GraphBuilder.to_json(output_path, export_config)
     logger.info(f'Model graphs compared successfully, the result file is saved in {output_path}')
 
 
-def _build_graph(dump_path, out_path, overflow_check=False, output_file_name=f'build_{current_time}.vis'):
+def _build_graph(dump_path, args):
     logger.info('Start building model graph...')
     construct_path = FileChecker(os.path.join(dump_path, GraphConst.CONSTRUCT_FILE), FileCheckConst.FILE,
                                  FileCheckConst.READ_ABLE).common_check()
@@ -92,14 +98,19 @@ def _build_graph(dump_path, out_path, overflow_check=False, output_file_name=f'b
                             FileCheckConst.READ_ABLE).common_check()
     stack_path = FileChecker(os.path.join(dump_path, GraphConst.STACK_FILE), FileCheckConst.FILE,
                              FileCheckConst.READ_ABLE).common_check()
-    create_directory(out_path)
-    output_path = os.path.join(out_path, output_file_name)
-    graph = GraphBuilder.build(construct_path, data_path, stack_path)
+    graph = GraphBuilder.build(construct_path, data_path, stack_path, complete_stack=args.complete_stack)
     micro_steps = graph.paging_by_micro_step()
     # 开启溢出检测
-    if overflow_check:
+    if args.overflow_check:
         graph.overflow_check()
-    GraphBuilder.to_json(output_path, GraphExportConfig(graph, micro_steps=micro_steps))
+    return BuildGraphResult(graph, micro_steps)
+
+
+def _export_build_graph_result(out_path, graph, micro_steps, overflow_check,
+                               output_file_name=f'build_{current_time}.vis'):
+    create_directory(out_path)
+    output_path = os.path.join(out_path, output_file_name)
+    GraphBuilder.to_json(output_path, GraphExportConfig(graph, micro_steps=micro_steps, overflow_check=overflow_check))
     logger.info(f'Model graph built successfully, the result file is saved in {output_path}')
 
 
@@ -111,12 +122,33 @@ def _compare_graph_ranks(input_param, args, step=None):
     if npu_ranks != bench_ranks:
         logger.error('The number of ranks in the two runs are different. Unable to match the ranks.')
         raise CompareException(CompareException.INVALID_PATH_ERROR)
+    compare_graph_results = []
     for nr, br in zip(npu_ranks, bench_ranks):
         logger.info(f'Start processing data for {nr}...')
         input_param['npu_path'] = os.path.join(dump_rank_n, nr)
         input_param['bench_path'] = os.path.join(dump_rank_b, br)
         output_file_name = f'compare_{step}_{nr}_{current_time}.vis' if step else f'compare_{nr}_{current_time}.vis'
-        _compare_graph(input_param, args, output_file_name=output_file_name)
+        result = _compare_graph(input_param, args)
+        result.output_file_name = output_file_name
+        if nr != Const.RANK:
+            try:
+                result.rank = int(nr.replace(Const.RANK, ""))
+            except Exception as e:
+                logger.error('The folder name format is incorrect, expected rank+number.')
+                raise CompareException(CompareException.INVALID_PATH_ERROR) from e
+        # 暂存所有rank的graph，用于匹配rank间的分布式节点
+        compare_graph_results.append(result)
+
+    # 匹配rank间的分布式节点
+    if len(compare_graph_results) > 1:
+        DistributedAnalyzer({obj.rank: obj.graph_n for obj in compare_graph_results},
+                            args.overflow_check).distributed_match()
+        DistributedAnalyzer({obj.rank: obj.graph_b for obj in compare_graph_results},
+                            args.overflow_check).distributed_match()
+
+    for result in compare_graph_results:
+        _export_compare_graph_result(args, [result.graph_n, result.graph_b], result.graph_comparator,
+                                     result.micro_steps, output_file_name=result.output_file_name)
 
 
 def _compare_graph_steps(input_param, args):
@@ -138,21 +170,38 @@ def _compare_graph_steps(input_param, args):
         _compare_graph_ranks(input_param, args, step=folder_step)
 
 
-def _build_graph_ranks(dump_ranks_path, out_path, overflow_check=False, step=None):
+def _build_graph_ranks(dump_ranks_path, args, step=None):
     ranks = sorted(check_and_return_dir_contents(dump_ranks_path, Const.RANK))
+    build_graph_results = []
     for rank in ranks:
         logger.info(f'Start processing data for {rank}...')
         dump_path = os.path.join(dump_ranks_path, rank)
         output_file_name = f'build_{step}_{rank}_{current_time}.vis' if step else f'build_{rank}_{current_time}.vis'
-        _build_graph(dump_path, out_path, overflow_check, output_file_name)
+        result = _build_graph(dump_path, args)
+        result.output_file_name = output_file_name
+        if rank != Const.RANK:
+            try:
+                result.rank = int(rank.replace(Const.RANK, ""))
+            except Exception as e:
+                logger.error('The folder name format is incorrect, expected rank+number.')
+                raise CompareException(CompareException.INVALID_PATH_ERROR) from e
+        build_graph_results.append(result)
+
+    if len(build_graph_results) > 1:
+        DistributedAnalyzer({obj.rank: obj.graph for obj in build_graph_results},
+                            args.overflow_check).distributed_match()
+
+    for result in build_graph_results:
+        _export_build_graph_result(args.output_path, result.graph, result.micro_steps, args.overflow_check,
+                                   result.output_file_name)
 
 
-def _build_graph_steps(dump_steps_path, out_path, overflow_check=False):
+def _build_graph_steps(dump_steps_path, args):
     steps = sorted(check_and_return_dir_contents(dump_steps_path, Const.STEP))
     for step in steps:
         logger.info(f'Start processing data for {step}...')
         dump_ranks_path = os.path.join(dump_steps_path, step)
-        _build_graph_ranks(dump_ranks_path, out_path, overflow_check, step)
+        _build_graph_ranks(dump_ranks_path, args, step)
 
 
 def _graph_service_parser(parser):
@@ -161,14 +210,17 @@ def _graph_service_parser(parser):
     parser.add_argument("-o", "--output_path", dest="output_path", type=str,
                         help="<Required> The compare task result out path.", required=True)
     parser.add_argument("-lm", "--layer_mapping", dest="layer_mapping", type=str,
-                        help="<optional> The layer mapping file path.", required=False)
+                        help="<Optional> The layer mapping file path.", required=False)
     parser.add_argument("-oc", "--overflow_check", dest="overflow_check", action="store_true",
                         help="<Optional> whether open overflow_check for graph.", 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("-cs", "--complete_stack", dest="complete_stack", action="store_true",
+                        help="<Optional> Whether to use complete stack information.", required=False)
 
 
 def _graph_service_command(args):
-    with FileOpen(args.input_path, "r") as file:
-        input_param = json.load(file)
+    input_param = load_json(args.input_path)
     npu_path = input_param.get("npu_path")
     bench_path = input_param.get("bench_path")
     check_file_or_directory_path(npu_path, isdir=True)
@@ -177,11 +229,12 @@ def _graph_service_command(args):
     if check_file_type(npu_path) == FileCheckConst.DIR and not bench_path:
         content = check_directory_content(npu_path)
         if content == GraphConst.RANKS:
-            _build_graph_ranks(npu_path, args.output_path, args.overflow_check)
+            _build_graph_ranks(npu_path, args)
         elif content == GraphConst.STEPS:
-            _build_graph_steps(npu_path, args.output_path, args.overflow_check)
+            _build_graph_steps(npu_path, args)
         else:
-            _build_graph(npu_path, args.output_path, args.overflow_check)
+            result = _build_graph(npu_path, args)
+            _export_build_graph_result(args.output_path, result.graph, result.micro_steps, args.overflow_check)
     elif check_file_type(npu_path) == FileCheckConst.DIR and check_file_type(bench_path) == FileCheckConst.DIR:
         content_n = check_directory_content(npu_path)
         content_b = check_directory_content(bench_path)
@@ -192,7 +245,9 @@ def _graph_service_command(args):
         elif content_n == GraphConst.STEPS:
             _compare_graph_steps(input_param, args)
         else:
-            _compare_graph(input_param, args)
+            result = _compare_graph(input_param, args)
+            _export_compare_graph_result(args, [result.graph_n, result.graph_b],
+                                         result.graph_comparator, result.micro_steps)
     else:
         logger.error("The npu_path or bench_path should be a folder.")
         raise CompareException(CompareException.INVALID_COMPARE_MODE)
@@ -212,3 +267,21 @@ def _ms_graph_service_parser(parser):
 
 def _ms_graph_service_command(args):
     _graph_service_command(args)
+
+
+class CompareGraphResult:
+    def __init__(self, graph_n, graph_b, graph_comparator, micro_steps, rank=0, output_file_name=''):
+        self.graph_n = graph_n
+        self.graph_b = graph_b
+        self.graph_comparator = graph_comparator
+        self.micro_steps = micro_steps
+        self.rank = rank
+        self.output_file_name = output_file_name
+
+
+class BuildGraphResult:
+    def __init__(self, graph, micro_steps, rank=0, output_file_name=''):
+        self.graph = graph
+        self.micro_steps = micro_steps
+        self.rank = rank
+        self.output_file_name = output_file_name