diff --git a/onnxruntime/python/tools/transformers/fusion_base.py b/onnxruntime/python/tools/transformers/fusion_base.py index 43969aa04b..d53a2f4ba4 100644 --- a/onnxruntime/python/tools/transformers/fusion_base.py +++ b/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) diff --git a/onnxruntime/python/tools/transformers/fusion_layernorm.py b/onnxruntime/python/tools/transformers/fusion_layernorm.py index e817e1a589..ec485e0dfa 100644 --- a/onnxruntime/python/tools/transformers/fusion_layernorm.py +++ b/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] diff --git a/onnxruntime/python/tools/transformers/fusion_reshape.py b/onnxruntime/python/tools/transformers/fusion_reshape.py index 853038f746..b872347738 100644 --- a/onnxruntime/python/tools/transformers/fusion_reshape.py +++ b/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: diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 2737369361..4b771c5bee 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/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" diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index ead61df9f3..8e1e21e9f8 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/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 diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py index e1bff135db..8719380ff2 100644 --- a/onnxruntime/python/tools/transformers/optimizer.py +++ b/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,