Fix transformers optimizations for GPT-NeoX (#16743)

### Description
Fix some issues found in GPT-NeoX graph fusion:
(1) GPT-NeoX uses float16 weights. The step of using onnxruntime with
opt_level==1 uses CPU provider. Since most operators does not have fp16
in CPU EP, so extra Cast nodes are added to up cast to fp32.
(2) Add is shared by two LayerNormalization children, and
SkipLayerNormalization might cause invalid graph.
(3) Reshape fusion might miss since some part only check initializer but
not Constant.

This PR adds a check whether model uses FP16, and output a warning when
use_gpu is not True, and use GPU provider for graph optimization when use_gpu=True.

onnxruntime/python/tools/transformers/fusion_base.py

@@ -4,20 +4,25 @@
 # --------------------------------------------------------------------------
 from collections import defaultdict
 from logging import getLogger
-from typing import List, Union
+from typing import Dict, List, Optional, Union
 
+from onnx import NodeProto
 from onnx_model import OnnxModel
 
 logger = getLogger(__name__)
 
 
 class Fusion:
+    """
+    Base class for Graph Fusion
+    """
+
     def __init__(
         self,
         model: OnnxModel,
         fused_op_type: str,
         search_op_types: Union[str, List[str]],
-        description: str = None,
+        description: str = "",
     ):
         self.search_op_types: List[str] = [search_op_types] if isinstance(search_op_types, str) else search_op_types
         self.fused_op_type: str = fused_op_type
@@ -27,14 +32,30 @@ class Fusion:
         self.nodes_to_add: List = []
         self.prune_graph: bool = False
         self.node_name_to_graph_name: dict = {}
-        self.this_graph_name: str = None
+        self.this_graph_name: Optional[str] = None
         # It is optional that subclass updates fused_count since we will also check nodes_to_add to get counter.
         self.fused_count: defaultdict = defaultdict(int)
 
-    def increase_counter(self, fused_op_name):
+    def increase_counter(self, fused_op_name: str):
+        """
+        Increase counter of a fused operator.
+        """
         self.fused_count[fused_op_name] += 1
 
+    def fuse(
+        self,
+        node: NodeProto,
+        input_name_to_nodes: Dict[str, List[NodeProto]],
+        output_name_to_node: Dict[str, NodeProto],
+    ):
+        """Interface for fusion that starts from a node"""
+        raise NotImplementedError
+
     def apply(self):
+        """
+        Apply graph fusion on the whole model graph.
+        It searched nodes of given operators, and start fusion on each of those nodes.
+        """
         logger.debug(f"start {self.description} fusion...")
         input_name_to_nodes = self.model.input_name_to_nodes()
         output_name_to_node = self.model.output_name_to_node()
@@ -44,7 +65,7 @@ class Fusion:
             for node in self.model.get_nodes_by_op_type(search_op_type):
                 graph = self.model.get_graph_by_node(node)
                 if graph is None:
-                    raise Exception("Can not find node in any graphs")
+                    raise Exception("Can not find node in any graph")
                 self.this_graph_name = graph.name
                 self.fuse(node, input_name_to_nodes, output_name_to_node)
 

onnxruntime/python/tools/transformers/fusion_layernorm.py

@@ -80,7 +80,7 @@ class FusionLayerNormalization(Fusion):
         second_add_node = parent_nodes[1]
         i, add_weight = self.model.get_constant_input(second_add_node)
         if add_weight is None or add_weight <= 0 or add_weight > 1.0e-4:
-            logger.warning(f"epsilon value is not expeced: {add_weight}")
+            logger.debug(f"skip SkipLayerNormalization fusion since epsilon value is not expected: {add_weight}")
             return
 
         pow_node = parent_nodes[3]

onnxruntime/python/tools/transformers/fusion_reshape.py

@@ -7,7 +7,7 @@ from logging import getLogger
 
 import numpy as np
 from fusion_base import Fusion
-from onnx import TensorProto, helper, numpy_helper
+from onnx import TensorProto, helper
 from onnx_model import OnnxModel
 
 logger = getLogger(__name__)
@@ -83,7 +83,7 @@ class FusionReshape(Fusion):
         path2 = []
         path3 = []
         shape_nodes = [shape_0, shape_1]
-        if len(concat_node.input) == 3 and self.model.get_initializer(concat_node.input[2]) is None:
+        if len(concat_node.input) == 3 and self.model.get_constant_value(concat_node.input[2]) is None:
             path2 = self.model.match_parent_path(
                 concat_node,
                 ["Unsqueeze", "Mul", "Gather", "Shape"],
@@ -149,11 +149,10 @@ class FusionReshape(Fusion):
             shape_nodes.extend([path2[-1]])
             shape.append(-1)
         elif len(concat_node.input) > 3:
-            concat_3 = self.model.get_initializer(concat_node.input[3])
-            if concat_3 is None:
+            concat_value = self.model.get_constant_value(concat_node.input[3])
+            if concat_value is None:
                 return
 
-            concat_value = numpy_helper.to_array(concat_3)
             if isinstance(concat_value, np.ndarray):
                 shape.extend(concat_value.tolist())
             else:

onnxruntime/python/tools/transformers/fusion_skiplayernorm.py

@@ -49,6 +49,10 @@ class FusionSkipLayerNormalization(Fusion):
         if len(self.model.get_parents(add)) != 2:
             return
 
+        # To avoid an Add node have two children of LayerNormalization, we shall only fuse one SkipLayerNormalization
+        if add in self.nodes_to_remove:
+            return
+
         # Root Mean Square Layer Normalization
         simplified = node.op_type == "SimplifiedLayerNormalization"
 

onnxruntime/python/tools/transformers/onnx_model.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
+import itertools
 import logging
 import os
 import sys
@@ -1183,3 +1184,48 @@ class OnnxModel:
 
     def clean_shape_infer(self):
         self.model.graph.ClearField("value_info")
+
+    def use_float16(self):
+        """Check whether the model uses float16"""
+        queue = []  # queue for BFS
+        queue.append(self.model.graph)
+        while queue:
+            sub_graphs = []
+            for graph in queue:
+                if not isinstance(graph, GraphProto):
+                    continue
+
+                for v in itertools.chain(graph.input, graph.output, graph.value_info):
+                    if v.type.tensor_type.elem_type == TensorProto.FLOAT16:
+                        return True
+                    if v.type.HasField("sequence_type"):
+                        if v.type.sequence_type.elem_type.tensor_type.elem_type == TensorProto.FLOAT16:
+                            return True
+
+                for t in graph.initializer:
+                    if t.data_type == TensorProto.FLOAT16:
+                        return True
+
+                for node in graph.node:
+                    if node.op_type == "Cast":
+                        for attr in node.attribute:
+                            if attr.name == "to" and attr.i == TensorProto.FLOAT16:
+                                return True
+
+                    for attr in node.attribute:
+                        if attr.type == AttributeProto.GRAPH:
+                            sub_graphs.append(attr.g)
+
+                        for g in attr.graphs:
+                            sub_graphs.append(g)
+
+                        if isinstance(attr.t, TensorProto) and attr.t.data_type == TensorProto.FLOAT16:
+                            return True
+
+                        for t in attr.tensors:
+                            if isinstance(t, TensorProto) and t.data_type == TensorProto.FLOAT16:
+                                return True
+
+            queue = sub_graphs
+
+        return False

onnxruntime/python/tools/transformers/optimizer.py

@@ -26,6 +26,7 @@ from typing import Dict, Optional
 import coloredlogs
 from fusion_options import FusionOptions
 from onnx import ModelProto, load_model
+from onnx_model import OnnxModel
 from onnx_model_bart import BartOnnxModel
 from onnx_model_bert import BertOnnxModel
 from onnx_model_bert_keras import BertOnnxModelKeras
@@ -45,20 +46,16 @@ MODEL_TYPES = {
     "bert": (BertOnnxModel, "pytorch", 1),
     "bert_tf": (BertOnnxModelTF, "tf2onnx", 0),
     "bert_keras": (BertOnnxModelKeras, "keras2onnx", 0),
+    "clip": (ClipOnnxModel, "pytorch", 1),  # Clip in Stable Diffusion
     "gpt2": (Gpt2OnnxModel, "pytorch", 1),
-    "gpt2_tf": (
-        Gpt2OnnxModel,
-        "tf2onnx",
-        0,
-    ),  # might add a class for GPT2OnnxModel for TF later.
+    "gpt2_tf": (Gpt2OnnxModel, "tf2onnx", 0),  # might add a class for GPT2OnnxModel for TF later.
+    "gpt_neox": (BertOnnxModel, "pytorch", 0),  # GPT-NeoX
+    "swin": (BertOnnxModel, "pytorch", 1),
     "tnlr": (TnlrOnnxModel, "pytorch", 1),
     "t5": (T5OnnxModel, "pytorch", 2),
-    # Stable Diffusion models
-    "unet": (UnetOnnxModel, "pytorch", 1),
-    "vae": (VaeOnnxModel, "pytorch", 1),
-    "clip": (ClipOnnxModel, "pytorch", 1),
+    "unet": (UnetOnnxModel, "pytorch", 1),  # UNet in Stable Diffusion
+    "vae": (VaeOnnxModel, "pytorch", 1),  # UAE in Stable Diffusion
     "vit": (BertOnnxModel, "pytorch", 1),
-    "swin": (BertOnnxModel, "pytorch", 1),
 }
 
 
@@ -93,6 +90,15 @@ def optimize_by_onnxruntime(
         logger.error("There is no gpu for onnxruntime to do optimization.")
         return onnx_model_path
 
+    model = OnnxModel(load_model(onnx_model_path, format=None, load_external_data=False))
+    if model.use_float16() and not use_gpu:
+        logger.warning(
+            "This model uses float16 in the graph, use_gpu=False might cause extra Cast nodes. "
+            "Most operators have no float16 implementation in CPU, so Cast nodes are added to compute them in float32. "
+            "If the model is intended to use in GPU, please set use_gpu=True. "
+            "Otherwise, consider exporting onnx in float32 and optional int8 quantization for better performance. "
+        )
+
     sess_options = onnxruntime.SessionOptions()
     if opt_level == 1:
         sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
@@ -120,15 +126,13 @@ def optimize_by_onnxruntime(
     else:
         gpu_ep = []
 
-        if torch_version.cuda:
-            gpu_ep.append("CUDAExecutionProvider")
-        elif torch_version.hip:
+        if torch_version.hip:
             gpu_ep.append("MIGraphXExecutionProvider")
             gpu_ep.append("ROCMExecutionProvider")
+        else:
+            gpu_ep.append("CUDAExecutionProvider")
+
         onnxruntime.InferenceSession(onnx_model_path, sess_options, providers=gpu_ep, **kwargs)
-        assert not set(onnxruntime.get_available_providers()).isdisjoint(
-            ["CUDAExecutionProvider", "ROCMExecutionProvider", "MIGraphXExecutionProvider"]
-        )
 
     assert os.path.exists(optimized_model_path) and os.path.isfile(optimized_model_path)
     logger.debug("Save optimized model by onnxruntime to %s", optimized_model_path)
@@ -279,10 +283,12 @@ def optimize_model(
         )
     elif opt_level == 1:
         # basic optimizations (like constant folding and cast elimination) are not specified to execution provider.
-        # CPU provider is used here so that there is no extra node for GPU memory copy.
+        # Note that use_gpu=False might cause extra Cast nodes for float16 model since most operators does not support float16 in CPU.
+        # Sometime, use_gpu=True might cause extra memory copy nodes when some operators are supported only in CPU.
+        # We might need remove GPU memory copy nodes as preprocess of optimize_by_fusion if they cause no matching in fusion.
         temp_model_path = optimize_by_onnxruntime(
             input,
-            use_gpu=False,
+            use_gpu=use_gpu,
             opt_level=1,
             disabled_optimizers=disabled_optimizers,
             verbose=verbose,