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

msprobe/visualization/builder/graph_builder.py

@@ -14,26 +14,29 @@
 # 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
+from msprobe.visualization.db_utils import node_to_db, config_to_db
 
 
 class GraphBuilder:
     backward_pattern = re.compile(r"(\.backward\.)(\d+)$")
     forward_pattern = re.compile(r"(\.forward\.)(\d+)$")
     # 匹配以大写字母开头，后接任意字母，并以Template(结尾，或包含api_template(的字符串
-    template_pattern = re.compile(r'\b([A-Z][a-zA-Z]*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:
@@ -41,75 +44,38 @@ 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 = {}
+    def to_db(filename, config):
+        config.graph_n.step = config.step
+        config.graph_n.rank = config.rank
+        config.graph_n.compare_mode = config.compare_mode
+        node_to_db(config.graph_n, filename)
         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 _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
+            config.graph_b.data_source = GraphConst.JSON_BENCH_KEY
+            config.graph_b.step = config.step
+            config.graph_b.rank = config.rank
+            node_to_db(config.graph_b, filename)
+        config_to_db(config, filename)
 
     @staticmethod
     def _handle_backward_upnode_missing(construct_dict, subnode_id, upnode_id):
@@ -132,10 +98,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)
@@ -171,6 +174,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
@@ -230,6 +235,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)
@@ -276,10 +283,128 @@ 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='',
-                 overflow_check=False, compare_mode=None):
+                 overflow_check=False, compare_mode=None, step=0, rank=0, step_list=None, rank_list=None):
         self.graph_n = graph_n
         self.graph_b = graph_b
         self.tool_tip = tool_tip
@@ -288,6 +413,10 @@ class GraphExportConfig:
         self.task = task
         self.overflow_check = overflow_check
         self.compare_mode = compare_mode
+        self.step = step
+        self.rank = rank
+        self.step_list = step_list
+        self.rank_list = rank_list
 
 
 @dataclass