onnxruntime-directml 1.20.0__cp313-cp313-win_amd64.whl

onnxruntime/transformers/fusion_embedlayer.py

@@ -0,0 +1,811 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from logging import getLogger
+from typing import Dict, List, Optional, Tuple, Union
+
+from fusion_base import Fusion
+from fusion_utils import FusionUtils
+from onnx import NodeProto, TensorProto, helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionEmbedLayerNoMask(Fusion):
+    """
+    Fuse embedding layer into one node (EmbedLayerNormalization).
+    It supports the following model types: BERT, DistilBert, ALBert.
+    """
+
+    def __init__(self, model: OnnxModel, description: str = "no mask"):
+        super().__init__(
+            model,
+            "EmbedLayerNormalization",
+            ["LayerNormalization", "SkipLayerNormalization"],
+            description,
+        )
+        self.utils = FusionUtils(model)
+        self.shape_infer = None
+        self.shape_infer_done = False
+
+        # The following will be reset in each fuse call of FusionEmbedLayerNormalization
+        self.attention = None
+        self.embed_node = 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
+
+        gather_1_path = self.model.match_parent_path(add, ["Gather"], [1])
+        if gather_1_path is None:
+            return None
+
+        return gather_0_path[0], gather_1_path[0]
+
+    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.
+
+        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
+
+        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 not None:
+            return True
+
+        if layernorm.output[0] not in input_name_to_nodes:
+            return False
+        children = input_name_to_nodes[layernorm.output[0]]
+        children_types = sorted([child.op_type for child in children])
+
+        # Try find MultiHeadAttention
+        if children_types == ["MatMul", "MatMul", "MatMul", "SkipLayerNormalization"]:
+            for node in children:
+                if node.op_type == "SkipLayerNormalization":
+                    path1 = self.model.match_parent_path(
+                        node,
+                        ["Add", "MatMul", "MultiHeadAttention", "MatMul"],
+                        [None, None, 0, 0],
+                    )
+                    if path1 is not None and path1[-1].input[0] == layernorm.output[0]:
+                        self.cross_attention = path1[2]
+                        return True
+
+        # In case user disables attention fusion, check whether subgraph looks like Attention.
+        # 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])
+
+        # 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
+        """
+        # remove after tests pass
+        path1 = self.model.match_parent_path(position_embedding_gather, ["Expand", "Shape"], [1, 1])
+        if path1 is None:
+            path1 = self.model.match_parent_path(
+                position_embedding_gather,
+                ["Expand", "Where", "Reshape", "Shape"],
+                [1, 1, 2, 0],
+            )
+            if path1 is None:
+                return False
+
+        expand, shape = path1[0], path1[-1]
+        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]
+
+        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:
+            input_ids_shape = self.shape_infer.get_edge_shape(input_ids)
+            position_ids_shape = self.shape_infer.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(
+                    f"Cannot fuse EmbedLayerNormalization: input_ids and position_ids not matched in 2nd dimension: {input_ids_shape} vs {position_ids_shape}"
+                )
+                return False
+
+            if segment_ids and not self.shape_infer.compare_shape(input_ids, segment_ids):
+                logger.info(
+                    f"Cannot fuse EmbedLayerNormalization: input_ids and segment_ids does not have same shape: {input_ids_shape} != {self.shape_infer.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 embedding table is the largest, and segment embedding table is the smallest, while position 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.warning(
+                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.warning(
+                    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.warning(
+                    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:
+                int32_output = input_name
+        else:
+            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],
+        position_ids: Optional[str] = None,
+        embedding_sum_output=False,
+        embedding_sum_name=None,
+    ):
+        """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")
+
+        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 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_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,
+            ]
+
+        if position_ids is not None:
+            # Adding an empty input for mask before position_ids
+            embed_node_inputs.append("")
+            position_ids, _ = self.cast_to_int32(position_ids)
+            embed_node_inputs.append(position_ids)
+
+        embed_node_outputs = [node_name + "_output", node_name + "_dummy_mask_index"]
+        if embedding_sum_output:
+            name = embedding_sum_name if embedding_sum_name is not None else node_name + "_embedding_sum"
+            embed_node_outputs.append(name)
+
+        embed_node = helper.make_node(
+            "EmbedLayerNormalization",
+            embed_node_inputs,
+            outputs=embed_node_outputs,
+            name=node_name,
+        )
+
+        embed_node.domain = "com.microsoft"
+
+        # Pass attribute "epsilon" from normalize node to EmbedLayerNormalization.
+        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.
+        # OnnxRuntime 1.2.0 or older has no epsilon attribute. The optimized model can only work for 1.3.0 or later.
+        if len(embed_node.attribute) == 0:
+            embed_node.attribute.extend([helper.make_attribute("epsilon", 1.0e-12)])
+
+        # 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 is_skip_layer_norm_with_sum_output(self, node):
+        return (node.op_type == "SkipLayerNormalization") and len(node.output) > 3 and len(node.output[3]) > 0
+
+    def fuse_gpt2(
+        self, layernorm, add_before_layernorm, input_name_to_nodes, output_name_to_node, optional_segment_gather=None
+    ):
+        # graph checks
+        # gpt2 has optional segment embedding, subgraph pattern is like
+        #                      input_ids  position_ids
+        #                         |        |
+        #  token_ids           Gather    Gather
+        #       |                   \   /
+        #   Gather (optional)        Add _ _ _ _ _
+        #                   \         |           |
+        #                     LayerNormalization  |
+        #                             |           |
+        #                          Attention      |
+        #                             |           |
+        #                           Matmul        |
+        #                             |          /
+        #                            Add        /
+        #                              \       /
+        #                                 Add
+        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]
+        position_ids = position_embedding_gather.input[1]
+
+        if not self.check_attention_subgraph(layernorm, input_name_to_nodes, is_distil_bert=False):
+            return False
+
+        if not self.check_embedding(word_embedding_gather, None, position_embedding_gather):
+            return False
+
+        # If layernorm node is SkipLayerNormalization, we need look at its optional fourth output.
+        # If the add_before_layernorm node is an Add node, then the add_output output is the first output of this node.
+        # If the add_before_layernorm node is a SkipLayerNormalization node, then the add_output output
+        # is the (optional) fourth index output of this node.
+        # When add_before_layernorm is SkipLayerNormalization, add_before_layernorm and layernorm are same node.
+        if layernorm.op_type == "SkipLayerNormalization":
+            need_embedding_sum_output = self.is_skip_layer_norm_with_sum_output(layernorm)
+            sum_output_index = 3
+            node_with_sum_output = layernorm
+            sum_output = layernorm.output[3] if need_embedding_sum_output else None
+            is_sum_graph_output = (sum_output is not None) and (self.model.find_graph_output(sum_output) is not None)
+        else:  # layernorm.op_type == "LayerNormalization"
+            node_with_sum_output = add_before_layernorm
+            sum_output_index = 0 if add_before_layernorm.op_type == "Add" else 3
+            sum_output = (
+                add_before_layernorm.output[sum_output_index]
+                if len(add_before_layernorm.output) > sum_output_index
+                else None
+            )
+            is_sum_graph_output = (sum_output is not None) and (self.model.find_graph_output(sum_output) is not None)
+            is_sum_used_by_multiple_nodes = (
+                sum_output and (sum_output in input_name_to_nodes) and len(input_name_to_nodes[sum_output]) > 1
+            )
+            need_embedding_sum_output = (sum_output is not None) and (
+                add_before_layernorm.op_type != "Add" or is_sum_graph_output or is_sum_used_by_multiple_nodes
+            )
+
+        # make the fused node
+        embed_node = self.create_fused_node(
+            input_ids,
+            layernorm,
+            word_embedding_gather,
+            position_embedding_gather,
+            optional_segment_gather,
+            position_ids,
+            embedding_sum_output=need_embedding_sum_output,
+            embedding_sum_name=sum_output if is_sum_graph_output else None,
+        )
+
+        if need_embedding_sum_output:
+            node_with_sum_output.output[sum_output_index] = "_no_use__to_be_removed_"
+            if not is_sum_graph_output:
+                self.model.replace_input_of_all_nodes(sum_output, embed_node.output[2])
+
+        self.finish_fusion(layernorm, embed_node)
+        return 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):
+        first_add_path = self.model.match_parent_path(node, ["Add"], [0])
+        if node.op_type == "LayerNormalization":
+            if first_add_path is None:
+                return
+            add_before_layernorm = first_add_path[0]
+            optional_segment_gather = None
+        else:  # SkipLayerNormalization
+            gather_0_path = self.model.match_parent_path(node, ["Gather"], [0])
+            gather_1_path = self.model.match_parent_path(node, ["Gather"], [1])
+            if gather_0_path is None and gather_1_path is not None:
+                if first_add_path is None:
+                    return
+                add_before_layernorm = first_add_path[0]
+                optional_segment_gather = gather_1_path[0]
+            elif gather_0_path is not None and gather_1_path is None:
+                first_add_path = self.model.match_parent_path(node, ["Add"], [1])
+                if first_add_path is None:
+                    return
+                add_before_layernorm = first_add_path[0]
+                optional_segment_gather = gather_0_path[0]
+            else:
+                add_before_layernorm = node  # Add is fused into SkipLayerNormalization
+                optional_segment_gather = None
+
+        if self.fuse_gpt2(
+            node, add_before_layernorm, input_name_to_nodes, output_name_to_node, optional_segment_gather
+        ):
+            return
+
+        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):
+    def __init__(self, model: OnnxModel, use_mask_index=False):
+        super().__init__(model, "with mask")
+        self.use_mask_index = use_mask_index
+
+    def replace_mask(self, mask_int32, attention_nodes):
+        # Inputs of EmbedLayerNorm: input_ids, segment_ids (optional), word_embedding, position_embedding,
+        #           segment_embedding (optional), gamma, beta, mask (optional), position_ids (optional)
+        embed_node = self.embed_node
+        if len(embed_node.input) == 7:
+            embed_node.input.append(mask_int32)
+            logger.debug("append mask to %s", embed_node.name)
+        elif len(embed_node.input) > 7 and not embed_node.input[7]:
+            embed_node.input[7] = mask_int32
+            logger.debug("replace mask in %s", embed_node.name)
+        else:
+            logger.debug("skip mask in %s", embed_node.name)
+            return
+
+        for attention_node in attention_nodes:
+            logger.debug("update mask_index in %s", attention_node.name)
+            if attention_node.op_type == "Attention":
+                attention_node.input[3] = embed_node.output[1]
+            elif attention_node.op_type == "MultiHeadAttention":
+                attention_node.input[4] = embed_node.output[1]
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        # Reset attention and embed_node so that we know fusion is successful when they are not None.
+        self.attention = None
+        self.cross_attention = None
+        self.embed_node = None
+        super().fuse(node, input_name_to_nodes, output_name_to_node)
+
+        if self.embed_node is None:
+            return
+
+        if not self.use_mask_index:
+            logger.debug("--use_mask_index is not set: EmbedLayerNormalization will not have mask")
+            self.increase_counter("EmbedLayerNormalization(no mask)")
+            return
+
+        if self.attention is None and self.cross_attention is None:
+            logger.debug("EmbedLayerNormalization will not have mask since attention node is not found")
+            self.increase_counter("EmbedLayerNormalization(no mask)")
+            return
+
+        if self.attention:
+            mask_int32 = self.attention.input[3]
+        else:
+            mask_int32 = self.cross_attention.input[4]
+
+        children_nodes = input_name_to_nodes[mask_int32]
+        if self.model.find_graph_input(mask_int32):
+            attention_nodes = [node for node in children_nodes if node.op_type in ["Attention", "MultiHeadAttention"]]
+            self.replace_mask(mask_int32, attention_nodes)
+            self.increase_counter("EmbedLayerNormalization(with mask)")
+            return
+
+        if mask_int32 not in output_name_to_node:
+            logger.debug("EmbedLayerNormalization will not have mask since %s is not a node output", mask_int32)
+            self.increase_counter("EmbedLayerNormalization(no mask)")
+            return
+
+        node = output_name_to_node[mask_int32]
+        if node.op_type in ["ReduceSum", "Cast"]:
+            attention_nodes = [node for node in children_nodes if node.op_type in ["Attention", "MultiHeadAttention"]]
+            if node.op_type == "ReduceSum":
+                mask_int32 = node.input[0]
+                if len(children_nodes) == len(attention_nodes):
+                    self.nodes_to_remove.append(node)
+            self.replace_mask(mask_int32, attention_nodes)
+            self.increase_counter("EmbedLayerNormalization(with mask)")