From b1bfff34e02c9176090469a23815e28e4f8a0a5c Mon Sep 17 00:00:00 2001 From: Ye Wang <52801275+wangyems@users.noreply.github.com> Date: Fri, 31 Jul 2020 17:57:54 -0700 Subject: [PATCH] Support distill-bert fusion in transformers tool (#4631) * checkin attention * checkin embedlayer but cause invalid onnx model * resolve comments * fix comments * check return values * add version limit * fix comments * add warning --- .../tools/transformers/fusion_attention.py | 40 +++- .../tools/transformers/fusion_embedlayer.py | 211 +++++++++++------- 2 files changed, 160 insertions(+), 91 deletions(-) diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index f5e87c08f0..55d1a96626 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -178,19 +178,33 @@ class FusionAttention(Fusion): return (_, _, add_v, matmul_v) = v_nodes + is_distill = False; qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Div', 'MatMul'], [0, 0, 0, 0]) if qk_nodes is None: qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Mul', 'MatMul'], [0, 0, 0, 0]) if qk_nodes is None: - logger.debug("fuse_attention: failed to match qk path") - return - (_, add_qk, _, matmul_qk) = qk_nodes + qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'MatMul', 'Div'], [0, 0, 2, 0]) + is_distill = True + if qk_nodes is None: + logger.debug("fuse_attention: failed to match qk path") + return + + add_qk = None + matmul_qk = None + where_qk = None + if is_distill: + (_, where_qk, matmul_qk, _) = qk_nodes + else: + (_, add_qk, _, matmul_qk) = qk_nodes q_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [0, 0, 0, 0]) if q_nodes is None: - logger.debug("fuse_attention: failed to match q path") - return - (_, _, add_q, matmul_q) = q_nodes + q_nodes = self.model.match_parent_path(matmul_qk, ['Div', 'Transpose', 'Reshape', 'Add', 'MatMul'], [0, 0, 0, 0, 0]) + if q_nodes is None: + logger.debug("fuse_attention: failed to match q path") + return + add_q = q_nodes[-2] + matmul_q = q_nodes[-1] k_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, 0]) if k_nodes is None: @@ -203,9 +217,15 @@ class FusionAttention(Fusion): matmul_k = k_nodes[-1] # Note that Cast might be removed by OnnxRuntime so we match two patterns here. - _, mask_nodes, _ = self.model.match_parent_paths( - add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]), - (['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], output_name_to_node) + mask_nodes = None + if is_distill: + _, mask_nodes, _ = self.model.match_parent_paths( + where_qk, [(['Expand', 'Reshape', 'Equal'], [0, 0, 0]), + (['Cast', 'Expand', 'Reshape', 'Equal'], [0, 0, 0, 0])], output_name_to_node) + else : + _, mask_nodes, _ = self.model.match_parent_paths( + add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]), + (['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], output_name_to_node) if mask_nodes is None: logger.debug("fuse_attention: failed to match mask path") return @@ -228,4 +248,4 @@ class FusionAttention(Fusion): # Use prune graph to remove mask nodes since they are shared by all attention nodes. #self.nodes_to_remove.extend(mask_nodes) - self.prune_graph = True + self.prune_graph = True \ No newline at end of file diff --git a/onnxruntime/python/tools/transformers/fusion_embedlayer.py b/onnxruntime/python/tools/transformers/fusion_embedlayer.py index d2e7d9ee09..ddc833a3dd 100644 --- a/onnxruntime/python/tools/transformers/fusion_embedlayer.py +++ b/onnxruntime/python/tools/transformers/fusion_embedlayer.py @@ -41,60 +41,12 @@ class FusionEmbedLayerNoMask(Fusion): self.utils = FusionUtils(model) self.attention = None - def fuse(self, node, input_name_to_nodes, output_name_to_node): - if self.model.match_parent_path(node, ['Add', 'Gather'], [0, 0]) is None: - logger.debug("Failed to match path SkipLayerNormalization[0] <-- Add <-- Gather") - return - - self.attention = self.model.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False) - if self.attention is None: - # In case user disables attention fusion, check whether subgraph looks like Attention. - if node.output[0] not in input_name_to_nodes: - return - children = input_name_to_nodes[node.output[0]] - children_types = sorted([child.op_type for child in children]) - if children_types != ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization']: - logger.debug("No Attention like subgraph in children of SkipLayerNormalization") - return - - # Assume the order of embeddings are word_embedding + position_embedding + segment_embedding - normalize_node = node - word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 0]) - if word_embedding_path is None: - logger.info("Word embedding path is not found. Embed layer cannot be fused.") - return - add_node, word_embedding_gather = word_embedding_path - input_ids = word_embedding_gather.input[1] - - position_embedding_expand = None - position_embedding_shape = None - - position_embedding_path = self.model.match_parent_path(normalize_node, ['Reshape', 'Slice'], [1, 0]) - if position_embedding_path is not None: - _, position_embedding_weight_node = position_embedding_path - else: - position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand', 'Shape'], [1, 1, 1]) - if position_embedding_path is not None: - position_embedding_weight_node, position_embedding_expand, position_embedding_shape = position_embedding_path - else: - position_embedding_path = self.model.match_parent_path( - add_node, ['Gather', 'Expand', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [1, 1, 1, 1, 0, 0]) - if position_embedding_path is not None: - position_embedding_weight_node, position_embedding_expand, _, _, _, position_embedding_shape = position_embedding_path - else: - # Here we will not try to get exact match. Instead, we only try identify position embedding weights. - position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand'], [1, 1]) - if position_embedding_path is not None: - position_embedding_weight_node, position_embedding_expand = position_embedding_path - else: - logger.info("Position embedding path is not found. Embed layer cannot be fused.") - return - - if position_embedding_shape is not None and position_embedding_shape.input[0] != input_ids: - logger.info("position and word embedding is expected to be applied on same input") - return + def match_segment_path(self, normalize_node, input_name_to_nodes, output_name_to_node, input_ids_cast_node): + segment_ids = None + segment_embedding_gather = None segment_embedding_path = self.model.match_parent_path(normalize_node, ['Gather'], [1]) + if segment_embedding_path is None: segment_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 1]) if segment_embedding_path is None: @@ -106,6 +58,99 @@ class FusionEmbedLayerNoMask(Fusion): segment_ids = segment_embedding_gather.input[1] + self.nodes_to_remove.extend(segment_embedding_path) + + if self.model.find_graph_input(segment_ids): + casted, segment_ids = self.utils.cast_graph_input_to_int32(segment_ids) + else: + segment_ids, segment_ids_cast_node = self.utils.cast_input_to_int32(segment_ids) + + # Cast might be removed by OnnxRuntime. + _, segment_id_path, _ = self.model.match_parent_paths( + segment_ids_cast_node, + [(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], [0, 0, 1, 0, 0, 0]), + (['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], output_name_to_node) + + if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[0] == segment_id_path[-1].input[0]: + logger.debug("Simplify semgent id path...") + self.model.add_node( + helper.make_node('Shape', inputs=[input_ids_cast_node.input[0]], outputs=["input_shape"])) + self.model.add_node( + helper.make_node('ConstantOfShape', + inputs=["input_shape"], + outputs=["zeros_for_input_shape"], + value=helper.make_tensor("value", onnx.TensorProto.INT32, [1], [1]))) + segment_ids = "zeros_for_input_shape" + + return segment_ids, segment_embedding_gather + + def fuse(self, node, input_name_to_nodes, output_name_to_node): + is_distill = False; + + if self.model.match_parent_path(node, ['Add', 'Gather'], [0, 0]) is None and self.model.match_parent_path(node, ['Gather'], [0]) is None: + logger.debug("Failed to match path SkipLayerNormalization[0] <-- Add <-- Gather or SkipLayerNormalization[0] <-- Gather") + return + + self.attention = self.model.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False) + if self.attention is None: + # In case user disables attention fusion, check whether subgraph looks like Attention. + if node.output[0] not in input_name_to_nodes: + return + children = input_name_to_nodes[node.output[0]] + children_types = sorted([child.op_type for child in children]) + if children_types != ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization'] and children_types != ['MatMul', 'MatMul', 'MatMul', 'Shape', 'Shape', 'SkipLayerNormalization']: + logger.debug("No Attention like subgraph in children of SkipLayerNormalization") + return + + # Assume the order of embeddings are word_embedding + position_embedding + segment_embedding + normalize_node = node + add_node = None + word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 0]) + if word_embedding_path is not None: + add_node, word_embedding_gather = word_embedding_path + else: + word_embedding_path = self.model.match_parent_path(normalize_node, ['Gather'], [0]) + if word_embedding_path is not None: + word_embedding_gather = word_embedding_path[0] + is_distill = True; + else: + logger.info("Word embedding path is not found. Embed layer cannot be fused.") + return + + input_ids = word_embedding_gather.input[1] + + position_embedding_expand = None + position_embedding_shape = None + + position_embedding_path = self.model.match_parent_path(normalize_node, ['Gather', 'Expand'], [1, 1]) # for distill-bert + if position_embedding_path is not None: + position_embedding_weight_node, position_embedding_expand = position_embedding_path + else: + position_embedding_path = self.model.match_parent_path(normalize_node, ['Reshape', 'Slice'], [1, 0]) + if position_embedding_path is not None: + _, position_embedding_weight_node = position_embedding_path + else: + position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand', 'Shape'], [1, 1, 1]) + if position_embedding_path is not None: + position_embedding_weight_node, position_embedding_expand, position_embedding_shape = position_embedding_path + else: + position_embedding_path = self.model.match_parent_path( + add_node, ['Gather', 'Expand', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [1, 1, 1, 1, 0, 0]) + if position_embedding_path is not None: + position_embedding_weight_node, position_embedding_expand, _, _, _, position_embedding_shape = position_embedding_path + else: + # Here we will not try to get exact match. Instead, we only try identify position embedding weights. + position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand'], [1, 1]) + if position_embedding_path is not None: + position_embedding_weight_node, position_embedding_expand = position_embedding_path + else: + logger.info("Position embedding path is not found. Embed layer cannot be fused.") + return + + if position_embedding_shape is not None and position_embedding_shape.input[0] != input_ids: + logger.info("position and word embedding is expected to be applied on same input") + return + if position_embedding_expand and position_embedding_shape: input_parent = self.model.get_parent(position_embedding_shape, 0, output_name_to_node) subgraph_nodes = self.model.get_parent_subgraph_nodes(position_embedding_expand, @@ -115,7 +160,6 @@ class FusionEmbedLayerNoMask(Fusion): self.nodes_to_remove.extend(word_embedding_path) self.nodes_to_remove.extend(position_embedding_path) - self.nodes_to_remove.extend(segment_embedding_path) self.nodes_to_remove.extend([normalize_node]) @@ -126,41 +170,46 @@ class FusionEmbedLayerNoMask(Fusion): else: input_ids, input_ids_cast_node = self.utils.cast_input_to_int32(input_ids) - if self.model.find_graph_input(segment_ids): - casted, segment_ids = self.utils.cast_graph_input_to_int32(segment_ids) - else: - segment_ids, segment_ids_cast_node = self.utils.cast_input_to_int32(segment_ids) - - # Cast might be removed by OnnxRuntime. - _, segment_id_path, _ = self.model.match_parent_paths( - segment_ids_cast_node, - [(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], [0, 0, 1, 0, 0, 0]), - (['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], output_name_to_node) - - if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[0] == segment_id_path[-1].input[0]: - logger.debug("Simplify semgent id path...") - self.model.add_node( - helper.make_node('Shape', inputs=[input_ids_cast_node.input[0]], outputs=["input_shape"])) - self.model.add_node( - helper.make_node('ConstantOfShape', - inputs=["input_shape"], - outputs=["zeros_for_input_shape"], - value=helper.make_tensor("value", onnx.TensorProto.INT32, [1], [1]))) - segment_ids = "zeros_for_input_shape" - node_name = self.model.create_node_name('EmbedLayerNormalization') output_name = node_name + "_output" - embed_node = helper.make_node( - 'EmbedLayerNormalization', - inputs=[ + + embed_node_inputs = None + if is_distill == False: + segment_path = self.match_segment_path(normalize_node, input_name_to_nodes, output_name_to_node, input_ids_cast_node) + if segment_path is None: + return + else: + from packaging.version import Version + import onnxruntime + if Version(onnxruntime.__version__) <= Version("1.4.0"): + logger.warning('Please install onnxruntime with version > 1.4.0 for embedlayer fusion support for distilbert') + return + + segment_ids, segment_embedding_gather = segment_path + + embed_node_inputs=[ + input_ids, + segment_ids, + word_embedding_gather.input[0], + position_embedding_weight_node.input[0], + segment_embedding_gather.input[0], + normalize_node.input[2], + normalize_node.input[3] # gamma and beta + ] + else: + embed_node_inputs=[ input_ids, - segment_ids, + '', word_embedding_gather.input[0], position_embedding_weight_node.input[0], - segment_embedding_gather.input[0], + '', normalize_node.input[2], normalize_node.input[3] # gamma and beta - ], + ] + + embed_node = helper.make_node( + 'EmbedLayerNormalization', + embed_node_inputs, outputs=[node_name + "_output", node_name + "_dummy_mask_index"], name=node_name)