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

msprobe/visualization/builder/graph_builder.py

@@ -14,14 +14,14 @@
 # limitations under the License.
 
 import re
+import copy
 from dataclasses import dataclass
 
 from msprobe.core.common.const import Const
-from msprobe.core.common.file_utils import load_json, save_json
+from msprobe.core.common.file_utils import load_json, load_construct_json
 from msprobe.core.common.utils import load_stack_json
 from msprobe.core.common.log import logger
 from msprobe.visualization.builder.msprobe_adapter import get_input_output
-from msprobe.visualization.builder.msprobe_adapter import op_patterns
 from msprobe.visualization.graph.graph import Graph
 from msprobe.visualization.graph.node_op import NodeOp
 from msprobe.visualization.utils import GraphConst
@@ -33,9 +33,10 @@ class GraphBuilder:
     forward_pattern = re.compile(r"(\.forward\.)(\d+)$")
     # 匹配以大写字母开头，后接任意字母，并以Template(结尾，或包含api_template(的字符串
     template_pattern = re.compile(r'\b([A-Z][a-zA-Z]*Template|api_template|api_instance)\(')
+    micro_step_dict = {}
 
     @staticmethod
-    def build(construct_path, data_path, stack_path, model_name='DefaultModel', complete_stack=False):
+    def build(construct_path, data_path, stack_path, model_name='DefaultModel'):
         """
         GraphBuilder的对外提供的构图方法
         Args:
@@ -43,48 +44,26 @@ class GraphBuilder:
             data_path: dump.json路径
             stack_path: stack.json路径
             model_name: 模型名字，依赖外部输入
-            complete_stack: 完整的堆栈信息
         Returns: Graph，代表图的数据结构
         """
-        construct_dict = load_json(construct_path)
+        construct_dict, micro_step_dict = load_construct_json(construct_path)
         if not construct_dict:
             logger.error("The content of 'construct.json' is empty, failed to build graph. "
                          "When dumping data, it is necessary to select level L0 or mix in order to "
                          "collect model structure data, that is, the content of 'construct.json' is not empty.")
             raise RuntimeError
+        GraphBuilder.micro_step_dict = micro_step_dict
         dump_dict = load_json(data_path)
         stack_dict = load_stack_json(stack_path)
-        if not complete_stack:
-            GraphBuilder._simplify_stack(stack_dict)
         data_dict = dump_dict.get(GraphConst.DATA_KEY, {})
-        graph = Graph(model_name, data_path=dump_dict.get('dump_data_dir', ''), dump_data=data_dict)
+        graph = Graph(model_name, data_path=dump_dict.get('dump_data_dir', ''), dump_data=data_dict,
+                      micro_step_num=micro_step_dict.get(Const.MEGATRON_MICRO_STEP_NUMBER))
         GraphBuilder._init_nodes(graph, construct_dict, data_dict, stack_dict)
+        GraphBuilder._handle_recompute(graph)
         GraphBuilder._collect_apis_between_modules(graph)
         GraphBuilder._add_parameters_grad(graph, data_dict)
         return graph
 
-    @staticmethod
-    def to_json(filename, config):
-        """
-        将graph导出成.vis文件的接口
-        """
-        result = {}
-        if config.graph_b:
-            result[GraphConst.JSON_NPU_KEY] = config.graph_n.to_dict(config.compare_mode)
-            result[GraphConst.JSON_BENCH_KEY] = config.graph_b.to_dict(config.compare_mode)
-        else:
-            result = config.graph_n.to_dict(config.compare_mode)
-        if config.tool_tip:
-            result[GraphConst.JSON_TIP_KEY] = config.tool_tip
-        if config.node_colors:
-            result[GraphConst.COLORS] = config.node_colors
-        if config.micro_steps:
-            result[GraphConst.MICRO_STEPS] = config.micro_steps
-        if config.task:
-            result[GraphConst.JSON_TASK_KEY] = config.task
-        result[GraphConst.OVERFLOW_CHECK] = config.overflow_check
-        save_json(filename, result, indent=4)
-
     @staticmethod
     def to_db(filename, config):
         config.graph_n.step = config.step
@@ -95,42 +74,10 @@ class GraphBuilder:
             config.graph_b.data_source = GraphConst.JSON_BENCH_KEY
             config.graph_b.step = config.step
             config.graph_b.rank = config.rank
+            config.graph_b.compare_mode = config.compare_mode
             node_to_db(config.graph_b, filename)
         config_to_db(config, filename)
 
-    @staticmethod
-    def _simplify_stack(stack_dict):
-        """
-        精简堆栈内容，模块级保留包含"模块名("的堆栈，api级保留"xxxTemplate("的下一行堆栈
-
-        例如模块 Module.layer3.0.bn2.BatchNorm2d.forward.0，模块名为bn2，匹配"bn2("，
-        保留堆栈"File /home/models/resnet.py, line 97, in forward, \n out = self.bn2(out)"
-
-        例如Api Tensor.__iadd__.4.forward，堆栈为：
-        "File /home/wrap_tensor.py, line 61,  return TensorOPTemplate(op_name, hook)(*args, **kwargs)",
-        "File /home/torchvision/models/resnet.py, line 102, in forward, \n out += identity",
-        匹配到第一行的"TensorOPTemplate("，保留下一行堆栈
-        """
-        module_pattern = re.compile(op_patterns[0])
-        for dump_name, stack_list in stack_dict.items():
-            if not isinstance(stack_list, list):
-                continue
-            if module_pattern.match(dump_name):
-                parts = dump_name.split(Const.SEP)
-                if len(parts) < abs(Const.LAYER_NAME_INDEX):
-                    continue
-                module_name = parts[Const.LAYER_NAME_INDEX]
-                for stack in stack_list:
-                    if re.search(module_name + r'\(', stack):
-                        stack_list = [stack]
-                        break
-            else:
-                for index, stack in enumerate(stack_list):
-                    if GraphBuilder.template_pattern.search(stack) and index < len(stack_list) - 1:
-                        stack_list = [stack_list[index + 1]]
-                        break
-            stack_dict[dump_name] = stack_list
-
     @staticmethod
     def _handle_backward_upnode_missing(construct_dict, subnode_id, upnode_id):
         """
@@ -152,10 +99,47 @@ class GraphBuilder:
                     return new_upnode_id
         return upnode_id
 
+    @staticmethod
+    def _handle_backward_inplace(construct_dict, sub_node_id, up_node_id):
+        """
+        如果当前backward节点的父层级信息不等于其父级节点的层级信息，则尝试从同名的forward节点寻找父级节点
+        主要针对的场景：inplace层会无法触发backward hook导致反向层级错误
+
+        example:
+            正确的层级关系：
+                父层：Module.layer4.1.BasicBlock.backward.0的层级信息为Module.layer4.1
+                子层：Module.layer4.1.conv2.Conv2d.backward.0的父层级信息为Module.layer4.1
+
+            错误的层级关系：
+                父层：Module.layer4.1.relu.ReLU.backward.1的层级信息为Module.layer4.1.relu
+                子层：Module.layer4.1.conv2.Conv2d.backward.0的父层级信息为Module.layer4.1
+        """
+        if GraphBuilder.backward_pattern.search(sub_node_id) and up_node_id:
+            sub_split = sub_node_id.split(Const.SEP)
+            if len(sub_split) < 5:
+                return up_node_id
+            up_split = up_node_id.split(Const.SEP)
+            if len(up_split) < 4:
+                return up_node_id
+            sub_node_prefix = Const.SEP.join(sub_split[:-4])
+            up_node_prefix = Const.SEP.join(up_split[:-3])
+            if sub_node_prefix != up_node_prefix:
+                forward_sub_node_id = GraphBuilder.backward_pattern.sub(r".forward.\2", sub_node_id)
+                if forward_sub_node_id in construct_dict:
+                    forward_up_node_id = construct_dict.get(forward_sub_node_id)
+                    # forward_up_node_id ---> null
+                    if not forward_up_node_id:
+                        return forward_up_node_id
+                    new_up_node_id = GraphBuilder.forward_pattern.sub(r".backward.\2", forward_up_node_id)
+                    if new_up_node_id in construct_dict:
+                        return new_up_node_id
+        return up_node_id
+
     @staticmethod
     def _init_nodes(graph, construct_dict, data_dict, stack_dict):
         for subnode_id, upnode_id in construct_dict.items():
-            upnode_id = GraphBuilder._handle_backward_upnode_missing(construct_dict, subnode_id, upnode_id)
+            upnode_id = GraphBuilder._handle_backward_inplace(construct_dict, subnode_id, upnode_id) if upnode_id \
+                else GraphBuilder._handle_backward_upnode_missing(construct_dict, subnode_id, upnode_id)
             if upnode_id:
                 upnode_op = NodeOp.get_node_op(upnode_id)
                 upnode = GraphBuilder._create_or_get_node(graph, [data_dict, stack_dict], upnode_op, upnode_id)
@@ -191,6 +175,8 @@ class GraphBuilder:
                 node_stack_info = forward_node.stack_info if forward_node \
                     else ['This backward node cannot find the forward node and cannot retrieve stack information.']
             node.stack_info = node_stack_info
+            if GraphBuilder.micro_step_dict:
+                node.micro_step_id = GraphBuilder.micro_step_dict.get(node.id, 0)
         # 添加节点
         node.add_upnode(upnode)
         return node
@@ -250,6 +236,8 @@ class GraphBuilder:
                         node.upnode = api_collection_node
                     api_collection_node.upnode = graph.root
                     output.append(api_collection_node)
+                    if temp_nodes[0].micro_step_id is not None:
+                        api_collection_node.micro_step_id = temp_nodes[0].micro_step_id
                 else:
                     # 如果连续的api节点不足2个，将它们原样添加到输出列表
                     output.extend(temp_nodes)
@@ -296,6 +284,124 @@ class GraphBuilder:
                 # 更新数据
                 graph.get_node(parameters_grad_node_id).set_input_output(input_data, output_data)
 
+    @staticmethod
+    def _handle_recompute(graph):
+        """
+        1. 通过_get_recompute_map获得重计算节点映射recompute_map: dict(node_id: node_id_prefix)
+        2. 通过_get_no_recompute_map获得非重计算节点映射no_recompute_map: dict(node_id_prefix: list(node_id))
+        3. 遍历recompute_map，通过node_id_prefix与no_recompute_map建立连接，通过非重计算节点找到自身的父节点
+        """
+        recompute_map, recompute_id_map = GraphBuilder._get_recompute_map(graph.root.subnodes)
+        if not recompute_map:
+            return
+        id_prefixes = set(recompute_map.values())
+        no_recompute_map = GraphBuilder._get_no_recompute_map(graph, id_prefixes)
+        if not no_recompute_map:
+            return
+        # 深拷贝非重计算节点字典用于反向模式
+        no_recompute_ids_b = copy.deepcopy(no_recompute_map)
+
+        del_indexes = []
+        for node_id, id_prefix in recompute_map.items():
+            if id_prefix not in no_recompute_map:
+                continue
+            node_list = no_recompute_map.get(id_prefix) if GraphBuilder.forward_pattern.search(node_id) else \
+                no_recompute_ids_b.get(id_prefix)
+            if not node_list:
+                continue
+            no_recompute_node = node_list.pop()
+            recompute_node = graph.node_map.get(node_id)
+            if not recompute_node:
+                continue
+            # 通过非重计算forward节点的父节点，找到对应的backward父节点
+            new_up_node = graph.node_map.get(
+                GraphBuilder.forward_pattern.sub(r".backward.\2", no_recompute_node.upnode.id))
+            if not new_up_node:
+                continue
+
+            # 更新节点连接关系
+            recompute_node.upnode = new_up_node
+            new_up_node.subnodes.append(recompute_node)
+
+            del_indexes.append(recompute_id_map.get(node_id))
+
+        # 从后往前删除graph首层中已更新父节点的重计算节点
+        del_indexes.sort(reverse=True)
+        for index in del_indexes:
+            if 0 <= index <= len(graph.root.subnodes):
+                del graph.root.subnodes[index]
+
+    @staticmethod
+    def _get_recompute_map(node_list: list):
+        """
+        找到graph首层的重计算层
+
+        return: dict(node_id: node_id_prefix), dict(node_id: index)
+
+        example:
+        {Module.0.module.decoder.layers.0.TransformerLayer.forward.4: Module.0.module.decoder.layers.0.TransformerLayer}
+        """
+        recompute_map = {}
+        recompute_id_map = {}
+        node_id_set = set([node.id for node in node_list])
+        node_id_cache = set()
+        for i, node in enumerate(node_list):
+            if NodeOp.get_node_op(node.id) != NodeOp.module:
+                continue
+            id_segments = node.id.split(Const.SEP)
+            prefix = Const.SEP.join(id_segments[:-2])
+            if node.id in node_id_cache:
+                recompute_map[node.id] = prefix
+                recompute_id_map[node.id] = i
+                continue
+            is_recompute = GraphBuilder._is_recompute_node_id(id_segments)
+            if not is_recompute:
+                continue
+            # 重计算层必然是一组对应的前反向节点
+            id_segments[-2] = Const.BACKWARD if id_segments[-2] == Const.FORWARD else Const.FORWARD
+            relative_node_id = Const.SEP.join(id_segments)
+            if relative_node_id in node_id_set:
+                recompute_map[node.id] = prefix
+                recompute_id_map[node.id] = i
+                # 对应节点id放入缓存避免后续重复判断
+                node_id_cache.add(relative_node_id)
+        return recompute_map, recompute_id_map
+
+    @staticmethod
+    def _is_recompute_node_id(id_segments):
+        """
+        非重计算首层节点命名必然是：Module/Cell.{number(可选)}.module_name.{number(可选)}.class_name.forward/backward.number
+        如果不符合，则判断为重计算节点
+        """
+        if len(id_segments) > 7:
+            return True
+        if len(id_segments) == 7 and not (id_segments[1].isdigit() and id_segments[3].isdigit()):
+            return True
+        if len(id_segments) == 6 and not id_segments[1].isdigit():
+            return True
+        return False
+
+    @staticmethod
+    def _get_no_recompute_map(graph, recompute_id_prefixes):
+        """
+        寻找与重计算层id前缀相同的非重计算forward层，按顺序排列，重计算层按照顺序使用非重计算forward层的父节点对应的backward节点
+
+        return: dict(node_id_prefix: list(node_id))
+        """
+        no_recompute_map = {}
+        for node_id, node in graph.node_map.items():
+            if NodeOp.get_node_op(node_id) == NodeOp.module and GraphBuilder.forward_pattern.search(node_id):
+                if not node.upnode or node.upnode.id == graph.root.id:
+                    continue
+                id_prefix = GraphBuilder.forward_pattern.sub('', node_id)
+                if id_prefix not in recompute_id_prefixes:
+                    continue
+                no_recompute_map.setdefault(id_prefix, []).append(node)
+        for node_list in no_recompute_map.values():
+            # 方便按顺序pop弹出
+            node_list.reverse()
+        return no_recompute_map
+
 
 class GraphExportConfig:
     def __init__(self, graph_n, graph_b=None, tool_tip=None, node_colors=None, micro_steps=None, task='',

msprobe/visualization/builder/graph_merger.py

@@ -212,7 +212,6 @@ class BaseGraphMerger:
                     if compare_data and not self.compare_param_same(main_param, other_param, has_uncertainty=True):
                         same_flag = False
                 if not same_flag:
-                    # {input.0: [{"Max": 0, "Min": 0, ...}, {"Max": 0.1, "Min": 0, ...}, ...]}
                     data_dict[key.replace(main_node.id + Const.SEP, '')] = tp_need_merge_params
         return data_types.get('input_data'), data_types.get('output_data')
 

msprobe/visualization/builder/msprobe_adapter.py

@@ -28,7 +28,7 @@ op_patterns = [
     # NodeOp.module
     r'^(Module.|Cell.|optimizer|clip_grad)',
     # NodeOp.function_api
-    r'^(Tensor.|Torch.|Functional.|NPU.|VF.|Distributed.|Aten.|Mint.|Primitive.|Jit.|MintFunctional.)'
+    r'^(Tensor.|Torch.|Functional.|NPU.|VF.|Distributed.|Aten.|Mint.|Primitive.|Jit.|MintFunctional.|MindSpeed.)'
 ]
 
 

msprobe/visualization/db_utils.py

@@ -41,7 +41,7 @@ node_columns = {
     'overflow_level': TEXT,
     'micro_step_id': INTEGER_NOT_NULL,
     'matched_node_link': TEXT,
-    'stack_info': TEXT,
+    'stack_id': TEXT,
     'parallel_merge_info': TEXT,
     'matched_distributed': TEXT,
     'modified': INTEGER_NOT_NULL,
@@ -65,6 +65,11 @@ config_columns = {
     'step_list': TEXT_NOT_NULL
 }
 
+stack_columns = {
+    'id': TEXT_PRIMARY_KEY,
+    'stack_info': TEXT
+}
+
 indexes = {
     "index1": ["step", "rank", "data_source", "up_node", "node_order"],
     "index2": ["step", "rank", "data_source", "node_name"],
@@ -197,19 +202,24 @@ def node_to_db(graph, db_name):
     create_table_sql = create_table_sql_from_dict('tb_nodes', node_columns)
     insert_sql = create_insert_sql_from_dict('tb_nodes', node_columns)
     data = []
+    stack_dict = {}
     for i, node in enumerate(graph.get_sorted_nodes()):
+        stack_info_text = json.dumps(node.stack_info)
+        if stack_info_text not in stack_dict:
+            stack_dict[stack_info_text] = get_stack_unique_id(graph, stack_dict)
         data.append((get_node_unique_id(graph, node), get_graph_unique_id(graph), i, node.id, node.op.value,
                      node.upnode.id if node.upnode else '',
                      json.dumps([node.id for node in node.subnodes]) if node.subnodes else '',
                      node.data.get(GraphConst.JSON_INDEX_KEY), node.data.get(GraphConst.OVERFLOW_LEVEL),
                      node.micro_step_id if node.micro_step_id is not None else 0, json.dumps(node.matched_node_link),
-                     json.dumps(node.stack_info),
+                     stack_dict.get(stack_info_text),
                      json.dumps(node.parallel_merge_info) if node.parallel_merge_info else '',
                      json.dumps(node.matched_distributed), 0,
                      json.dumps(format_node_data(node.input_data, node.id, graph.compare_mode)),
                      json.dumps(format_node_data(node.output_data, node.id, graph.compare_mode)),
                      graph.data_source, graph.data_path, graph.step, graph.rank))
     to_db(db_name, create_table_sql, insert_sql, data)
+    stack_to_db(stack_dict, db_name)
 
 
 def config_to_db(config, db_name):
@@ -221,9 +231,22 @@ def config_to_db(config, db_name):
     to_db(db_name, create_table_sql, insert_sql, data)
 
 
+def stack_to_db(stack_dict, db_name):
+    create_table_sql = create_table_sql_from_dict('tb_stack', stack_columns)
+    insert_sql = create_insert_sql_from_dict('tb_stack', stack_columns)
+    data = []
+    for stack_info_text, unique_id in stack_dict.items():
+        data.append((unique_id, stack_info_text))
+    to_db(db_name, create_table_sql, insert_sql, data)
+
+
 def get_graph_unique_id(graph):
     return f'{graph.data_source}_{graph.step}_{graph.rank}'
 
 
 def get_node_unique_id(graph, node):
     return f'{get_graph_unique_id(graph)}_{node.id}'
+
+
+def get_stack_unique_id(graph, stack_dict):
+    return f'{get_graph_unique_id(graph)}_{len(stack_dict)}'

msprobe/visualization/graph/base_node.py

@@ -89,30 +89,6 @@ class BaseNode:
         self.matched_node_link = ancestors
         node.matched_node_link = ancestors
 
-    def to_dict(self, compare_mode=None):
-        """
-        输出数据
-        """
-        result = {
-            'id': self.id,
-            'node_type': self.op.value,
-            'output_data': format_node_data(self.output_data, self.id, compare_mode),
-            'input_data': format_node_data(self.input_data, self.id, compare_mode),
-            'upnode': self.upnode.id if self.upnode else 'None',
-            'subnodes': [node.id for node in self.subnodes],
-            'matched_node_link': self.matched_node_link,
-            'suggestions': self.suggestions,
-            'stack_info': self.stack_info
-        }
-        if self.micro_step_id is not None:
-            result['micro_step_id'] = self.micro_step_id
-        result['data'] = self.data
-        if self.matched_distributed:
-            result[GraphConst.MATCHED_DISTRIBUTED] = self.matched_distributed
-        if self.parallel_merge_info:
-            result['parallel_merge_info'] = self.parallel_merge_info
-        return result
-
     def get_ancestors(self):
         """
         获取节点所有祖先的列表

msprobe/visualization/graph/graph.py

@@ -22,7 +22,7 @@ from msprobe.core.common.decorator import recursion_depth_decorator
 
 
 class Graph:
-    def __init__(self, model_name, data_path='', dump_data=None):
+    def __init__(self, model_name, data_path='', dump_data=None, micro_step_num=None):
         self.node_map = {}
         self.node_id_map = {}
         self.add_node(NodeOp.module, model_name)
@@ -33,6 +33,7 @@ class Graph:
         self.step = 0
         self.rank = 0
         self.compare_mode = GraphConst.SUMMARY_COMPARE
+        self.micro_step_num = micro_step_num
 
     def __str__(self):
         infos = [f'{str(self.node_map.get(node_id))}' for node_id in self.node_map]
@@ -172,19 +173,6 @@ class Graph:
         """
         return self.node_map.get(node_id, None)
 
-    def to_dict(self, compare_mode=None):
-        """
-        用于数据输出
-        """
-        result = {}
-        result[GraphConst.JSON_ROOT_KEY] = self.root.id if self.root else 'None'
-        result[GraphConst.JSON_DATA_KEY] = self.data_path
-        result[GraphConst.JSON_NODE_KEY] = {}
-        for node_id in self.node_map:
-            info = self.node_map.get(node_id).to_dict(compare_mode)
-            result[GraphConst.JSON_NODE_KEY][node_id] = info
-        return result
-
     def paging_by_micro_step(self, graph_other=None):
         """
         给graph首层节点增加micro step标记，供前端分页展示，有助于在处理大规模图数据时进行优化和管理
@@ -203,6 +191,9 @@ class Graph:
             for sub_node in node.subnodes:
                 propagate_micro_step_id(sub_node)
 
+        if self.micro_step_num is not None:
+            return self.micro_step_num + 1
+
         batches_n = Graph.split_nodes_by_micro_step(self.root.subnodes)
         for batch_number, nodes in batches_n.items():
             for node in nodes: