diff --git a/onnxruntime/python/tools/transformers/fusion_utils.py b/onnxruntime/python/tools/transformers/fusion_utils.py index 9f2fd05172..a3ae4fb60e 100644 --- a/onnxruntime/python/tools/transformers/fusion_utils.py +++ b/onnxruntime/python/tools/transformers/fusion_utils.py @@ -105,45 +105,6 @@ class FusionUtils: else: return value == expected_value - def get_dtype(self, shape_infer_helper, input_or_output_name: str) -> int: - """Get data type of an input or output. - - Args: - shape_infer_helper (SymbolicShapeInferenceHelper): object of symbolic shape inference - input_or_output_name (str): name of input or output - - Returns: - int: tensor data type - """ - dtype = self.model.get_dtype(input_or_output_name) - if dtype is not None: - return dtype - - if shape_infer_helper: - tensor_proto = shape_infer_helper.known_vi_[input_or_output_name] - if tensor_proto.type.tensor_type.HasField("elem_type"): - return tensor_proto.type.tensor_type.elem_type - - return None - - def remove_cascaded_cast_nodes(self): - """Remove Cast node that are overrided by another Cast node like --> Cast --> Cast --> - Note that this shall be used carefully since it might introduce semantic change. - For example, float -> int -> float could get different value than the original float value. - So, it is recommended to used only in post-processing of mixed precision conversion. - """ - removed_count = 0 - for node in self.model.nodes(): - if node.op_type == "Cast": - parent = self.model.get_parent(node, 0) - if parent and parent.op_type == "Cast": - node.input[0] = parent.input[0] - removed_count += 1 - - if removed_count > 0: - logger.info(f"Removed {removed_count} cascaded Cast nodes") - self.model.prune_graph() - def remove_identity_nodes(self): """Remove Identity nodes, except those right before graph output.""" nodes_to_remove = [] @@ -157,34 +118,11 @@ class FusionUtils: self.model.remove_nodes(nodes_to_remove) logger.info(f"Removed {len(nodes_to_remove)} Identity nodes") + def remove_cascaded_cast_nodes(self): + self.model.remove_cascaded_cast_nodes() + def remove_useless_cast_nodes(self): - """Remove cast nodes that are not needed: input and output has same data type.""" - shape_infer = self.model.infer_runtime_shape(update=True) - if shape_infer is None: - return - - nodes_to_remove = [] - for node in self.model.nodes(): - if node.op_type == "Cast": - input_dtype = self.get_dtype(shape_infer, node.input[0]) - output_dtype = self.get_dtype(shape_infer, node.output[0]) - if input_dtype and input_dtype == output_dtype: - nodes_to_remove.append(node) - - if nodes_to_remove: - graph_input_names = set(self.model.get_graphs_input_names()) - graph_output_names = set(self.model.get_graphs_output_names()) - for node in nodes_to_remove: - if bool(set(node.output) & graph_output_names): - if not bool(set(node.input) & graph_input_names): - self.model.replace_output_of_all_nodes(node.input[0], node.output[0]) - else: - continue - else: - self.model.replace_input_of_all_nodes(node.output[0], node.input[0]) - self.model.remove_node(node) - - logger.info(f"Removed {len(nodes_to_remove)} Cast nodes with output type same as input") + self.model.remove_useless_cast_nodes() def remove_useless_reshape_nodes(self): """Remove reshape node that is not needed based on symbolic shape inference: input and output has same shape""" diff --git a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py index df84574f56..a154a52858 100644 --- a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py @@ -24,7 +24,6 @@ sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) from benchmark_helper import Precision from float16 import float_to_float16_max_diff -from fusion_utils import FusionUtils from io_binding_helper import IOBindingHelper from onnx_model import OnnxModel from torch_onnx_export_helper import torch_onnx_export @@ -599,10 +598,6 @@ class Gpt2Helper: logger.info(f"auto_mixed_precision parameters: {parameters}") onnx_model.convert_float_to_float16(use_symbolic_shape_infer=True, **parameters) - fusion_utils = FusionUtils(onnx_model) - fusion_utils.remove_cascaded_cast_nodes() - fusion_utils.remove_useless_cast_nodes() - return parameters @staticmethod diff --git a/onnxruntime/python/tools/transformers/models/t5/t5_helper.py b/onnxruntime/python/tools/transformers/models/t5/t5_helper.py index d4d0b82625..4d853a6544 100644 --- a/onnxruntime/python/tools/transformers/models/t5/t5_helper.py +++ b/onnxruntime/python/tools/transformers/models/t5/t5_helper.py @@ -20,7 +20,6 @@ from onnxruntime import InferenceSession sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) from float16 import float_to_float16_max_diff -from fusion_utils import FusionUtils from onnx_model import OnnxModel from optimizer import optimize_model @@ -218,10 +217,6 @@ class T5Helper: logger.info(f"auto_mixed_precision parameters: {parameters}") onnx_model.convert_float_to_float16(use_symbolic_shape_infer=True, **parameters) - fusion_utils = FusionUtils(onnx_model) - fusion_utils.remove_cascaded_cast_nodes() - fusion_utils.remove_useless_cast_nodes() - return parameters @staticmethod diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index 2ac06ba6b5..2e8cb5048f 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -8,7 +8,7 @@ import os import sys from collections import deque from pathlib import Path -from typing import Dict, List, Tuple +from typing import Dict, List, Optional, Tuple import numpy as np from onnx import AttributeProto, GraphProto, ModelProto, NodeProto, TensorProto, helper, numpy_helper, save_model @@ -24,9 +24,9 @@ class OnnxModel: def initialize(self, model): self.model: ModelProto = model self._node_name_suffix: Dict[str, int] = {} # key is node name prefix, value is the last suffix generated - self.shape_infer_helper = None - self.enable_shape_infer = True - self.all_graphs = None + self.shape_infer_helper: SymbolicShapeInferenceHelper = None + self.enable_shape_infer: bool = True + self.all_graphs: Optional[List[GraphProto]] = None def disable_shape_inference(self): self.enable_shape_infer = False @@ -503,6 +503,62 @@ class OnnxModel: return value return None + def remove_cascaded_cast_nodes(self): + """Remove Cast node that are followed by another Cast node like --> Cast --> Cast --> + Note that this shall be used carefully since it might introduce semantic change. + For example, float -> int -> float could get different value than the original float value. + So, it is recommended to used only in post-processing of mixed precision conversion. + """ + removed_count = 0 + for node in self.nodes(): + if node.op_type == "Cast": + parent = self.get_parent(node, 0) + if parent and parent.op_type == "Cast": + node.input[0] = parent.input[0] + removed_count += 1 + + if removed_count > 0: + logger.info("Removed %d cascaded Cast nodes", removed_count) + self.prune_graph() + + def remove_useless_cast_nodes(self): + """Remove cast nodes that are not needed: input and output has same data type.""" + shape_infer = self.infer_runtime_shape(update=True) + if shape_infer is None: + logger.info(f"Skip removing useless cast nodes since shape inference failed.") + return + + def get_data_type(input_or_output_name): + dtype = self.get_dtype(input_or_output_name) + if dtype: + return dtype + if shape_infer.known_vi_[input_or_output_name].type.tensor_type.HasField("elem_type"): + return shape_infer.known_vi_[input_or_output_name].type.tensor_type.elem_type + return None + + nodes_to_remove = [] + for node in self.nodes(): + if node.op_type == "Cast": + input_dtype = get_data_type(node.input[0]) + output_dtype = get_data_type(node.output[0]) + if input_dtype and input_dtype == output_dtype: + nodes_to_remove.append(node) + + if nodes_to_remove: + graph_input_names = set(self.get_graphs_input_names()) + graph_output_names = set(self.get_graphs_output_names()) + for node in nodes_to_remove: + if bool(set(node.output) & graph_output_names): + if not bool(set(node.input) & graph_input_names): + self.replace_output_of_all_nodes(node.input[0], node.output[0]) + else: + continue + else: + self.replace_input_of_all_nodes(node.output[0], node.input[0]) + self.remove_node(node) + + logger.info("Removed %d Cast nodes with output type same as input", len(nodes_to_remove)) + def convert_model_float32_to_float16(self, cast_input_output=True): logger.warning( "The function convert_model_float32_to_float16 is deprecated. Use convert_float_to_float16 instead!" @@ -562,37 +618,9 @@ class OnnxModel: fp16_model = convert_float_to_float16(model, **parameters) self.initialize(fp16_model) - # Convert_float_to_float16 might add Cast(to=10) --> Cast(to=1) when two consequent nodes are computed in FP32. - # Below are post-processing that removes those Cast nodes. - # Remove first Cast nodes in path like --> Cast --> Cast --> - nodes_to_remove = [] - for node in self.nodes(): - if node.op_type == "Cast": - parent = self.get_parent(node, 0) - if parent and parent.op_type == "Cast": - if self.get_children(parent) == 1: # cannot be removed if its output is used by multiple nodes - self.replace_input_of_all_nodes(parent.output[0], parent.input[0]) - nodes_to_remove.append(parent) + self.remove_cascaded_cast_nodes() - # Remove the second cast node. - for node in self.nodes(): - if ( - node.op_type == "Cast" - and OnnxModel.get_node_attribute(node, "to") == int(TensorProto.FLOAT) - and self.get_dtype(node.input[0]) == int(TensorProto.FLOAT) - ): - - if self.find_graph_output(node.output[0]): - self.replace_output_of_all_nodes(node.input[0], node.output[0]) - else: - self.replace_input_of_all_nodes(node.output[0], node.input[0]) - nodes_to_remove.append(node) - - self.remove_nodes(nodes_to_remove) - - if nodes_to_remove: - self.prune_graph() - print(f"removed {len(nodes_to_remove)} Cast nodes from float16 model") + self.remove_useless_cast_nodes() def create_node_name(self, op_type, name_prefix=None): """Create a unique node name that starts with a prefix (default is operator type).