diff --git a/onnxruntime/python/tools/transformers/fusion_base.py b/onnxruntime/python/tools/transformers/fusion_base.py index 08b8ee2990..0469e90f48 100644 --- a/onnxruntime/python/tools/transformers/fusion_base.py +++ b/onnxruntime/python/tools/transformers/fusion_base.py @@ -25,6 +25,8 @@ class Fusion: self.prune_graph: bool = False self.node_name_to_graph_name: dict = {} self.this_graph_name: str = None + # It is optional that subclass updates fused_count since we will also check nodes_to_add to get counter. + self.fused_count: int = 0 def apply(self): logger.debug(f"start {self.description} fusion...") @@ -41,7 +43,7 @@ class Fusion: self.fuse(node, input_name_to_nodes, output_name_to_node) op_list = [node.op_type for node in self.nodes_to_add] - count = op_list.count(self.fused_op_type) + count = max(self.fused_count, op_list.count(self.fused_op_type)) if count > 0: logger.info(f"Fused {self.description} count: {count}") diff --git a/onnxruntime/python/tools/transformers/fusion_embedlayer.py b/onnxruntime/python/tools/transformers/fusion_embedlayer.py index dbf403128f..74b19942e7 100644 --- a/onnxruntime/python/tools/transformers/fusion_embedlayer.py +++ b/onnxruntime/python/tools/transformers/fusion_embedlayer.py @@ -3,9 +3,9 @@ # Licensed under the MIT License. #-------------------------------------------------------------------------- -from typing import Dict +from typing import Dict, List, Tuple, Union from logging import getLogger -from onnx import helper +from onnx import helper, TensorProto, NodeProto from onnx_model import OnnxModel from fusion_base import Fusion from fusion_utils import FusionUtils @@ -15,239 +15,375 @@ logger = getLogger(__name__) class FusionEmbedLayerNoMask(Fusion): """ - Embed Layer Normalization will fuse embeddings and mask processing into one node. - The embeddings before conversion: - - (input_ids) --------> Gather ----------+ (segment_ids) - | | | - | v v - +--> Shape --> Expand -> Gather---->Add Gather - | ^ | | - | | v v - +---(optional graph) SkipLayerNormalization - - Optional graph is used to generate position list (0, 1, ...) per batch. It can be a constant in some model. - - (input_ids) --> Gather -----+ Slice - | | - v v - (segment_ids)--> Gather --->Add Reshape - | | - v v - SkipLayerNormalization + Fuse embedding layer into one node (EmbedLayerNormalization). + It supports the following model types: BERT, DistilBert, ALBert. """ - def __init__(self, model: OnnxModel, description='no mask'): - super().__init__(model, "EmbedLayerNormalization", ["SkipLayerNormalization", "LayerNormalization"], + def __init__(self, model: OnnxModel, description: str = 'no mask'): + super().__init__(model, "EmbedLayerNormalization", ["LayerNormalization", "SkipLayerNormalization"], description) self.utils = FusionUtils(model) + self.shape_infer_helper = self.model.infer_runtime_shape({}, update=True) + # The following will be reset in each fuse call of FusionEmbedLayerNormalization self.attention = None + self.embed_node = None - 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 + def match_two_gather(self, add: NodeProto) -> Union[None, Tuple[NodeProto, NodeProto]]: + gather_0_path = self.model.match_parent_path(add, ['Gather'], [0]) + if gather_0_path is None: + return None - if normalize_node.op_type == "SkipLayerNormalization": - segment_embedding_path = self.model.match_parent_path(normalize_node, ['Gather'], [1]) + gather_1_path = self.model.match_parent_path(add, ['Gather'], [1]) + if gather_1_path is None: + return None - 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: - logger.info("Segment embedding is not found. Embed layer cannot be fused.") - return - _, segment_embedding_gather = segment_embedding_path - else: - segment_embedding_gather = segment_embedding_path[0] - elif normalize_node.op_type == "LayerNormalization": - segment_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Add', 'Gather'], [0, 0, 1]) + return gather_0_path[0], gather_1_path[0] - if segment_embedding_path is None: - logger.info("Segment embedding is not found. Embed layer cannot be fused.") - return - _, _, segment_embedding_gather = segment_embedding_path + def check_attention_subgraph(self, layernorm: NodeProto, input_name_to_nodes: Dict[str, List[NodeProto]], + is_distil_bert: bool) -> bool: + """Check that LayerNormalization has a child of Attention node or subgraph like Attention. - segment_ids = segment_embedding_gather.input[1] + Args: + layernorm (NodeProto): LayerNormalization node + input_name_to_nodes (Dict[str, List[NodeProto]]): map from input name to nodes + is_distil_bert (bool): whether it is DistilBert or not - 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.this_graph_name) - 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])), - self.this_graph_name) - 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" - ) - if node.op_type != "LayerNormalization" or self.model.match_parent_path(node, ['Add', 'Gather'], - [0, 1]) is None: - return - - self.attention = self.model.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False) + Returns: + bool: whether there is Attention node or subgraph like Attention + """ + self.attention = self.model.find_first_child_by_type(layernorm, + '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 + if layernorm.output[0] not in input_name_to_nodes: + return False + children = input_name_to_nodes[layernorm.output[0]] - # 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, ['Add', 'Add', 'Gather'], [0, 0, 0]) - if word_embedding_path is not None: - _, add_node, word_embedding_gather = word_embedding_path + # For Albert, there is MatMul+Add after embedding layer before attention. + if len(children) == 1 and children[0].op_type == "MatMul" and children[0].output[0] in input_name_to_nodes: + grandchildren = input_name_to_nodes[children[0].output[0]] + if len(grandchildren) == 1 and grandchildren[0].op_type == "Add" and grandchildren[0].output[ + 0] in input_name_to_nodes: + nodes = input_name_to_nodes[grandchildren[0].output[0]] + for node in nodes: + if node.op_type == "Attention": + self.attention = node + return True + children_types = sorted([child.op_type for child in nodes]) 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 - 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 - else: - logger.info("Word embedding path is not found. Embed layer cannot be fused.") - return + children_types = sorted([child.op_type for child in children]) + # Two Shape nodes might be merged by ORT + if is_distil_bert: + # SkipLayerNormailization might exist when model has been optimized by ORT first. + if children_types != ['MatMul', 'MatMul', 'MatMul', 'Shape', 'SkipLayerNormalization'] and \ + children_types != ['Add', 'MatMul', 'MatMul', 'MatMul', 'Shape', 'Shape'] and \ + children_types != ['Add', 'MatMul', 'MatMul', 'MatMul', 'Shape']: + logger.debug("No Attention like subgraph in children of LayerNormalization") + return False + else: + if children_types != ['Add', 'MatMul', 'MatMul', 'MatMul'] and \ + children_types != ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization']: + logger.debug("No Attention like subgraph in children of LayerNormalization") + return False + return True + + def match_position_embedding_distilbert(self, position_embedding_gather, input_ids, output_name_to_node): + """ Match position embedding path from input_ids to Gather for DistilBert. + + Pattern is like the following: + (input_ids) + | + Shape + | \ + | Gather (indices=1) + | | + | Cast (optional) + | | + | Range (start=0, end=*, delta=1) + | | + | Unsqueeze + | / + Expand + | + Gather + """ + path1 = self.model.match_parent_path(position_embedding_gather, ['Expand', 'Shape'], [1, 1]) + if path1 is None: + return False + + expand, shape = path1 + if shape.input[0] != input_ids: + return False + + _, path2, _ = self.model.match_parent_paths(expand, [(['Unsqueeze', 'Range', 'Cast', 'Gather', 'Shape'], [0, 0, 1, 0, 0]), \ + (['Unsqueeze', 'Range', 'Gather', 'Shape'], [0, 0, 1, 0])], output_name_to_node) + if path2 is None: + return False + + range_node = path2[1] + if not (self.utils.check_node_input_value(range_node, 0, 0) + and self.utils.check_node_input_value(range_node, 2, 1)): + return False + + gather_node = path2[-2] + if not (self.utils.check_node_input_value(gather_node, 1, 1)): + return False + + shape_node = path2[-1] + if shape_node.input[0] != input_ids: + return False + + return True + + def match_position_embedding_roberta(self, position_embedding_gather, input_ids, output_name_to_node): + """ Match position embedding path from input_ids to Gather for Roberta. + + Roberta Embedding Layer Pattern (* is optional since it might be removed by ORT, ? is the padding word id): + (input_ids) --> Equal(B=?) -- Not -- Cast(to=6) -- CumSum(axis=1) -- Mul -- Cast(to=7) -- Add(B=1) -- Cast(to=7)* --> Gather + | ^ + V | + +------------------------------+ + + Roberta new pattern from transformers v4.9: + (input_ids) --> Equal(B=?) -- Not -- Cast(to=6) -- CumSum(axis=1) -- Add(B=0) -- Mul -- Cast(to=7) -- Add(B=1) --> Gather + | ^ + V | + +-------------------------------------------+ + + start_node = position_embedding_gather + start_index = 1 + + # match optional Cast node. + parent = self.model.get_parent(start_node, start_index, output_name_to_node) + if parent is None: + return + if parent.op_type == "Cast": + if OnnxModel.get_node_attribute(parent, "to") != 7: + return + start_node = parent + start_index = 0 + + i, path, return_indices = self.model.match_parent_paths( + start_node, + [ (['Add', 'Cast', 'Mul', 'CumSum', 'Cast', 'Not', 'Equal'], [start_index, 0, 0, 0, 0, 0, 0]), + (['Add', 'Cast', 'Mul', 'Add', 'CumSum', 'Cast', 'Not', 'Equal'], [start_index, 0, 0, 0, 0, 0, 0, 0])], + output_name_to_node) + + if path is not None: + # constant input of Add shall be 1. + i, value = self.model.get_constant_input(path[0]) + if value != 1: + return False + + _, self.padding_word_id = self.model.get_constant_input(path[-1]) + + return input_ids == path[-1].input[0] + """ + + return False + + def match_position_embedding_bert(self, position_embedding_gather, input_ids, output_name_to_node): + """ Match position embedding path from input_ids to Gather for BERT. + + BERT Embedding Layer Pattern: + (input_ids) + / \ + / Shape + / | + / Gather (indices=1) + / | + / Add (optional, B=0) + / | + Gather (segment_ids) Unsqueeze (axes=0) + \ | | + \ Gather Slice (data[1,512], starts=0, ends=*, axes=1, steps=1) + \ / | + Add Gather + \ / + Add + | + LayerNormalization + """ + path = self.model.match_parent_path(position_embedding_gather, ['Slice', 'Unsqueeze'], [1, 2], + output_name_to_node) + if path is None: + return False + + slice, unsqueeze = path + slice_weight = self.model.get_constant_value(slice.input[0]) + if not (slice_weight is not None and len(slice_weight.shape) == 2 and slice_weight.shape[0] == 1 \ + and self.utils.check_node_input_value(slice, 1, [0]) \ + and self.utils.check_node_input_value(slice, 3, [1]) \ + and (len(slice.input) == 4 or self.utils.check_node_input_value(slice, 4, [1]))): + return False + + opset_version = self.model.get_opset_version() + if opset_version < 13: + if not FusionUtils.check_node_attribute(unsqueeze, 'axes', [0]): + return False + else: + if not self.utils.check_node_input_value(unsqueeze, 1, [0]): + return False + + node = self.model.get_parent(unsqueeze, 0, output_name_to_node) + if node is None: + return False + if node.op_type == "Add": + if not self.utils.check_node_input_value(node, 1, 0): + return False + gather = self.model.get_parent(node, 0, output_name_to_node) + else: + gather = node + + if gather is None or gather.op_type != "Gather": + return False + if not (self.utils.check_node_input_value(gather, 1, 1)): + return False + + shape = self.model.get_parent(gather, 0, output_name_to_node) + if shape is None or shape.op_type != "Shape": + return False + + return input_ids == shape.input[0] + + def match_position_embedding(self, position_embedding_gather, input_ids, output_name_to_node): + if self.match_position_embedding_bert(position_embedding_gather, input_ids, output_name_to_node): + return True + + # TODO: Support roberta (position starts from 2 instead of 0) in EmbedLayerNormalization kernel + # related: https://github.com/huggingface/transformers/issues/10736 + #if self.match_position_embedding_roberta(position_embedding_gather, input_ids, output_name_to_node): + # return True + + if self.match_position_embedding_distilbert(position_embedding_gather, input_ids, output_name_to_node): + return True + + return False + + def check_embedding(self, word_embedding_gather, segment_embedding_gather, position_embedding_gather): + """Sanity check of embedding weights, and match hidden_size of weights and shape of inputs. + """ input_ids = word_embedding_gather.input[1] + segment_ids = segment_embedding_gather.input[1] if segment_embedding_gather else None + position_ids = position_embedding_gather.input[1] - position_embedding_node_before_gather = None - position_embedding_shape = None + if self.shape_infer_helper is not None: + input_ids_shape = self.shape_infer_helper.get_edge_shape(input_ids) + position_ids_shape = self.shape_infer_helper.get_edge_shape(position_ids) + assert input_ids_shape and position_ids_shape + if not (len(input_ids_shape) == 2 and len(position_ids_shape) == 2 + and input_ids_shape[1] == position_ids_shape[1]): + logger.info( + "Cannot fuse EmbedLayerNormalization: input_ids and position_ids not matched in 2nd dimension: {} vs {}" + .format(input_ids_shape, position_ids_shape)) + return False - 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_node_before_gather = 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 + if segment_ids and not self.shape_infer_helper.compare_shape(input_ids, segment_ids): + logger.info( + "Cannot fuse EmbedLayerNormalization: input_ids and segment_ids does not have same shape: {} != {}". + format(input_ids_shape, self.shape_infer_helper.get_edge_shape(segment_ids))) + return False + + word_embedding_table = self.model.get_constant_value(word_embedding_gather.input[0]) + if word_embedding_table is None or len(word_embedding_table.shape) != 2: + logger.info("Cannot fuse EmbedLayerNormalization: word embedding table is not expected") + return False + + position_embedding_table = self.model.get_constant_value(position_embedding_gather.input[0]) + if position_embedding_table is None or len(position_embedding_table.shape) != 2 or ( + word_embedding_table.shape[1] != position_embedding_table.shape[1]): + logger.info("Cannot fuse EmbedLayerNormalization: position embedding table is not expected") + return False + + if segment_ids: + segment_embedding_table = self.model.get_constant_value(segment_embedding_gather.input[0]) + if segment_embedding_table is None or len(segment_embedding_table.shape) != 2 or ( + word_embedding_table.shape[1] != segment_embedding_table.shape[1]): + logger.info("Cannot fuse EmbedLayerNormalization: segment embedding table is not expected") + return False + + # In normal case, word embeding table is the largest, and segment embedding table is the smallest, while postion embedding table is in between. + # TODO: use other information (like initializer names) to identify different embedding weights automatically. + if word_embedding_table.shape[0] <= position_embedding_table.shape[0]: + logger.warn( + f"word_embedding_table ({word_embedding_gather.input[0]}) size {word_embedding_table.shape[0]} <= position_embedding_table ({position_embedding_gather.input[0]}) size {position_embedding_table.shape[0]}" + ) + + if segment_ids: + if word_embedding_table.shape[0] <= segment_embedding_table.shape[0]: + logger.warn( + f"word_embedding_table ({word_embedding_gather.input[0]}) size {word_embedding_table.shape[0]} <= segment_embedding_table ({segment_embedding_gather.input[0]}) size {segment_embedding_table.shape[0]}" + ) + + if position_embedding_table.shape[0] <= segment_embedding_table.shape[0]: + logger.warn( + f"position_embedding_table ({position_embedding_gather.input[0]}) size {position_embedding_table.shape[0]} <= segment_embedding_table ({segment_embedding_gather.input[0]}) size {segment_embedding_table.shape[0]}" + ) + + return True + + def cast_to_int32(self, input_name: str) -> Tuple[str, Union[None, NodeProto]]: + """Cast a graph input or node input to int32. + + Args: + input_name (str): name of graph input or node input + + Returns: + A tuple of casted input name and the cast node. + int32_output (str): If input is int32, it is the input name, Otherwise it is output name of Cast node. + input_cast_node (Union[None, NodeProto]): Cast node. It could be None if input is int32. + """ + input_cast_node = None + graph_input = self.model.find_graph_input(input_name) + if graph_input is not None: + if graph_input.type.tensor_type.elem_type != TensorProto.INT32: + int32_output, input_cast_node = self.utils.cast_input_to_int32(input_name) 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_node_before_gather, 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_node_before_gather, _, _, _, 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_node_before_gather = position_embedding_path - else: - position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Slice'], - [1, 1]) - if position_embedding_path is not None: - position_embedding_weight_node, position_embedding_node_before_gather = position_embedding_path - else: - position_embedding_path = self.model.match_parent_path( - normalize_node, ['Add', 'Gather', 'Slice'], [0, 1, 1]) - if position_embedding_path is not None: - _, position_embedding_weight_node, position_embedding_node_before_gather = 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_node_before_gather 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_node_before_gather, - [input_parent] if input_parent else [], - output_name_to_node) - self.nodes_to_remove.extend(subgraph_nodes) - - self.nodes_to_remove.extend(word_embedding_path) - self.nodes_to_remove.extend(position_embedding_path) - - self.nodes_to_remove.extend([normalize_node]) - - # Cast input_ids and segment_ids to int32. - input_ids_cast_node = None - if self.model.find_graph_input(input_ids): - casted, input_ids = self.utils.cast_graph_input_to_int32(input_ids) + int32_output = input_name else: - input_ids, input_ids_cast_node = self.utils.cast_input_to_int32(input_ids) + int32_output, input_cast_node = self.utils.cast_input_to_int32(input_name) + + return int32_output, input_cast_node + + def create_fused_node(self, input_ids: str, layernorm: NodeProto, word_embedding_gather: NodeProto, + position_embedding_gather: NodeProto, segment_embedding_gather: Union[None, NodeProto]): + """Create an EmbedLayerNormalization node. Note that segment embedding is optional. + + Args: + input_ids (str): input_ids for word embeddings + layernorm (NodeProto): LayerNormalization or SkipLayerNormalization node. + word_embedding_gather (NodeProto): the Gather node for word embedding + position_embedding_gather (NodeProto): the Gather node for position embedding + segment_embedding_gather (Union[None, NodeProto]): the Gather node for segment embedding, or None. + + Returns: + NodeProto: the EmbedLayerNormalization node created. + """ + nodes_to_add = [] + input_ids, _ = self.cast_to_int32(input_ids) node_name = self.model.create_node_name('EmbedLayerNormalization') - output_name = node_name + "_output" - if normalize_node.op_type == "LayerNormalization": - gamma = normalize_node.input[1] - beta = normalize_node.input[2] - elif normalize_node.op_type == "SkipLayerNormalization": - gamma = normalize_node.input[2] - beta = normalize_node.input[3] + if layernorm.op_type == "LayerNormalization": + gamma = layernorm.input[1] + beta = layernorm.input[2] + else: # SkipLayerNormalization + gamma = layernorm.input[2] + beta = layernorm.input[3] 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: - segment_ids, segment_embedding_gather = segment_path + if segment_embedding_gather is not None: + segment_ids, _ = self.cast_to_int32(segment_embedding_gather.input[1]) - embed_node_inputs = [ - input_ids, - segment_ids, - word_embedding_gather.input[0], - position_embedding_weight_node.input[0], - segment_embedding_gather.input[0], - gamma, - beta # gamma and beta - ] - else: embed_node_inputs = [ - input_ids, - '', - word_embedding_gather.input[0], - position_embedding_weight_node.input[0], - '', - gamma, - beta # gamma and beta + input_ids, segment_ids, word_embedding_gather.input[0], position_embedding_gather.input[0], + segment_embedding_gather.input[0], gamma, beta + ] + else: # no segment embedding + embed_node_inputs = [ + input_ids, '', word_embedding_gather.input[0], position_embedding_gather.input[0], '', gamma, beta ] embed_node = helper.make_node('EmbedLayerNormalization', @@ -258,7 +394,7 @@ class FusionEmbedLayerNoMask(Fusion): embed_node.domain = "com.microsoft" # Pass attribute "epsilon" from normalize node to EmbedLayerNormalization. - for att in normalize_node.attribute: + for att in layernorm.attribute: if att.name == 'epsilon': embed_node.attribute.extend([att]) # Set default value to 1e-12 if no attribute is found. @@ -266,9 +402,119 @@ class FusionEmbedLayerNoMask(Fusion): if len(embed_node.attribute) == 0: embed_node.attribute.extend([helper.make_attribute("epsilon", 1.0E-12)]) - self.model.replace_input_of_all_nodes(normalize_node.output[0], output_name) - self.nodes_to_add.append(embed_node) - self.node_name_to_graph_name[embed_node.name] = self.this_graph_name + # Make sure new EmbedLayerNormalization node is the last one in self.nodes_to_add. + nodes_to_add.append(embed_node) + for node in nodes_to_add: + self.node_name_to_graph_name[node.name] = self.this_graph_name + self.nodes_to_add.extend(nodes_to_add) + + self.embed_node = embed_node + return embed_node + + def finish_fusion(self, layernorm, embed_node): + self.model.replace_input_of_all_nodes(layernorm.output[0], embed_node.output[0]) + # use prune graph to remove nodes that is not needed + self.prune_graph = True + + def fuse_distilbert(self, layernorm, add_before_layernorm, input_name_to_nodes, output_name_to_node): + """Fuse embedding layer for DistilBert + Args: + layernorm (NodeProto): node of LayerNormalization or SkipLayerNormalization + add_before_layernorm (NodeProto): the Add node before LayerNormalization, or the SkipLayerNormalization itself + input_name_to_nodes (Dict[str, List[NodeProto]]): map from input name to nodes + output_name_to_node (Dict[str, List[NodeProto]]): map from output name to nodes + """ + + # DistilBert has no segment embedding, subgraph pattern is like + # input_ids + # | \ + # | (position_embedding_subgraph) + # | | + # Gather Gather + # \ / + # Add + # | + # LayerNormalization + two_gather = self.match_two_gather(add_before_layernorm) + if two_gather is None: + return False + + word_embedding_gather, position_embedding_gather = two_gather + input_ids = word_embedding_gather.input[1] + + if not self.check_attention_subgraph(layernorm, input_name_to_nodes, is_distil_bert=True): + return False + + if not self.match_position_embedding(position_embedding_gather, input_ids, output_name_to_node): + return False + + if not self.check_embedding(word_embedding_gather, None, position_embedding_gather): + return False + + embed_node = self.create_fused_node(input_ids, layernorm, word_embedding_gather, position_embedding_gather, + None) + self.finish_fusion(layernorm, embed_node) + return True + + def fuse_bert(self, layernorm, add_before_layernorm, input_name_to_nodes, output_name_to_node): + """Fuse embedding layer for Bert + Args: + layernorm (NodeProto): node of LayerNormalization or SkipLayerNormalization + add_before_layernorm (NodeProto): the Add node before LayerNormalization, or the SkipLayerNormalization itself + input_name_to_nodes (Dict[str, List[NodeProto]]): map from input name to nodes + output_name_to_node (Dict[str, List[NodeProto]]): map from output name to nodes + """ + + add_2_gather = self.model.match_parent_path(add_before_layernorm, ['Add'], [0]) + if add_2_gather is None: + return False + + two_gather = self.match_two_gather(add_2_gather[0]) + if two_gather is None: + return False + + word_embedding_gather, segment_embedding_gather = two_gather + + input_ids = word_embedding_gather.input[1] + + if not self.check_attention_subgraph(layernorm, input_name_to_nodes, is_distil_bert=False): + return False + + position_embedding_path = self.model.match_parent_path(add_before_layernorm, ['Gather'], [1]) + if position_embedding_path is None: + return False + + position_embedding_gather = position_embedding_path[0] + if not self.match_position_embedding(position_embedding_gather, input_ids, output_name_to_node): + if not self.match_position_embedding(segment_embedding_gather, input_ids, output_name_to_node): + return False + # position and segment are switched + temp = segment_embedding_gather + segment_embedding_gather = position_embedding_gather + position_embedding_gather = temp + + if not self.check_embedding(word_embedding_gather, segment_embedding_gather, position_embedding_gather): + return False + + embed_node = self.create_fused_node(input_ids, layernorm, word_embedding_gather, position_embedding_gather, + segment_embedding_gather) + self.finish_fusion(layernorm, embed_node) + return True + + def fuse(self, node, input_name_to_nodes, output_name_to_node): + if node.op_type == "LayerNormalization": + first_add_path = self.model.match_parent_path(node, ['Add'], [0]) + if first_add_path is None: + return + add_before_layernorm = first_add_path[0] + else: # SkipLayerNormalization + add_before_layernorm = node # Add is fused into SkipLayerNormalization + + if self.fuse_distilbert(node, add_before_layernorm, input_name_to_nodes, output_name_to_node): + return + + if self.fuse_bert(node, add_before_layernorm, input_name_to_nodes, output_name_to_node): + return class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): @@ -276,22 +522,18 @@ class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): super().__init__(model, "with mask") def fuse(self, node, input_name_to_nodes, output_name_to_node): - old_count = len(self.nodes_to_add) - + # Reset attention and embed_node so that we know fusion is successful when they are not None. + self.attention = None + self.embed_node = None super().fuse(node, input_name_to_nodes, output_name_to_node) - if len(self.nodes_to_add) == old_count: - return - if self.attention is not None: + if self.attention and self.embed_node: mask_index = self.attention.input[3] if mask_index in output_name_to_node: node = output_name_to_node[mask_index] if node.op_type == "ReduceSum": - embed_node = self.nodes_to_add.pop() + embed_node = self.embed_node mask_input_name = node.input[0] self.nodes_to_remove.extend([node]) embed_node.input.append(mask_input_name) embed_node.output[1] = mask_index - self.nodes_to_add.append(embed_node) - - self.prune_graph = True diff --git a/onnxruntime/python/tools/transformers/fusion_shape.py b/onnxruntime/python/tools/transformers/fusion_shape.py new file mode 100644 index 0000000000..ea2b278bf9 --- /dev/null +++ b/onnxruntime/python/tools/transformers/fusion_shape.py @@ -0,0 +1,100 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +#-------------------------------------------------------------------------- + +from fusion_base import Fusion +from logging import getLogger +from onnx import TensorProto, NodeProto +from onnx_model import OnnxModel +from fusion_utils import FusionUtils +from typing import Union, Dict, List + +logger = getLogger(__name__) + + +class FusionShape(Fusion): + def __init__(self, model: OnnxModel): + super().__init__(model, "Shape", "Concat") + self.utils = FusionUtils(model) + self.shape_infer = None + self.shape_infer_done = False + + def get_dimensions_from_tensor_proto(self, tensor_proto: TensorProto) -> Union[int, None]: + if tensor_proto.type.tensor_type.HasField('shape'): + return len(tensor_proto.type.tensor_type.shape.dim) + else: + return None + + def get_dimensions(self, input_name: str) -> Union[int, None]: + graph_input = self.model.find_graph_input(input_name) + if graph_input: + return self.get_dimensions_from_tensor_proto(graph_input) + + if not self.shape_infer_done: + self.shape_infer = self.model.infer_runtime_shape({}, update=True) + self.shape_infer_done = True + + if self.shape_infer is not None: + return self.get_dimensions_from_tensor_proto(self.shape_infer.known_vi_[input_name]) + + return None + + def fuse(self, concat_node: NodeProto, input_name_to_nodes: Dict[str, List[NodeProto]], + output_name_to_node: Dict[str, NodeProto]): + """ + Smplify subgraph like + + (2d_input) + / \ + Shape shape + / \ + Gather(indices=0) Gather(indices=1) + | | + Unsqueeze(axes=0) Unsqueeze(axes=0) + \ / + Concat + | + + into (2d_input) --> Shape --> + """ + opset_version = self.model.get_opset_version() + + inputs = len(concat_node.input) + root = None + shape_output = None + for i in range(inputs): + path = self.model.match_parent_path(concat_node, ['Unsqueeze', 'Gather', 'Shape'], [i, 0, 0], + output_name_to_node) + if path is None: + return + + unsqueeze, gather, shape = path + if i == 0: + shape_output = shape.output[0] + if root is None: + root = shape.input[0] + if self.get_dimensions(root) != inputs: + return + elif shape.input[0] != root: + return + + if not FusionUtils.check_node_attribute(unsqueeze, 'axis', 0, default_value=0): + return + + if opset_version < 13: + if not FusionUtils.check_node_attribute(unsqueeze, 'axes', [0]): + return + else: + if not self.utils.check_node_input_value(unsqueeze, 1, [0]): + return + + value = self.model.get_constant_value(gather.input[1]) + from numpy import ndarray, array_equal + if not (isinstance(value, ndarray) and value.size == 1 and value.item() == i): + return + + if self.model.find_graph_output(concat_node.output[0]) is None: + self.model.replace_input_of_all_nodes(concat_node.output[0], shape_output) + self.fused_count += 1 + self.prune_graph = True diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 5a036ac0f5..01cbe9ae83 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py @@ -19,7 +19,8 @@ class FusionSkipLayerNormalization(Fusion): """ def __init__(self, model: OnnxModel): super().__init__(model, "SkipLayerNormalization", "LayerNormalization") - self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}) + # Update shape inference is needed since other fusions might add new edge which does not have shape info yet. + self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True) def fuse(self, node, input_name_to_nodes, output_name_to_node): add = self.model.get_parent(node, 0, output_name_to_node) diff --git a/onnxruntime/python/tools/transformers/gpt2_helper.py b/onnxruntime/python/tools/transformers/gpt2_helper.py index 2a259ce152..b15340892d 100644 --- a/onnxruntime/python/tools/transformers/gpt2_helper.py +++ b/onnxruntime/python/tools/transformers/gpt2_helper.py @@ -786,7 +786,7 @@ class Gpt2Helper: logger.info(f"Removed the existed directory: {new_dir}") except OSError as e: logger.info(f"Failed to remove the directory {new_dir}: {e.strerror}") - + # store each model to its own directory (for external data format). return { "raw": os.path.join(os.path.join(output_dir, model_name), model_name + ".onnx"), diff --git a/onnxruntime/python/tools/transformers/onnx_exporter.py b/onnxruntime/python/tools/transformers/onnx_exporter.py index 97f1d3adbf..d12c1d1307 100644 --- a/onnxruntime/python/tools/transformers/onnx_exporter.py +++ b/onnxruntime/python/tools/transformers/onnx_exporter.py @@ -9,7 +9,7 @@ import numpy import os import torch from pathlib import Path -from transformers import AutoConfig, AutoTokenizer, AutoModel, LxmertConfig, TransfoXLConfig +from transformers import AutoConfig, AutoTokenizer, LxmertConfig, TransfoXLConfig from affinity_helper import AffinitySetting from benchmark_helper import create_onnxruntime_session, Precision from gpt2_helper import GPT2ModelNoPastState, PRETRAINED_GPT2_MODELS, TFGPT2ModelNoPastState diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index 75647f52a8..cfd924046e 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -698,13 +698,9 @@ class OnnxModel: Args: outputs (list): a list of graph outputs to retain. If it is None, all graph outputs will be kept. """ - - for node in self.model.graph.node: - # Some operators with inner graph in attributes like 'body' 'else_branch' or 'then_branch' - if node.op_type in ['Loop', 'Scan', 'If']: - # TODO: handle inner graph - logger.debug(f"Skip prune_graph since graph has operator: {node.op_type}") - return + if len(self.graphs()) > 1: + logger.debug(f"Skip prune_graph since graph has subgraph") + return if outputs is None: outputs = [output.name for output in self.model.graph.output] diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert.py b/onnxruntime/python/tools/transformers/onnx_model_bert.py index a937201497..6718c3acfc 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert.py @@ -8,10 +8,11 @@ from typing import List from onnx import GraphProto, ModelProto, TensorProto, ValueInfoProto, helper from onnx_model import OnnxModel from fusion_reshape import FusionReshape +from fusion_shape import FusionShape from fusion_layernorm import FusionLayerNormalization, FusionLayerNormalizationTF from fusion_skiplayernorm import FusionSkipLayerNormalization, FusionBiasSkipLayerNormalization from fusion_embedlayer import FusionEmbedLayerNormalization -from fusion_attention import FusionAttention, AttentionMask, AttentionMaskFormat +from fusion_attention import FusionAttention, AttentionMask from fusion_gelu import FusionGelu from fusion_fastgelu import FusionFastGelu from fusion_biasgelu import FusionBiasGelu @@ -74,6 +75,10 @@ class BertOnnxModel(OnnxModel): fusion = FusionReshape(self) fusion.apply() + def fuse_shape(self): + fusion = FusionShape(self) + fusion.apply() + def fuse_embed_layer(self): fusion = FusionEmbedLayerNormalization(self) fusion.apply() @@ -313,6 +318,8 @@ class BertOnnxModel(OnnxModel): self.attention_mask.set_mask_format(options.attention_mask_format) self.fuse_attention() + self.fuse_shape() + if (options is None) or options.enable_embed_layer_norm: self.fuse_embed_layer() @@ -340,7 +347,7 @@ class BertOnnxModel(OnnxModel): if add_dynamic_axes: self.use_dynamic_axes() - logger.info(f"opset verion: {self.model.opset_import[0].version}") + logger.info(f"opset verion: {self.get_opset_version()}") def get_fused_operator_statistics(self): """ diff --git a/onnxruntime/python/tools/transformers/onnx_model_gpt2.py b/onnxruntime/python/tools/transformers/onnx_model_gpt2.py index 7069100469..40ef47869a 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_gpt2.py +++ b/onnxruntime/python/tools/transformers/onnx_model_gpt2.py @@ -4,11 +4,6 @@ #-------------------------------------------------------------------------- import logging import onnx -import sys -import argparse -import numpy as np -from collections import deque -from onnx import ModelProto, TensorProto, numpy_helper from onnx_model_bert import BertOnnxModel from fusion_gpt_attention_no_past import FusionGptAttentionNoPast from fusion_gpt_attention import FusionGptAttention diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py index ddaf027de0..9dd691b2bd 100644 --- a/onnxruntime/python/tools/transformers/optimizer.py +++ b/onnxruntime/python/tools/transformers/optimizer.py @@ -19,16 +19,12 @@ import logging import coloredlogs -import onnx import os -import sys import argparse -import numpy as np from typing import Dict -from collections import deque -from onnx import ModelProto, TensorProto, numpy_helper, load_model +from onnx import load_model from onnx_model_bart import BartOnnxModel -from onnx_model_bert import BertOnnxModel, BertOptimizationOptions +from onnx_model_bert import BertOnnxModel from onnx_model_bert_tf import BertOnnxModelTF from onnx_model_bert_keras import BertOnnxModelKeras from onnx_model_gpt2 import Gpt2OnnxModel diff --git a/onnxruntime/test/python/transformers/test_optimizer.py b/onnxruntime/test/python/transformers/test_optimizer.py index ac76331f2a..b3e7f83fdd 100644 --- a/onnxruntime/test/python/transformers/test_optimizer.py +++ b/onnxruntime/test/python/transformers/test_optimizer.py @@ -37,13 +37,17 @@ BERT_TEST_MODELS = { } +def _get_fusion_test_model(file): + relative_path = os.path.join(os.path.dirname(__file__), '..', '..', 'testdata', 'transform', 'fusion', file) + if (os.path.exists(relative_path)): + return relative_path + return os.path.join('.', 'testdata', 'transform', 'fusion', file) + + def _get_test_model_path(name): sub_dir, file = BERT_TEST_MODELS[name] if sub_dir == "FUSION": - relative_path = os.path.join(os.path.dirname(__file__), '..', '..', 'testdata', 'transform', 'fusion', file) - if (os.path.exists(relative_path)): - return relative_path - return os.path.join('.', 'testdata', 'transform', 'fusion', file) + return _get_fusion_test_model(file) else: relative_path = os.path.join(os.path.dirname(__file__), 'test_data', sub_dir, file) if (os.path.exists(relative_path)): @@ -58,7 +62,8 @@ class TestBertOptimization(unittest.TestCase): print(f"Counters is not expected in test: {test_name}") for op, counter in expected_node_count.items(): print("{}: {} expected={}".format(op, len(bert_model.get_nodes_by_op_type(op)), counter)) - self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count) + + self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count) # add test function for huggingface pytorch model def _test_optimizer_on_huggingface_model(self, @@ -198,6 +203,20 @@ class TestBertOptimization(unittest.TestCase): } self.verify_node_count(model, expected_node_count, 'test_multiple_embed') + def test_embed_layer_norm_fusion(self): + onnx_files = [] + for i in [3, 8, 9]: + onnx_files.append(f"embed_layer_norm_format{i}.onnx") + onnx_files.append(f"embed_layer_norm_format{i}_opset13.onnx") + onnx_files.append('embed_layer_norm_format3_no_cast.onnx') + onnx_files.append('embed_layer_norm_format3_no_cast_opset13.onnx') + + for file in onnx_files: + input_model_path = _get_fusion_test_model(file) + model = optimize_model(input_model_path, 'bert') + expected_node_count = {'EmbedLayerNormalization': 1, 'Attention': 1, 'ReduceSum': 0} + self.verify_node_count(model, expected_node_count, file) + # def test_bert_tf2onnx_0(self): # input = _get_test_model_path('bert_tf2onnx_0') # model = optimize_model(input, 'bert_tf', num_heads=2, hidden_size=8)