From 95f96287762f3626f68fd09834f0b4db35da4abf Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 18 Jul 2023 10:29:08 -0700 Subject: [PATCH] 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. --- .../python/tools/transformers/fusion_base.py | 31 +++++++++++-- .../tools/transformers/fusion_layernorm.py | 2 +- .../tools/transformers/fusion_reshape.py | 9 ++-- .../transformers/fusion_skiplayernorm.py | 4 ++ .../python/tools/transformers/onnx_model.py | 46 +++++++++++++++++++ .../python/tools/transformers/optimizer.py | 42 +++++++++-------- 6 files changed, 105 insertions(+), 29 deletions(-) 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,