diff --git a/onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py index c6564adce8..efce064ac5 100644 --- a/onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py @@ -76,6 +76,12 @@ def parse_arguments(): parser.add_argument('-w', '--overwrite', required=False, action='store_true', help="overwrite existing ONNX model") parser.set_defaults(overwrite=False) + parser.add_argument('--disable_auto_mixed_precision', + required=False, + action='store_true', + help="use pure fp16 instead of mixed precision") + parser.set_defaults(disable_auto_mixed_precision=False) + args = parser.parse_args() return args @@ -91,7 +97,8 @@ def export_onnx_models(model_name_or_path, verbose, use_decoder_start_token: bool = True, merge_encoder_and_decoder_init: bool = True, - overwrite: bool = False): + overwrite: bool = False, + disable_auto_mixed_precision: bool = False): device = torch.device("cuda:0" if use_gpu else "cpu") models = T5Helper.load_model(model_name_or_path, cache_dir, device, merge_encoder_and_decoder_init) @@ -102,6 +109,7 @@ def export_onnx_models(model_name_or_path, output_paths = [] for name, model in models.items(): + model.to(device) filename_suffix = "_" + name onnx_path = T5Helper.get_onnx_path(output_dir, @@ -112,7 +120,8 @@ def export_onnx_models(model_name_or_path, if overwrite or not os.path.exists(onnx_path): logger.info(f"Exporting ONNX model to {onnx_path}") # We have to clone model before exporting onnx, otherwise verify_onnx will report large difference. - T5Helper.export_onnx(copy.deepcopy(model), + cloned_model = copy.deepcopy(model).to(device) + T5Helper.export_onnx(cloned_model, device, onnx_path, verbose, @@ -130,8 +139,13 @@ def export_onnx_models(model_name_or_path, if overwrite or not os.path.exists(output_path): logger.info(f"Optimizing model to {output_path}") - T5Helper.optimize_onnx(onnx_path, output_path, precision == Precision.FLOAT16, config.num_heads, - config.hidden_size, use_external_data_format) + T5Helper.optimize_onnx(onnx_path, + output_path, + precision == Precision.FLOAT16, + config.num_heads, + config.hidden_size, + use_external_data_format, + auto_mixed_precision=not disable_auto_mixed_precision) else: logger.info(f"Skip optimizing: existed ONNX model {onnx_path}") else: @@ -175,7 +189,7 @@ def main(): output_paths = export_onnx_models(args.model_name_or_path, cache_dir, output_dir, args.use_gpu, args.use_external_data_format, args.optimize_onnx, args.precision, args.verbose, args.use_decoder_start_token, merge_encoder_and_decoder_init, - args.overwrite) + args.overwrite, args.disable_auto_mixed_precision) logger.info(f"Done! Outputs: {output_paths}") diff --git a/onnxruntime/python/tools/transformers/models/t5/t5_decoder.py b/onnxruntime/python/tools/transformers/models/t5/t5_decoder.py index 87c0019fbc..b4c44fc10e 100644 --- a/onnxruntime/python/tools/transformers/models/t5/t5_decoder.py +++ b/onnxruntime/python/tools/transformers/models/t5/t5_decoder.py @@ -18,6 +18,7 @@ from past_helper import PastKeyValuesHelper sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) from torch_onnx_export_helper import torch_onnx_export +from io_binding_helper import TypeHelper logger = logging.getLogger(__name__) @@ -26,6 +27,7 @@ class T5DecoderInit(torch.nn.Module): """ A T5 decoder with LM head to create initial past key values. This model is only called once during starting decoding. """ + def __init__(self, decoder: torch.nn.Module, lm_head: torch.nn.Module, @@ -62,6 +64,7 @@ class T5DecoderInit(torch.nn.Module): class T5Decoder(torch.nn.Module): """ A T5 decoder with LM head and past key values""" + def __init__(self, decoder, lm_head, config): super().__init__() self.decoder = decoder @@ -92,6 +95,7 @@ class T5Decoder(torch.nn.Module): class T5DecoderInputs: + def __init__(self, decoder_input_ids, encoder_attention_mask, encoder_hidden_states, past_key_values=None): self.decoder_input_ids: torch.LongTensor = decoder_input_ids self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask @@ -99,8 +103,12 @@ class T5DecoderInputs: self.past_key_values: Union[List[torch.FloatTensor], List[torch.HalfTensor], None] = past_key_values @staticmethod - def create_dummy(config: T5Config, batch_size: int, encode_sequence_length: int, past_decode_sequence_length: int, - device: torch.device): # -> T5DecoderInputs: + def create_dummy(config: T5Config, + batch_size: int, + encode_sequence_length: int, + past_decode_sequence_length: int, + device: torch.device, + float16: bool = False): # -> T5DecoderInputs: """ Create dummy inputs for T5Decoder. Args: @@ -109,6 +117,7 @@ class T5DecoderInputs: encode_sequence_length (int): sequence length of input_ids for encoder past_decode_sequence_length (int): past sequence length of input_ids for decoder device (torch.device): device of output tensors + float16 (bool): whether the model uses float32 or float16 in input Returns: T5DecoderInputs: dummy inputs for decoder @@ -127,10 +136,11 @@ class T5DecoderInputs: encoder_inputs = T5EncoderInputs.create_dummy(batch_size, encode_sequence_length, vocab_size, device) + float_type = torch.float16 if float16 else torch.float32 encoder_hidden_state = torch.rand(batch_size, encode_sequence_length, hidden_size, - dtype=torch.float32, + dtype=float_type, device=device) if past_decode_sequence_length > 0: @@ -145,10 +155,10 @@ class T5DecoderInputs: past = [] for _ in range(2 * num_layers): - past.append(torch.rand(self_attention_past_shape, dtype=torch.float32, device=device)) + past.append(torch.rand(self_attention_past_shape, dtype=float_type, device=device)) for _ in range(2 * num_layers): - past.append(torch.rand(cross_attention_past_shape, dtype=torch.float32, device=device)) + past.append(torch.rand(cross_attention_past_shape, dtype=float_type, device=device)) else: past = None @@ -160,8 +170,15 @@ class T5DecoderInputs: input_list.extend(self.past_key_values) return input_list + def to_fp32(self): + encoder_hidden_state = self.encoder_hidden_states.to(dtype=torch.float32) + past = [p.to(dtype=torch.float32) for p in self.past_key_values] + return T5DecoderInputs(self.decoder_input_ids.clone(), self.encoder_attention_mask.clone(), + encoder_hidden_state, past) + class T5DecoderHelper: + @staticmethod def export_onnx(decoder: Union[T5Decoder, T5DecoderInit], device: torch.device, @@ -292,6 +309,8 @@ class T5DecoderHelper: max_cases=4): """ Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good. """ + float16: bool = (TypeHelper.get_input_type(ort_session, "encoder_hidden_states") == "tensor(float16)") + test_cases = [(4, 11, 3), (1, 2, 5), (3, 1, 1), (8, 5, 2)] test_cases_max_diff = [] for (batch_size, encode_sequence_length, past_decode_sequence_length) in test_cases[:max_cases]: @@ -302,8 +321,11 @@ class T5DecoderHelper: batch_size, encode_sequence_length, past_decode_sequence_length, - device=device) - input_list = inputs.to_list() + device=device, + float16=float16) + + # We use fp32 PyTroch model as baseline even when ONNX model is fp16 + input_list = inputs.to_fp32().to_list() # Run inference of PyTorch model with torch.no_grad(): diff --git a/onnxruntime/python/tools/transformers/models/t5/t5_helper.py b/onnxruntime/python/tools/transformers/models/t5/t5_helper.py index 4bcb5d4284..3d4d4ab448 100644 --- a/onnxruntime/python/tools/transformers/models/t5/t5_helper.py +++ b/onnxruntime/python/tools/transformers/models/t5/t5_helper.py @@ -7,7 +7,7 @@ import os import sys from pathlib import Path -from typing import Union, Dict +from typing import Union, Dict, List import logging import torch from transformers import T5ForConditionalGeneration @@ -16,12 +16,19 @@ from t5_encoder import T5Encoder, T5EncoderHelper from t5_decoder import T5DecoderInit, T5Decoder, T5DecoderHelper from t5_encoder_decoder_init import T5EncoderDecoderInit, T5EncoderDecoderInitHelper +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) +from onnx_model import OnnxModel +from float16 import float_to_float16_max_diff +from fusion_utils import FusionUtils +from optimizer import optimize_model + logger = logging.getLogger(__name__) PRETRAINED_T5_MODELS = ["t5-small", "t5-base", "t5-large", "t5-3B", "t5-11B"] class T5Helper: + @staticmethod def get_onnx_path(output_dir: str, model_name_or_path: str, suffix: str = "", new_folder: bool = False) -> str: """Build onnx path @@ -96,17 +103,82 @@ class T5Helper: else: T5DecoderHelper.export_onnx(model, device, onnx_model_path, verbose, use_external_data_format) + @staticmethod + def auto_mixed_precision( + onnx_model: OnnxModel, + op_block_list: List[str] = ["Pow", "ReduceMean", "Add", "Sqrt", "Div", "Mul", "Softmax", "Relu"]): + """Convert model to mixed precision. + It detects whether original model has fp16 precision weights, and set parameters for float16 conversion automatically. + Args: + onnx_model (OnnxModel): optimized ONNX model + op_block_list (List[str], optional): . Defaults to ["Pow", "ReduceMean", "Add", "Sqrt", "Div", "Mul", "Softmax", "Relu"] + Returns: + parameters(dict): a dictionary of parameters used in float16 conversion + """ + op_full_set = set([node.op_type for node in onnx_model.nodes()]) + fp32_op_set = set(op_block_list) + fp16_op_set = op_full_set.difference(fp32_op_set) + logger.info(f"fp32 op: {fp32_op_set} fp16 op: {fp16_op_set}") + + # logits is the first output + logits_output_name = onnx_model.graph().output[0].name + + # We use the weight in last MatMul node to detect whether the model is stored with float16 weights from training. + is_weight_fp16_precision = False + output_name_to_node = onnx_model.output_name_to_node() + assert logits_output_name in output_name_to_node + node = output_name_to_node[logits_output_name] + last_matmul_node = None + if node.op_type == "MatMul": + last_matmul_node = node + logger.info(f"Found last MatMul node for logits: {node.name}") + initializer = None + for input in node.input: + initializer = onnx_model.get_initializer(input) + if initializer is not None: + break + + # when the max difference of value after converting float to float16 is lower than a threshold (1e-6), + # we can deduce that the weights are stored in float16 precision. + max_diff = float_to_float16_max_diff(initializer) + logger.debug(f"max diff of converting weights in last MatMul node {node.name}: {max_diff}") + is_weight_fp16_precision = (max_diff < 1E-6) + else: + logger.warning(f"Failed to find MatMul node for logits. Found {node.op_type} of node {node.name}") + + keep_io_types = [] + node_block_list = [] + if (not is_weight_fp16_precision) and (last_matmul_node is not None): + # When original weight is float32 precision, keep logits and last MatMul in float32 could get better precision. + keep_io_types = [logits_output_name] + node_block_list = [last_matmul_node.name] + + parameters = { + "keep_io_types": keep_io_types, + "op_block_list": op_block_list, + "node_block_list": node_block_list, + "force_fp16_initializers": is_weight_fp16_precision + } + + 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 def optimize_onnx(onnx_model_path: str, optimized_model_path: str, is_float16: bool, num_attention_heads: int, hidden_size: int, - use_external_data_format: bool = False): + use_external_data_format: bool = False, + auto_mixed_precision: bool = True): """ Optimize ONNX model with an option to convert it to use mixed precision. """ - sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) - from optimizer import optimize_model m = optimize_model(onnx_model_path, model_type='bert', num_heads=num_attention_heads, @@ -115,7 +187,10 @@ class T5Helper: optimization_options=None, use_gpu=False) if is_float16: - m.convert_model_float32_to_float16(cast_input_output=False) + if auto_mixed_precision: + T5Helper.auto_mixed_precision(m) + else: + m.convert_model_float32_to_float16(cast_input_output=False) m.save_model_to_file(optimized_model_path, use_external_data_format)