onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/onnx_model_gpt2.py

@@ -0,0 +1,101 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import logging
+
+import onnx
+from fusion_gpt_attention import FusionGptAttention
+from fusion_gpt_attention_megatron import FusionGptAttentionMegatron
+from fusion_gpt_attention_no_past import FusionGptAttentionNoPast
+from fusion_rotary_attention import FusionRotaryAttention
+from onnx_model_bert import BertOnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class Gpt2OnnxModel(BertOnnxModel):
+    def __init__(self, model, num_heads, hidden_size):
+        super().__init__(model, num_heads, hidden_size)
+
+    def fuse_attention(self):
+        if len(self.model.graph.input) == 1 or len(self.model.graph.output) == 1:
+            fusion = FusionGptAttentionNoPast(self, self.num_heads)
+            fusion.apply()
+        else:
+            fusion = FusionGptAttention(self, self.num_heads)
+            fusion.apply()
+            fusion = FusionGptAttentionMegatron(self, self.num_heads)
+            fusion.apply()
+
+        fusion = FusionRotaryAttention(self, self.hidden_size, self.num_heads)
+        fusion.apply()
+
+    def postprocess(self):
+        """
+        Remove extra reshape nodes.
+        """
+        logger.debug("start postprocessing...")
+
+        input_name_to_nodes = self.input_name_to_nodes()
+        output_name_to_node = self.output_name_to_node()
+
+        reshape_count = 0
+        for gemm_node in self.get_nodes_by_op_type("Gemm"):
+            reshape_after_gemm = self.find_first_child_by_type(
+                gemm_node, "Reshape", input_name_to_nodes, recursive=False
+            )
+
+            nodes = self.match_parent_path(gemm_node, ["Reshape", "FastGelu"], [0, 0], output_name_to_node)
+            if nodes is None:
+                nodes = self.match_parent_path(
+                    gemm_node,
+                    ["Reshape", "LayerNormalization"],
+                    [0, 0],
+                    output_name_to_node,
+                )
+
+                if nodes is None:
+                    nodes = self.match_parent_path(
+                        gemm_node,
+                        ["Reshape", "SkipLayerNormalization"],
+                        [0, 0],
+                        output_name_to_node,
+                    )
+
+                    if nodes is None:
+                        continue
+
+            (reshape_before_gemm, root_node) = nodes
+
+            matmul_node_name = self.create_node_name("MatMul", "FullyConnect_MatMul")
+            matmul_node = onnx.helper.make_node(
+                "MatMul",
+                inputs=[matmul_node_name + "_input", gemm_node.input[1]],
+                outputs=[matmul_node_name + "_output"],
+                name=matmul_node_name,
+            )
+
+            add_node_name = self.create_node_name("Add", "FullyConnect_Add")
+            add_node = onnx.helper.make_node(
+                "Add",
+                inputs=[matmul_node_name + "_output", gemm_node.input[2]],
+                outputs=[add_node_name + "_output"],
+                name=add_node_name,
+            )
+
+            self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output")
+
+            # Link root node output with MatMul
+            self.replace_input_of_all_nodes(root_node.output[0], matmul_node_name + "_input")
+            root_node.output[0] = matmul_node_name + "_input"
+
+            self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output")
+
+            self.add_node(matmul_node)
+            self.add_node(add_node)
+
+            reshape_count += 2
+
+        self.prune_graph()
+        logger.info(f"postprocess: remove Reshape count: {reshape_count}")

onnxruntime/transformers/onnx_model_mmdit.py

@@ -0,0 +1,112 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import logging
+
+from fusion_layernorm import FusionLayerNormalization
+from fusion_mha_mmdit import FusionMultiHeadAttentionMMDit
+from fusion_options import FusionOptions
+from import_utils import is_installed
+from onnx import ModelProto
+from onnx_model_bert import BertOnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class MmditOnnxModel(BertOnnxModel):
+    def __init__(self, model: ModelProto, num_heads: int = 0, hidden_size: int = 0):
+        """Initialize Multimodal Diffusion Transformer (MMDiT) ONNX Model.
+
+        Args:
+            model (ModelProto): the ONNX model
+            num_heads (int, optional): number of attention heads. Defaults to 0 (detect the parameter automatically).
+            hidden_size (int, optional): hidden dimension. Defaults to 0 (detect the parameter automatically).
+        """
+        assert (num_heads == 0 and hidden_size == 0) or (num_heads > 0 and hidden_size % num_heads == 0)
+        super().__init__(model, num_heads=num_heads, hidden_size=hidden_size)
+
+    def postprocess(self):
+        self.prune_graph()
+        self.remove_unused_constant()
+
+    def fuse_layer_norm(self):
+        layernorm_support_broadcast = True
+        logger.warning(
+            "The optimized model requires LayerNormalization with broadcast support. "
+            "Please use onnxruntime-gpu>=1.21 for inference."
+        )
+        fusion = FusionLayerNormalization(
+            self, check_constant_and_dimension=not layernorm_support_broadcast, force=True
+        )
+        fusion.apply()
+
+    def fuse_multi_head_attention(self):
+        fusion = FusionMultiHeadAttentionMMDit(self)
+        fusion.apply()
+
+    def optimize(self, options: FusionOptions | None = None, add_dynamic_axes: bool = False):
+        assert not add_dynamic_axes
+
+        if is_installed("tqdm"):
+            import tqdm  # noqa: PLC0415
+            from tqdm.contrib.logging import logging_redirect_tqdm  # noqa: PLC0415
+
+            with logging_redirect_tqdm():
+                steps = 5
+                progress_bar = tqdm.tqdm(range(steps), initial=0, desc="fusion")
+                self._optimize(options, progress_bar)
+        else:
+            logger.info("tqdm is not installed. Run optimization without progress bar")
+            self._optimize(options, None)
+
+    def _optimize(self, options: FusionOptions | None = None, progress_bar=None):
+        if (options is not None) and not options.enable_shape_inference:
+            self.disable_shape_inference()
+
+        # Remove cast nodes that having same data type of input and output based on symbolic shape inference.
+        self.utils.remove_useless_cast_nodes()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_layer_norm:
+            self.fuse_layer_norm()
+            self.fuse_simplified_layer_norm()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_gelu:
+            self.fuse_gelu()
+        if progress_bar:
+            progress_bar.update(1)
+
+        if (options is None) or options.enable_attention:
+            self.fuse_multi_head_attention()
+        if progress_bar:
+            progress_bar.update(1)
+
+        self.postprocess()
+        if progress_bar:
+            progress_bar.update(1)
+
+        logger.info(f"opset version: {self.get_opset_version()}")
+
+    def get_fused_operator_statistics(self):
+        """
+        Returns node count of fused operators.
+        """
+        op_count = {}
+        ops = [
+            "FastGelu",
+            "MultiHeadAttention",
+            "LayerNormalization",
+            "SimplifiedLayerNormalization",
+        ]
+
+        for op in ops:
+            nodes = self.get_nodes_by_op_type(op)
+            op_count[op] = len(nodes)
+
+        logger.info(f"Optimized operators:{op_count}")
+        return op_count