onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/fusion_gpt_attention_megatron.py

@@ -0,0 +1,355 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from logging import getLogger
+
+import numpy as np
+from fusion_gpt_attention import FusionGptAttentionPastBase
+from onnx import helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+def is_close(value, expected_value):
+    return abs(value - expected_value) <= 1e-6
+
+
+class FusionGptAttentionMegatron(FusionGptAttentionPastBase):
+    """
+    Fuse GPT-2 Attention with past state subgraph from Megatron into one Attention node.
+    """
+
+    def __init__(self, model: OnnxModel, num_heads: int):
+        super().__init__(model, num_heads)
+
+    def fuse_attention_node(
+        self,
+        matmul_before_split,
+        add_before_split,
+        past,
+        present,
+        input,
+        reshape_qkv,
+        mask,
+    ):
+        attention_node_name = self.model.create_node_name("GptAttention")
+        int32_mask = self.cast_attention_mask(mask)
+        output = reshape_qkv.output[0]
+        i = 1 if (add_before_split.input[0] == matmul_before_split.output[0]) else 0
+        attention_node = helper.make_node(
+            "Attention",
+            inputs=[
+                input,
+                matmul_before_split.input[1],
+                add_before_split.input[i],
+                int32_mask,
+                past,
+            ],
+            outputs=[output, present],
+            name=attention_node_name,
+        )
+        attention_node.domain = "com.microsoft"
+        attention_node.attribute.extend(
+            [
+                helper.make_attribute("num_heads", self.num_heads),
+                helper.make_attribute("unidirectional", 0),  # unidirectional shall not be ON for 4D attention mask
+            ]
+        )
+        if self.mask_filter_value is not None:
+            attention_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))])
+
+        nodes_to_add = [attention_node]
+        self.nodes_to_add.extend(nodes_to_add)
+
+        for node in nodes_to_add:
+            self.node_name_to_graph_name[node.name] = self.this_graph_name
+
+        self.nodes_to_remove.append(reshape_qkv)
+
+        # we rely on prune_graph() to clean old subgraph nodes
+        self.prune_graph = True
+
+    def match_mask(self, sub_qk, mul_qk, matmul_qk, layernorm_before_attention):
+        mask_nodes = self.model.match_parent_path(sub_qk, ["Mul", "Sub", "Slice", "Slice"], [1, 0, 1, 0])
+        if mask_nodes is None:
+            logger.debug("fuse_attention: failed to match unidirectional mask path")
+            return None
+        (mul_mask, sub_mask, last_slice_mask, slice_mask) = mask_nodes
+
+        if len(mask_nodes) > 1 and mask_nodes[0].op_type == "Mul":
+            _, mul_val = self.model.get_constant_input(mask_nodes[0])
+            if mul_val != 10000:
+                self.mask_filter_value = -mul_val
+
+        if mul_qk.input[1] != last_slice_mask.output[0]:
+            logger.debug("fuse_attention failed: mul_qk.input[1] != last_slice_mask.output[0]")
+            return None
+
+        if not self.utils.check_node_input_value(mul_mask, 1, 10000.0):
+            logger.debug("fuse_attention failed: mul_mask input 1 is not constant 10000.0")
+            return None
+
+        if not self.utils.check_node_input_value(sub_mask, 0, 1.0):
+            logger.debug("fuse_attention failed: sub_mask input 0 is not constant 1.0")
+            return None
+
+        if not self.model.find_graph_input(slice_mask.input[0]):
+            logger.info("expect slick_mask input 0 to be graph input")
+            return None
+
+        if not self.utils.check_node_input_value(last_slice_mask, 1, [0]):
+            logger.debug("fuse_attention failed: last_slice_mask input 1 (starts) is not constant [0]")
+            return None
+
+        if not self.utils.check_node_input_value(last_slice_mask, 3, [3]):
+            logger.debug("fuse_attention failed: last_slice_mask input 3 (axes) is not constant [3]")
+            return False
+
+        if not self.utils.check_node_input_value(last_slice_mask, 4, [1]):
+            logger.debug("fuse_attention failed: last_slice_mask input 4 (steps) is not constant [1]")
+            return False
+
+        if not self.utils.check_node_input_value(slice_mask, 3, [2]):
+            logger.debug("fuse_attention failed: slice_mask input 3 (axes) is not constant [2]")
+            return None
+
+        if not self.utils.check_node_input_value(slice_mask, 4, [1]):
+            logger.debug("fuse_attention failed: slice_mask input 4 (steps) is not constant [1]")
+            return None
+
+        last_slice_path = self.model.match_parent_path(
+            last_slice_mask, ["Unsqueeze", "Gather", "Shape", "MatMul"], [2, 0, 0, 0]
+        )
+        if last_slice_path is None or last_slice_path[-1] != matmul_qk:
+            logger.debug("fuse_attention: failed to match last slice path")
+            return None
+
+        first_slice_path = self.model.match_parent_path(
+            slice_mask, ["Unsqueeze", "Gather", "Shape", "MatMul"], [2, 0, 0, 0]
+        )
+        if first_slice_path is None or first_slice_path[-1] != matmul_qk:
+            logger.debug("fuse_attention: failed to match first slice path")
+            return None
+
+        first_slice_sub = self.model.match_parent_path(
+            slice_mask,
+            ["Unsqueeze", "Sub", "Gather", "Shape", "MatMul"],
+            [1, 0, 0, 0, 0],
+        )
+        if first_slice_sub is None or first_slice_sub[-1] != matmul_qk:
+            logger.debug("fuse_attention: failed to match last slice sub path")
+            return None
+
+        first_slice_sub_1 = self.model.match_parent_path(
+            slice_mask,
+            ["Unsqueeze", "Sub", "Gather", "Shape", "LayerNormalization"],
+            [1, 0, 1, 0, 0],
+        )
+
+        if first_slice_sub_1 is None:
+            first_slice_sub_1 = self.model.match_parent_path(
+                slice_mask,
+                ["Unsqueeze", "Sub", "Gather", "Shape", "SkipLayerNormalization"],
+                [1, 0, 1, 0, 0],
+            )
+
+        if first_slice_sub_1 is None or first_slice_sub_1[-1] != layernorm_before_attention:
+            logger.debug("fuse_attention: failed to match last slice sub path 1")
+            return None
+
+        return slice_mask.input[0]
+
+    def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
+        past = None
+        present = None
+
+        is_normalize_node_skiplayernorm = normalize_node.op_type == "SkipLayerNormalization"
+        qkv_nodes = None
+
+        if not is_normalize_node_skiplayernorm:
+            qkv_nodes = self.model.match_parent_path(
+                normalize_node,
+                ["Add", "Add", "MatMul", "Reshape", "Transpose", "MatMul"],
+                [0, 1, None, 0, 0, 0],
+                output_name_to_node=output_name_to_node,
+            )
+        else:
+            qkv_nodes = self.model.match_parent_path(
+                normalize_node,
+                ["Add", "MatMul", "Reshape", "Transpose", "MatMul"],
+                [1, None, 0, 0, 0],
+                output_name_to_node=output_name_to_node,
+            )
+
+        if qkv_nodes is None:
+            return
+
+        skip_input = None
+        if not is_normalize_node_skiplayernorm:
+            (
+                add_skip,
+                add_after_attention,
+                matmul_after_attention,
+                reshape_qkv,
+                transpose_qkv,
+                matmul_qkv,
+            ) = qkv_nodes
+
+            skip_input = add_skip.input[0]
+        else:
+            (
+                add_after_attention,
+                matmul_after_attention,
+                reshape_qkv,
+                transpose_qkv,
+                matmul_qkv,
+            ) = qkv_nodes
+
+            skip_input = normalize_node.input[0]
+
+        v_nodes = self.model.match_parent_path(
+            matmul_qkv,
+            [
+                "Concat",
+                "Transpose",
+                "Reshape",
+                "Split",
+                "Add",
+                "MatMul",
+                "LayerNormalization",
+            ],
+            [1, 1, 0, 0, 0, None, 0],
+        )
+
+        if v_nodes is None:
+            v_nodes = self.model.match_parent_path(
+                matmul_qkv,
+                [
+                    "Concat",
+                    "Transpose",
+                    "Reshape",
+                    "Split",
+                    "Add",
+                    "MatMul",
+                    "SkipLayerNormalization",
+                ],
+                [1, 1, 0, 0, 0, None, 0],
+            )
+
+        if v_nodes is None:
+            logger.debug("fuse_attention: failed to match v path")
+            return
+        (
+            concat_v,
+            transpose_v,
+            reshape_v,
+            split_v,
+            add_before_split,
+            matmul_before_split,
+            layernorm_before_attention,
+        ) = v_nodes
+
+        if (
+            layernorm_before_attention.op_type == "LayerNormalization"
+            and skip_input != layernorm_before_attention.input[0]
+        ):
+            logger.debug("fuse_attention: skip_input != layernorm_before_attention.input[0]")
+            return
+
+        if (
+            layernorm_before_attention.op_type == "SkipLayerNormalization"
+            and skip_input != layernorm_before_attention.output[3]
+        ):
+            logger.debug("fuse_attention: skip_input != layernorm_before_attention.input[0]")
+            return
+
+        qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Sub", "Mul", "MatMul"], [0, 0, 0, 0])
+        if qk_nodes is None:
+            logger.debug("fuse_attention: failed to match qk path")
+            return None
+        (softmax_qk, sub_qk, mul_qk, matmul_qk) = qk_nodes
+        if self.model.get_node_attribute(softmax_qk, "axis") != 3:
+            logger.debug("fuse_attention failed: softmax_qk axis != 3")
+            return None
+
+        attention_mask = self.match_mask(sub_qk, mul_qk, matmul_qk, layernorm_before_attention)
+
+        q_nodes = self.model.match_parent_path(matmul_qk, ["Div", "Transpose", "Reshape", "Split"], [0, 0, 0, 0])
+        if q_nodes is None:
+            logger.debug("fuse_attention: failed to match q path")
+            return
+        (div_q, transpose_q, reshape_q, split_q) = q_nodes
+        if split_v != split_q:
+            logger.debug("fuse_attention: skip since split_v != split_q")
+            return
+
+        k_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Div", "Transpose", "Concat", "Transpose", "Reshape", "Split"],
+            [1, 0, 0, 1, 0, 0],
+        )
+        if k_nodes is None:
+            logger.debug("fuse_attention: failed to match k path")
+            return
+        (div_k, _, concat_k, transpose_k, reshape_k, split_k) = k_nodes
+        if split_v != split_k:
+            logger.debug("fuse_attention: skip since split_v != split_k")
+            return
+
+        i, value = self.model.get_constant_input(reshape_k)
+        if not (
+            isinstance(value, np.ndarray)
+            and list(value.shape) == [4]
+            and value[0] == 0
+            and value[1] == 0
+            and value[2] > 0
+            and value[3] > 0
+        ):
+            logger.debug("fuse_attention: reshape constant input is not [0, 0, N, H]")
+            return
+
+        num_heads = value[2]
+        if num_heads != self.num_heads:
+            logger.info(f"Detected num_heads={num_heads}. Ignore user specified value {self.num_heads}")
+            self.num_heads = num_heads
+
+        hidden_size_per_head = value[3]
+        i, value = self.model.get_constant_input(div_k)
+        expected_value = float(np.sqrt(np.sqrt(hidden_size_per_head)))
+        if not is_close(value, expected_value):
+            logger.debug(f"fuse_attention: div_k value={value} expected={expected_value}")
+            return
+
+        i, value = self.model.get_constant_input(div_q)
+        if not is_close(value, expected_value):
+            logger.debug(f"fuse_attention: div_q value={value} expected={expected_value}")
+            return
+
+        # Match past and present paths
+        past = self.match_past_pattern_2(concat_k, concat_v, output_name_to_node)
+        if past is None:
+            logger.debug("fuse_attention: match past failed")
+            return
+        if not self.model.find_graph_input(past):
+            logger.debug("fuse_attention: past is not graph input.")
+            # For GPT2LMHeadModel_BeamSearchStep, there is an extra Gather node to select beam index so it is not graph input.
+
+        present = self.match_present(concat_v, input_name_to_nodes)
+        if present is None:
+            logger.debug("fuse_attention: match present failed")
+            return
+        if not self.model.find_graph_output(present):
+            logger.info("fuse_attention: expect present to be graph output")
+            return
+
+        self.fuse_attention_node(
+            matmul_before_split,
+            add_before_split,
+            past,
+            present,
+            layernorm_before_attention.output[0],
+            reshape_qkv,
+            attention_mask,
+        )

onnxruntime/transformers/fusion_gpt_attention_no_past.py

@@ -0,0 +1,260 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from logging import getLogger
+
+from fusion_base import Fusion
+from onnx import helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionGptAttentionNoPast(Fusion):
+    """
+    Fuse GPT-2 Attention without past state into one Attention node.
+    This does not support attention_mask graph input right now.
+    """
+
+    def __init__(self, model: OnnxModel, num_heads: int):
+        super().__init__(model, "Attention", ["LayerNormalization", "SkipLayerNormalization"], "without past")
+        # TODO: detect num_heads from graph like FusionAttention
+        self.num_heads = num_heads
+        self.mask_filter_value = None
+
+    def create_attention_node(self, gemm, gemm_qkv, input, output):
+        attention_node_name = self.model.create_node_name("Attention")
+        attention_node = helper.make_node(
+            "Attention",
+            inputs=[input, gemm.input[1], gemm.input[2]],
+            outputs=[attention_node_name + "_output"],
+            name=attention_node_name,
+        )
+        attention_node.domain = "com.microsoft"
+        attention_node.attribute.extend(
+            [
+                helper.make_attribute("num_heads", self.num_heads),
+                helper.make_attribute("unidirectional", 1),
+            ]
+        )
+        if self.mask_filter_value is not None:
+            attention_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))])
+
+        matmul_node = helper.make_node(
+            "MatMul",
+            inputs=[attention_node_name + "_output", gemm_qkv.input[1]],
+            outputs=[attention_node_name + "_matmul_output"],
+            name=attention_node_name + "_matmul",
+        )
+
+        add_node = helper.make_node(
+            "Add",
+            inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]],
+            outputs=[output],
+            name=attention_node_name + "_add",
+        )
+
+        self.nodes_to_add.extend([attention_node, matmul_node, add_node])
+        self.node_name_to_graph_name[attention_node.name] = self.this_graph_name
+        self.node_name_to_graph_name[matmul_node.name] = self.this_graph_name
+        self.node_name_to_graph_name[add_node.name] = self.this_graph_name
+
+    def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
+        # (TODO) hasesh/tlwu: Investigate what fixes the following logic needs in order
+        # to fuse the Attention sub-graph. With some changes to other fusions, this stopped
+        # working.
+        return_indice = []
+
+        is_normalize_node_skiplayernorm = normalize_node.op_type == "SkipLayerNormalization"
+        qkv_nodes = None
+
+        if not is_normalize_node_skiplayernorm:
+            qkv_nodes = self.model.match_parent_path(
+                normalize_node,
+                ["Add", "Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
+                [0, None, 0, 0, 0, 0, 0],
+                output_name_to_node=output_name_to_node,
+                return_indice=return_indice,
+            )
+        else:
+            qkv_nodes = self.model.match_parent_path(
+                normalize_node,
+                ["Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
+                [None, 0, 0, 0, 0, 0],
+                output_name_to_node=output_name_to_node,
+                return_indice=return_indice,
+            )
+
+        if qkv_nodes is None:
+            return
+
+        another_input = None
+        if not is_normalize_node_skiplayernorm:
+            (
+                add_qkv,
+                reshape_qkv,
+                gemm_qkv,
+                reshape_1,
+                reshape_2,
+                transpose_qkv,
+                matmul_qkv,
+            ) = qkv_nodes
+
+            another_input = add_qkv.input[1 - return_indice[0]]
+        else:
+            (
+                reshape_qkv,
+                gemm_qkv,
+                reshape_1,
+                reshape_2,
+                transpose_qkv,
+                matmul_qkv,
+            ) = qkv_nodes
+
+        v_nodes = self.model.match_parent_path(
+            matmul_qkv,
+            ["Transpose", "Reshape", "Split", "Reshape", "Gemm", "Reshape"],
+            [1, 0, 0, 0, 0, 0],
+        )
+        if v_nodes is None:
+            logger.debug("fuse_attention: failed to match v path")
+            return
+        (
+            transpose_v,
+            reshape_v,
+            split_v,
+            reshape_after_gemm,
+            gemm,
+            reshape_before_gemm,
+        ) = v_nodes
+
+        layernorm_before_attention = self.model.get_parent(reshape_before_gemm, 0, output_name_to_node)
+        if layernorm_before_attention is None or (
+            layernorm_before_attention.op_type != "LayerNormalization"
+            and layernorm_before_attention.op_type != "SkipLayerNormalization"
+        ):
+            if layernorm_before_attention.op_type != "Add":
+                logger.debug(f"failed to get (skip)layernorm before gemm. Got {layernorm_before_attention.op_type}")
+                return
+
+        # `another_input` will be non-None only if
+        # (1) SkipLayerNorm fusion wasn't turned ON
+        # (2) SkipLayerNorm fusion was turned ON but upstream layer's LayerNorm + Add was not
+        # fused into a SkipLayerNorm. This can happen if the shapes to the Add node are different.
+        # So, keep the following check if SkipLayerNorm fusion is turned ON or OFF.
+        if another_input is not None:
+            if another_input not in layernorm_before_attention.input:
+                # match openai-gpt
+                if another_input not in layernorm_before_attention.output:
+                    logger.debug("Add and (Skip)LayerNormalization shall have one same input")
+                    return
+
+        qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Sub", "Mul", "Div", "MatMul"], [0, 0, 0, 0, 0])
+        if qk_nodes is not None:
+            (softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes
+            mask_nodes = self.model.match_parent_path(
+                sub_qk,
+                [
+                    "Mul",
+                    "Sub",
+                    "Slice",
+                    "Slice",
+                    "Unsqueeze",
+                    "Sub",
+                    "Squeeze",
+                    "Slice",
+                    "Shape",
+                    "Div",
+                ],
+                [1, 0, 1, 0, 1, 0, 0, 0, 0, 0],
+            )
+            if mask_nodes is None:
+                logger.debug("fuse_attention: failed to match mask path")
+                return
+            div_mask = mask_nodes[-1]
+
+            if div_qk != div_mask:
+                logger.debug("fuse_attention: skip since div_qk != div_mask")
+                return
+            if len(mask_nodes) > 1 and mask_nodes[0].op_type == "Mul":
+                _, mul_val = self.model.get_constant_input(mask_nodes[0])
+                if mul_val != -10000:
+                    self.mask_filter_value = mul_val
+
+        else:
+            # New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0.
+            qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Where", "Div", "MatMul"], [0, 0, 1, 0])
+            if qk_nodes is not None:
+                (softmax_qk, where_qk, div_qk, matmul_qk) = qk_nodes
+                mask_nodes = self.model.match_parent_path(
+                    where_qk,
+                    [
+                        "Cast",
+                        "Slice",
+                        "Slice",
+                        "Unsqueeze",
+                        "Sub",
+                        "Squeeze",
+                        "Slice",
+                        "Shape",
+                        "Div",
+                    ],
+                    [0, 0, 0, 1, 0, 0, 0, 0, 0],
+                )
+                if mask_nodes is None:
+                    logger.debug("fuse_attention: failed to match mask path")
+                    return
+                div_mask = mask_nodes[-1]
+
+                if div_qk != div_mask:
+                    logger.debug("fuse_attention: skip since div_qk != div_mask")
+                    return
+            else:
+                # match openai-gpt
+                qk_nodes = self.model.match_parent_path(
+                    matmul_qkv,
+                    ["Softmax", "Add", "Mul", "Div", "MatMul"],
+                    [0, 0, 0, 0, 0],
+                )
+                if qk_nodes is None:
+                    logger.debug("fuse_attention: failed to match qk path")
+                    return
+                (softmax_qk, add_qk, mul_qk, div_qk, matmul_qk) = qk_nodes
+                mask_nodes = self.model.match_parent_path(
+                    mul_qk,
+                    ["Slice", "Slice", "Unsqueeze", "Squeeze", "Slice", "Shape", "Div"],
+                    [1, 0, 2, 0, 0, 0, 0],
+                )
+                if mask_nodes is None:
+                    logger.debug("fuse_attention: failed to match mask path")
+                    return
+                div_mask = mask_nodes[-1]
+
+                if div_qk != div_mask:
+                    logger.debug("fuse_attention: skip since div_qk != div_mask")
+                    return
+
+        q_nodes = self.model.match_parent_path(matmul_qk, ["Transpose", "Reshape", "Split"], [0, 0, 0])
+        if q_nodes is None:
+            logger.debug("fuse_attention: failed to match q path")
+            return
+        (transpose_q, reshape_q, split_q) = q_nodes
+        if split_v != split_q:
+            logger.debug("fuse_attention: skip since split_v != split_q")
+            return
+
+        k_nodes = self.model.match_parent_path(matmul_qk, ["Transpose", "Reshape", "Split"], [1, 0, 0])
+        if k_nodes is None:
+            logger.debug("fuse_attention: failed to match k path")
+            return
+        (transpose_k, reshape_k, split_k) = k_nodes
+        if split_v != split_k:
+            logger.debug("fuse_attention: skip since split_v != split_k")
+            return
+
+        self.create_attention_node(gemm, gemm_qkv, layernorm_before_attention.output[0], reshape_qkv.output[0])
+
+        # we rely on prune_graph() to clean old subgraph nodes:
+        # qk_nodes + q_nodes + k_nodes + v_nodes + mask_nodes + [reshape_qkv, transpose_qkv, matmul_qkv]
+        self.prune_graph = True