diff --git a/onnxruntime/python/tools/transformers/models/t5/__init__.py b/onnxruntime/python/tools/transformers/models/t5/__init__.py new file mode 100644 index 0000000000..ad5632855c --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/t5/__init__.py @@ -0,0 +1,4 @@ +import os +import sys + +sys.path.append(os.path.dirname(__file__)) diff --git a/onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py new file mode 100644 index 0000000000..c6564adce8 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py @@ -0,0 +1,184 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import os +import sys +import argparse +import logging +import torch +import copy +from t5_helper import PRETRAINED_T5_MODELS, T5Helper + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) +from benchmark_helper import setup_logger, prepare_environment, create_onnxruntime_session, Precision + +logger = logging.getLogger('') + + +def parse_arguments(): + parser = argparse.ArgumentParser() + + parser.add_argument('-m', + '--model_name_or_path', + required=False, + default=PRETRAINED_T5_MODELS[0], + type=str, + help='Model path, or pretrained model name in the list: ' + ', '.join(PRETRAINED_T5_MODELS)) + + parser.add_argument('--cache_dir', + required=False, + type=str, + default=os.path.join('.', 'cache_models'), + help='Directory to cache pre-trained models') + + parser.add_argument('--output', + required=False, + type=str, + default=os.path.join('.', 'onnx_models'), + help='Output directory') + + parser.add_argument('-o', + '--optimize_onnx', + required=False, + action='store_true', + help='Use optimizer.py to optimize onnx model') + parser.set_defaults(optimize_onnx=False) + + parser.add_argument('--use_gpu', required=False, action='store_true', help="use GPU for inference") + parser.set_defaults(use_gpu=False) + + parser.add_argument("-p", + "--precision", + required=False, + type=Precision, + default=Precision.FLOAT32, + choices=[Precision.FLOAT32, Precision.FLOAT16], + help="Precision of model to run. fp32 for full precision, fp16 for half precision") + + parser.add_argument('--verbose', required=False, action='store_true') + parser.set_defaults(verbose=False) + + parser.add_argument('-e', '--use_external_data_format', required=False, action='store_true') + parser.set_defaults(use_external_data_format=False) + + parser.add_argument( + '-s', + '--use_decoder_start_token', + required=False, + action='store_true', + help="Use config.decoder_start_token_id in decoding. Otherwise, add an extra graph input for decoder_input_ids." + ) + parser.set_defaults(use_decoder_start_token=False) + + parser.add_argument('-w', '--overwrite', required=False, action='store_true', help="overwrite existing ONNX model") + parser.set_defaults(overwrite=False) + + args = parser.parse_args() + + return args + + +def export_onnx_models(model_name_or_path, + cache_dir, + output_dir, + use_gpu, + use_external_data_format, + optimize_onnx, + precision, + verbose, + use_decoder_start_token: bool = True, + merge_encoder_and_decoder_init: bool = True, + overwrite: 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) + config = models["decoder"].config + + if (not use_external_data_format) and (config.num_layers > 24): + logger.info(f"Try use_external_data_format when model size > 2GB") + + output_paths = [] + for name, model in models.items(): + filename_suffix = "_" + name + + onnx_path = T5Helper.get_onnx_path(output_dir, + model_name_or_path, + suffix=filename_suffix, + new_folder=use_external_data_format) + + 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), + device, + onnx_path, + verbose, + use_external_data_format, + use_decoder_input_ids=not use_decoder_start_token) + else: + logger.info(f"Skip exporting: existed ONNX model {onnx_path}") + + # Optimize ONNX graph. Note that we have not implemented graph optimization for T5 yet. + if optimize_onnx or precision != Precision.FLOAT32: + output_path = T5Helper.get_onnx_path(output_dir, + model_name_or_path, + suffix=filename_suffix + "_" + str(precision), + new_folder=use_external_data_format) + + 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) + else: + logger.info(f"Skip optimizing: existed ONNX model {onnx_path}") + else: + output_path = onnx_path + + ort_session = create_onnxruntime_session( + output_path, + use_gpu=use_gpu, + provider=['CUDAExecutionProvider', 'CPUExecutionProvider'] if use_gpu else ['CPUExecutionProvider']) + max_diff = T5Helper.verify_onnx(model, ort_session, device) + logger.info(f'PyTorch and OnnxRuntime results max difference = {max_diff}') + if max_diff > 1e-4: + logger.warn(f'PyTorch and OnnxRuntime results are NOT close') + + output_paths.append(output_path) + + return output_paths + + +def main(): + args = parse_arguments() + setup_logger(args.verbose) + + logger.info(f"Arguments:{args}") + + cache_dir = args.cache_dir + output_dir = args.output if not args.output.endswith(".onnx") else os.path.dirname(args.output) + prepare_environment(cache_dir, output_dir, args.use_gpu) + + if args.precision != Precision.FLOAT32: + assert args.optimize_onnx, "fp16/int8 requires --optimize_onnx" + + if args.precision == Precision.FLOAT16: + assert args.use_gpu, "fp16 requires --use_gpu" + + if args.optimize_onnx: + logger.warn(f'Graph optimization for T5 is not implemented yet.') + + with torch.no_grad(): + merge_encoder_and_decoder_init = True # Merge encoder and decoder initialization into one model is recommended. + 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) + + logger.info(f"Done! Outputs: {output_paths}") + + +if __name__ == '__main__': + main() diff --git a/onnxruntime/python/tools/transformers/models/t5/past_helper.py b/onnxruntime/python/tools/transformers/models/t5/past_helper.py new file mode 100644 index 0000000000..3c585c23c8 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/t5/past_helper.py @@ -0,0 +1,50 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import logging + +logger = logging.getLogger(__name__) + + +class PastKeyValuesHelper: + """ Helper functions to process past key values for encoder-decoder model""" + + @staticmethod + def get_past_names(num_layers, present: bool = False): + past_self_names = [] + past_cross_names = [] + for i in range(num_layers): + past_self_names.extend([f'present_key_self_{i}', f'present_value_self_{i}'] + if present else [f'past_key_self_{i}', f'past_value_self_{i}']) + past_cross_names.extend([f'present_key_cross_{i}', f'present_value_cross_{i}'] + if present else [f'past_key_cross_{i}', f'past_value_cross_{i}']) + return past_self_names + past_cross_names + + @staticmethod + def group_by_self_or_cross(present_key_values): + """Split present state from grouped by layer to grouped by self/cross attention. + Before: (past_key_self_0, past_value_self_0, past_key_cross_0, past_value_cross_0), (past_key_self_1, past_value_self_1, past_key_cross_1, past_value_cross_1), ... + After: (past_key_self_0, past_value_self_0, past_key_self_1, past_value_self_1, ...), (past_key_cross_0, past_value_cross_0, past_key_cross_1, past_value_cross_1, ...) + + """ + present_self = [] + present_cross = [] + for i, present_layer_i in enumerate(present_key_values): + assert len(present_layer_i) == 4, f"Expected to have four items. Got {len(present_layer_i)}" + present_key_self, present_value_self, present_key_cross, present_value_cross = present_layer_i + present_self.extend([present_key_self, present_value_self]) + present_cross.extend([present_key_cross, present_value_cross]) + return present_self, present_cross + + @staticmethod + def group_by_layer(past, num_layers): + """Reorder past state from grouped by self/cross attention to grouped by layer. + Before: past_key_self_0, past_value_self_0, past_key_self_1, past_value_self_1, ..., past_key_cross_0, past_value_cross_0, past_key_cross_1, past_value_cross_1, ... + After: (past_key_self_0, past_value_self_0, past_key_cross_0, past_value_cross_0), (past_key_self_1, past_value_self_1, past_key_cross_1, past_value_cross_1), + """ + assert len(past) == 4 * num_layers + return tuple([past[2 * i], past[2 * i + 1], past[2 * num_layers + 2 * i], past[2 * num_layers + 2 * i + 1]] + for i in range(num_layers)) diff --git a/onnxruntime/python/tools/transformers/models/t5/t5_decoder.py b/onnxruntime/python/tools/transformers/models/t5/t5_decoder.py new file mode 100644 index 0000000000..5bfd530581 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/t5/t5_decoder.py @@ -0,0 +1,335 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from pathlib import Path +from typing import List, Union +import logging +import numpy +import torch +from transformers import T5Config +from onnxruntime import InferenceSession +from t5_encoder import T5EncoderInputs +from past_helper import PastKeyValuesHelper + +logger = logging.getLogger(__name__) + + +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, + config: T5Config, + decoder_start_token_id: int = None): + super().__init__() + self.decoder = decoder + self.lm_head = lm_head + self.config = config + self.decoder_start_token_id = decoder_start_token_id if decoder_start_token_id is not None else self.config.decoder_start_token_id + + def forward(self, decoder_input_ids: torch.Tensor, encoder_attention_mask: torch.Tensor, + encoder_hidden_states: torch.FloatTensor): + if decoder_input_ids is None: + batch_size = encoder_attention_mask.shape[0] + decoder_input_ids = torch.ones( + (batch_size, 1), dtype=torch.long, device=encoder_attention_mask.device) * self.decoder_start_token_id + + decoder_outputs = self.decoder(input_ids=decoder_input_ids, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=True, + return_dict=True) + + sequence_output = decoder_outputs.last_hidden_state + present_key_values = decoder_outputs.past_key_values + + sequence_output = sequence_output * (self.config.d_model**-0.5) + + lm_logits = self.lm_head(sequence_output) + past_self, past_cross = PastKeyValuesHelper.group_by_self_or_cross(present_key_values) + return lm_logits, past_self, past_cross + + +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 + self.lm_head = lm_head + self.config = config + + def forward(self, decoder_input_ids, encoder_attention_mask, encoder_hidden_states, *past): + + past_key_values = PastKeyValuesHelper.group_by_layer(past, self.config.num_layers) + + decoder_outputs = self.decoder(input_ids=decoder_input_ids, + past_key_values=past_key_values, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=True, + return_dict=True) + + sequence_output = decoder_outputs.last_hidden_state + present_key_values = decoder_outputs.past_key_values + + sequence_output = sequence_output * (self.config.d_model**-0.5) + + lm_logits = self.lm_head(sequence_output) + present_self, _ = PastKeyValuesHelper.group_by_self_or_cross(present_key_values) + + # Do not return present_cross since they are identical to corresponding past_cross input + return lm_logits, present_self + + +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 + self.encoder_hidden_states: Union[torch.FloatTensor, torch.HalfTensor] = encoder_hidden_states + 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: + """ Create dummy inputs for T5Decoder. + + Args: + decoder: decoder + batch_size (int): batch size + 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 + + Returns: + T5DecoderInputs: dummy inputs for decoder + """ + hidden_size: int = config.d_model + num_attention_heads: int = config.num_heads + num_layers: int = config.num_layers + vocab_size: int = config.vocab_size + + sequence_length: int = 1 # fixed for decoding + decoder_input_ids = torch.randint(low=0, + high=vocab_size - 1, + size=(batch_size, sequence_length), + dtype=torch.int64, + device=device) + + encoder_inputs = T5EncoderInputs.create_dummy(batch_size, encode_sequence_length, vocab_size, device) + + encoder_hidden_state = torch.rand(batch_size, + encode_sequence_length, + hidden_size, + dtype=torch.float32, + device=device) + + if past_decode_sequence_length > 0: + self_attention_past_shape = [ + batch_size, num_attention_heads, past_decode_sequence_length, + int(hidden_size / num_attention_heads) + ] + cross_attention_past_shape = [ + batch_size, num_attention_heads, encode_sequence_length, + int(hidden_size / num_attention_heads) + ] + + past = [] + for _ in range(2 * num_layers): + past.append(torch.rand(self_attention_past_shape, dtype=torch.float32, device=device)) + + for _ in range(2 * num_layers): + past.append(torch.rand(cross_attention_past_shape, dtype=torch.float32, device=device)) + else: + past = None + + return T5DecoderInputs(decoder_input_ids, encoder_inputs.attention_mask, encoder_hidden_state, past) + + def to_list(self) -> List: + input_list = [self.decoder_input_ids, self.encoder_attention_mask, self.encoder_hidden_states] + if self.past_key_values: + input_list.extend(self.past_key_values) + return input_list + + +class T5DecoderHelper: + + @staticmethod + def export_onnx(decoder: Union[T5Decoder, T5DecoderInit], + device: torch.device, + onnx_model_path: str, + verbose: bool = True, + use_external_data_format: bool = False): + """Export decoder to ONNX + + Args: + decoder (Union[T5Decoder, T5DecoderNoPastState]): decoder object + device (torch.device): device of decoder object + onnx_model_path (str): onnx path + verbose (bool, optional): print verbose information. Defaults to True. + use_external_data_format (bool, optional): use external data format or not. Defaults to False. + """ + assert isinstance(decoder, (T5Decoder, T5DecoderInit)) + + inputs = T5DecoderInputs.create_dummy(decoder.config, + batch_size=2, + encode_sequence_length=3, + past_decode_sequence_length=5 if isinstance(decoder, T5Decoder) else 0, + device=device) + input_list = inputs.to_list() + with torch.no_grad(): + outputs = decoder(*input_list) + + past_names = PastKeyValuesHelper.get_past_names(decoder.config.num_layers, present=False) + present_names = PastKeyValuesHelper.get_past_names(decoder.config.num_layers, present=True) + present_self_names = present_names[:2 * decoder.config.num_layers] + + input_past_names = past_names if isinstance(decoder, T5Decoder) else [] + output_present_names = present_self_names if isinstance(decoder, T5Decoder) else present_names + output_names = ["logits"] + output_present_names + + # Shape of input tensors (sequence_length==1): + # input_ids: (batch_size, sequence_length) + # encoder_attention_mask: (batch_size, encode_sequence_length) + # encoder_hidden_states: (batch_size, encode_sequence_length, hidden_size) + # past_self_*: (batch_size, num_heads, past_decode_sequence_length, hidden_size/num_heads) + # past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads) + + # Shape of output tensors: + # logits: (batch_size, sequence_length, vocab_size) + # past_self_*: (batch_size, num_heads, past_decode_sequence_length + sequence_length, hidden_size/num_heads) + # past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads) + + input_names = ["input_ids"] + input_names.append("encoder_attention_mask") + input_names.append("encoder_hidden_states") + input_names.extend(input_past_names) + + dynamic_axes = { + 'input_ids': { + 0: 'batch_size', + #1: 'sequence_length' + }, + 'encoder_attention_mask': { + 0: 'batch_size', + 1: 'encode_sequence_length' + }, + 'encoder_hidden_states': { + 0: 'batch_size', + 1: 'encode_sequence_length' + }, + "logits": { + 0: 'batch_size', + #1: 'sequence_length' + } + } + + for name in input_past_names: + dynamic_axes[name] = { + 0: 'batch_size', + 2: 'past_decode_sequence_length' if "self" in name else "encode_sequence_length" + } + + for name in output_present_names: + if "cross" in name: + dynamic_axes[name] = {0: 'batch_size', 2: "encode_sequence_length"} + else: # self attention past state + if isinstance(decoder, T5Decoder): + dynamic_axes[name] = {0: 'batch_size', 2: 'past_decode_sequence_length + 1'} + else: + dynamic_axes[name] = { + 0: 'batch_size', + #2: 'sequence_length' + } + + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + torch.onnx.export(decoder, + args=tuple(input_list), + f=onnx_model_path, + export_params=True, + input_names=input_names, + output_names=output_names, + example_outputs=outputs, + dynamic_axes=dynamic_axes, + opset_version=12, + do_constant_folding=True, + use_external_data_format=use_external_data_format, + verbose=verbose) + + @staticmethod + def onnxruntime_inference(ort_session, inputs: T5DecoderInputs): + """ Run inference of ONNX model. + """ + logger.debug(f"start onnxruntime_inference") + + ort_inputs = { + 'input_ids': numpy.ascontiguousarray(inputs.decoder_input_ids.cpu().numpy()), + 'encoder_attention_mask': numpy.ascontiguousarray(inputs.encoder_attention_mask.cpu().numpy()), + 'encoder_hidden_states': numpy.ascontiguousarray(inputs.encoder_hidden_states.cpu().numpy()) + } + + if inputs.past_key_values: + assert len(inputs.past_key_values) % 4 == 0 + num_layers = int(len(inputs.past_key_values) / 4) + past_names = PastKeyValuesHelper.get_past_names(num_layers) + for i, past_tensor in enumerate(inputs.past_key_values): + ort_inputs[past_names[i]] = numpy.ascontiguousarray(past_tensor.cpu().numpy()) + + ort_outputs = ort_session.run(None, ort_inputs) + return ort_outputs + + @staticmethod + def verify_onnx(model: Union[T5Decoder, T5DecoderInit], + ort_session: InferenceSession, + device: torch.device, + max_cases=4): + """ Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good. + """ + 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]: + if isinstance(model, T5DecoderInit): + past_decode_sequence_length = 0 + + inputs = T5DecoderInputs.create_dummy(model.config, + batch_size, + encode_sequence_length, + past_decode_sequence_length, + device=device) + input_list = inputs.to_list() + + # Run inference of PyTorch model + with torch.no_grad(): + torch_outputs = model(*input_list) + + ort_outputs = T5DecoderHelper.onnxruntime_inference(ort_session, inputs) + + max_diff = numpy.amax(numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs[0])) + max_diff_all = max_diff + logger.debug(f"logits max_diff={max_diff}") + + for i in range(2 * model.config.num_layers): + max_diff = numpy.amax(numpy.abs(torch_outputs[1][i].cpu().numpy() - ort_outputs[1 + i])) + logger.debug(f"self attention past state {i} max_diff={max_diff}") + max_diff_all = max(max_diff_all, max_diff) + + if isinstance(model, T5DecoderInit): + for i in range(2 * model.config.num_layers): + max_diff = numpy.amax( + numpy.abs(torch_outputs[2][i].cpu().numpy() - ort_outputs[1 + 2 * model.config.num_layers + i])) + logger.debug(f"cross attention past state {i} max_diff={max_diff}") + max_diff_all = max(max_diff_all, max_diff) + + test_cases_max_diff.append(max_diff_all) + logger.info( + f"batch_size={batch_size} encode_sequence_length={encode_sequence_length}, past_decode_sequence_length={past_decode_sequence_length}, max_diff={max_diff_all}" + ) + + return max_diff_all diff --git a/onnxruntime/python/tools/transformers/models/t5/t5_encoder.py b/onnxruntime/python/tools/transformers/models/t5/t5_encoder.py new file mode 100644 index 0000000000..c0086896b7 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/t5/t5_encoder.py @@ -0,0 +1,150 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import random +from pathlib import Path +from typing import List +import logging +import numpy +import torch +from transformers import T5Config +from onnxruntime import InferenceSession + +logger = logging.getLogger(__name__) + + +class T5Encoder(torch.nn.Module): + """ T5 encoder outputs only the last hidden state""" + + def __init__(self, encoder, config: T5Config): + super().__init__() + self.encoder = encoder + self.config = config + + def forward(self, input_ids, attention_mask): + return self.encoder(input_ids, attention_mask)[0] + + +class T5EncoderInputs: + + def __init__(self, input_ids, attention_mask): + self.input_ids: torch.LongTensor = input_ids + self.attention_mask: torch.LongTensor = attention_mask + + @staticmethod + def create_dummy(batch_size: int, sequence_length: int, vocab_size: int, + device: torch.device): # -> T5EncoderInputs + """ Create dummy inputs for T5 encoder. + + Args: + batch_size (int): batch size + sequence_length (int): sequence length + vocab_size (int): vocaburary size + device (torch.device): device of output tensors + + Returns: + T5EncoderInputs: dummy inputs for encoder + """ + input_ids = torch.randint(low=0, + high=vocab_size - 1, + size=(batch_size, sequence_length), + dtype=torch.int64, + device=device) + + attention_mask = torch.ones([batch_size, sequence_length], dtype=torch.int64, device=device) + if sequence_length >= 2: + for i in range(batch_size): + padding_position = random.randint(0, sequence_length - 1) + attention_mask[i, :padding_position] = 0 + return T5EncoderInputs(input_ids, attention_mask) + + def to_list(self) -> List: + input_list = [v for v in [self.input_ids, self.attention_mask] if v is not None] + return input_list + + +class T5EncoderHelper: + + @staticmethod + def export_onnx(encoder: T5Encoder, + device: torch.device, + onnx_model_path: str, + verbose: bool = True, + use_external_data_format: bool = False): + """Export encoder to ONNX + + Args: + encoder (T5Encoder): encoder object + device (torch.device): device of encoder object + onnx_model_path (str): onnx path + verbose (bool, optional): print verbose information. Defaults to True. + use_external_data_format (bool, optional): use external data format or not. Defaults to False. + """ + config = encoder.config + encoder_inputs = T5EncoderInputs.create_dummy(batch_size=2, + sequence_length=4, + vocab_size=config.vocab_size, + device=device) + + with torch.no_grad(): + outputs = encoder(encoder_inputs.input_ids, encoder_inputs.attention_mask) + + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + torch.onnx.export(encoder, + args=tuple(encoder_inputs.to_list()), + f=onnx_model_path, + export_params=True, + input_names=['input_ids', 'attention_mask'], + output_names=['hidden_states'], + example_outputs=outputs, + dynamic_axes={ + 'input_ids': { + 0: 'batch_size', + 1: 'sequence_length' + }, + 'attention_mask': { + 0: 'batch_size', + 1: 'sequence_length' + }, + 'hidden_states': { + 0: 'batch_size', + 1: 'sequence_length' + }, + }, + opset_version=12, + do_constant_folding=True, + use_external_data_format=use_external_data_format, + verbose=verbose) + + @staticmethod + def onnxruntime_inference(ort_session, inputs: T5EncoderInputs): + """ Run inference of ONNX model. + """ + ort_inputs = { + 'input_ids': numpy.ascontiguousarray(inputs.input_ids.cpu().numpy()), + 'attention_mask': numpy.ascontiguousarray(inputs.attention_mask.cpu().numpy()) + } + + return ort_session.run(None, ort_inputs) + + @staticmethod + def verify_onnx(model: T5Encoder, ort_session: InferenceSession, device: torch.device): + """ Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good. + """ + inputs = T5EncoderInputs.create_dummy(batch_size=4, + sequence_length=11, + vocab_size=model.config.vocab_size, + device=device) + input_list = inputs.to_list() + torch_outputs = model(*input_list) + + ort_outputs = T5EncoderHelper.onnxruntime_inference(ort_session, inputs) + + max_diff = numpy.amax(numpy.abs(torch_outputs.cpu().numpy() - ort_outputs[0])) + + logger.info(f'max_diff={max_diff}') + + return max_diff diff --git a/onnxruntime/python/tools/transformers/models/t5/t5_encoder_decoder_init.py b/onnxruntime/python/tools/transformers/models/t5/t5_encoder_decoder_init.py new file mode 100644 index 0000000000..29b82cda19 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/t5/t5_encoder_decoder_init.py @@ -0,0 +1,231 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +from pathlib import Path +from typing import List +import logging +import numpy +import torch +from transformers import T5Config +from onnxruntime import InferenceSession +from t5_encoder import T5Encoder, T5EncoderInputs +from t5_decoder import T5DecoderInit +from past_helper import PastKeyValuesHelper + +logger = logging.getLogger(__name__) + + +class T5EncoderDecoderInit(torch.nn.Module): + """ A combination of T5Encoder and T5DecoderInit. + """ + + def __init__(self, + encoder: torch.nn.Module, + decoder: torch.nn.Module, + lm_head: torch.nn.Module, + config: T5Config, + decoder_start_token_id: int = None): + super().__init__() + self.config = config + self.t5_encoder = T5Encoder(encoder, config) + self.t5_decoder_init = T5DecoderInit(decoder, lm_head, config, decoder_start_token_id) + + def forward(self, + encoder_input_ids: torch.Tensor, + encoder_attention_mask: torch.Tensor, + decoder_input_ids: torch.Tensor = None): + encoder_hidden_states: torch.FloatTensor = self.t5_encoder(encoder_input_ids, encoder_attention_mask) + lm_logits, past_self, past_cross = self.t5_decoder_init(decoder_input_ids, encoder_attention_mask, + encoder_hidden_states) + return lm_logits, encoder_hidden_states, past_self, past_cross + + +class T5EncoderDecoderInitInputs: + + def __init__(self, encoder_input_ids, encoder_attention_mask, decoder_input_ids=None): + self.encoder_input_ids: torch.LongTensor = encoder_input_ids + self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask + self.decoder_input_ids: torch.LongTensor = decoder_input_ids + + @staticmethod + def create_dummy(config: T5Config, batch_size: int, encode_sequence_length: int, use_decoder_input_ids: int, + device: torch.device): # -> T5EncoderDecoderInitInputs: + encoder_inputs: T5EncoderInputs = T5EncoderInputs.create_dummy(batch_size, encode_sequence_length, + config.vocab_size, device) + decoder_input_ids = None + if use_decoder_input_ids: + decoder_input_ids = torch.ones( + (batch_size, 1), dtype=torch.long, device=device) * config.decoder_start_token_id + + return T5EncoderDecoderInitInputs(encoder_inputs.input_ids, encoder_inputs.attention_mask, decoder_input_ids) + + def to_list(self) -> List: + input_list = [self.encoder_input_ids, self.encoder_attention_mask] + if self.decoder_input_ids is not None: + input_list.append(self.decoder_input_ids) + return input_list + + +class T5EncoderDecoderInitHelper: + + @staticmethod + def export_onnx(model: T5EncoderDecoderInit, + device: torch.device, + onnx_model_path: str, + use_decoder_input_ids: bool = True, + verbose: bool = True, + use_external_data_format: bool = False): + """Export decoder to ONNX + + Args: + model (T5EncoderDecoderInit): the model to export + device (torch.device): device of decoder object + onnx_model_path (str): onnx path + verbose (bool, optional): print verbose information. Defaults to True. + use_external_data_format (bool, optional): use external data format or not. Defaults to False. + """ + assert isinstance(model, T5EncoderDecoderInit) + + inputs = T5EncoderDecoderInitInputs.create_dummy(model.config, + batch_size=2, + encode_sequence_length=3, + use_decoder_input_ids=use_decoder_input_ids, + device=device) + input_list = inputs.to_list() + outputs = model(*input_list) + + present_names = PastKeyValuesHelper.get_past_names(model.config.num_layers, present=True) + + output_names = ["logits", "encoder_hidden_states"] + present_names + + # Shape of input tensors (sequence_length==1): + # input_ids: (batch_size, sequence_length) + # encoder_attention_mask: (batch_size, encode_sequence_length) + # encoder_hidden_states: (batch_size, encode_sequence_length, hidden_size) + # past_self_*: (batch_size, num_heads, past_decode_sequence_length, hidden_size/num_heads) + # past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads) + + # Shape of output tensors: + # logits: (batch_size, sequence_length, vocab_size) + # past_self_*: (batch_size, num_heads, past_decode_sequence_length + sequence_length, hidden_size/num_heads) + # past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads) + + input_names = ["encoder_input_ids", "encoder_attention_mask"] + + # ONNX exporter might mark dimension like 'Transposepresent_value_self_1_dim_2'. Use more friendly string here. + sequence_length = '1' + num_heads = str(model.config.num_heads) + hidden_size = str(model.config.d_model) + head_size = str(model.config.d_model // model.config.num_heads) + + dynamic_axes = { + 'encoder_input_ids': { + 0: 'batch_size', + 1: 'encode_sequence_length' + }, + 'encoder_attention_mask': { + 0: 'batch_size', + 1: 'encode_sequence_length' + }, + 'encoder_hidden_states': { + 0: 'batch_size', + 1: 'encode_sequence_length', + 2: hidden_size + }, + "logits": { + 0: 'batch_size', + 1: sequence_length + } + } + + if use_decoder_input_ids: + input_names.append("decoder_input_ids") + dynamic_axes["decoder_input_ids"] = {0: 'batch_size', 1: sequence_length} + + for name in present_names: + if "cross" in name: + dynamic_axes[name] = {0: 'batch_size', 1: num_heads, 2: 'encode_sequence_length', 3: head_size} + + else: # self attention past state + dynamic_axes[name] = {0: 'batch_size', 1: num_heads, 2: sequence_length, 3: head_size} + + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + torch.onnx.export(model, + args=tuple(input_list), + f=onnx_model_path, + export_params=True, + input_names=input_names, + output_names=output_names, + example_outputs=outputs, + dynamic_axes=dynamic_axes, + opset_version=12, + do_constant_folding=True, + use_external_data_format=use_external_data_format, + verbose=verbose) + + @staticmethod + def onnxruntime_inference(ort_session, inputs: T5EncoderDecoderInitInputs): + """ Run inference of ONNX model. + """ + logger.debug(f"start onnxruntime_inference") + + ort_inputs = { + 'encoder_input_ids': numpy.ascontiguousarray(inputs.encoder_input_ids.cpu().numpy()), + 'encoder_attention_mask': numpy.ascontiguousarray(inputs.encoder_attention_mask.cpu().numpy()), + } + if inputs.decoder_input_ids is not None: + ort_inputs['decoder_input_ids'] = numpy.ascontiguousarray(inputs.decoder_input_ids.cpu().numpy()) + + ort_outputs = ort_session.run(None, ort_inputs) + return ort_outputs + + @staticmethod + def verify_onnx(model: T5EncoderDecoderInit, ort_session: InferenceSession, device: torch.device, max_cases=4): + """ Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good. + """ + ort_inputs = ort_session.get_inputs() + use_decoder_input_ids = len(ort_inputs) == 3 + + test_cases = [(4, 11), (1, 2), (3, 1), (8, 5)] + test_cases_max_diff = [] + for (batch_size, encode_sequence_length) in test_cases[:max_cases]: + inputs = T5EncoderDecoderInitInputs.create_dummy(model.config, + batch_size, + encode_sequence_length, + use_decoder_input_ids=use_decoder_input_ids, + device=device) + + ort_outputs = T5EncoderDecoderInitHelper.onnxruntime_inference(ort_session, inputs) + + # Run inference of PyTorch model + input_list = inputs.to_list() + torch_outputs = model(*input_list) + + assert (torch_outputs[0].cpu().numpy().shape == ort_outputs[0].shape) + max_diff = numpy.amax(numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs[0])) + logger.debug(f"logits max_diff={max_diff}") + max_diff_all = max_diff + + assert (torch_outputs[1].cpu().numpy().shape == ort_outputs[1].shape) + max_diff = numpy.amax(numpy.abs(torch_outputs[1].cpu().numpy() - ort_outputs[1])) + logger.debug(f"encoder_hidden_states max_diff={max_diff}") + max_diff_all = max(max_diff_all, max_diff) + + for i in range(2 * model.config.num_layers): + max_diff = numpy.amax(numpy.abs(torch_outputs[2][i].cpu().numpy() - ort_outputs[2 + i])) + logger.debug(f"self attention past state {i} max_diff={max_diff}") + + for i in range(2 * model.config.num_layers): + max_diff = numpy.amax( + numpy.abs(torch_outputs[3][i].cpu().numpy() - ort_outputs[2 + 2 * model.config.num_layers + i])) + logger.debug(f"cross attention past state {i} max_diff={max_diff}") + max_diff_all = max(max_diff_all, max_diff) + + test_cases_max_diff.append(max_diff_all) + logger.info( + f"batch_size={batch_size} encode_sequence_length={encode_sequence_length}, max_diff={max_diff_all}") + + return max(test_cases_max_diff) diff --git a/onnxruntime/python/tools/transformers/models/t5/t5_helper.py b/onnxruntime/python/tools/transformers/models/t5/t5_helper.py new file mode 100644 index 0000000000..f04fa9941c --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/t5/t5_helper.py @@ -0,0 +1,133 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import os +import sys +from pathlib import Path +from typing import Union, Dict +import logging +import torch +from transformers import T5ForConditionalGeneration +from onnxruntime import InferenceSession +from t5_encoder import T5Encoder, T5EncoderHelper +from t5_decoder import T5DecoderInit, T5Decoder, T5DecoderHelper +from t5_encoder_decoder_init import T5EncoderDecoderInit, T5EncoderDecoderInitHelper + +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 + + Args: + output_dir (str): output directory + model_name_or_path (str): pretrained model name, or path to the model checkpoint + suffix (str, optional): suffix like "_encoder" or "_decoder_fp16" will be appended to file name. Defaults to None. + new_folder (bool, optional): create a new directory for the model. Defaults to False. + + Returns: + str: path of onnx model + """ + model_name = model_name_or_path + if os.path.isdir(model_name_or_path): + model_name = Path(model_name_or_path).parts[-1] + else: + model_name.split('/')[-1] + + model_name += suffix + + dir = os.path.join(output_dir, model_name) if new_folder else output_dir + return os.path.join(dir, model_name + ".onnx") + + @staticmethod + def load_model(model_name_or_path: str, + cache_dir: str, + device: torch.device, + merge_encoder_and_decoder_init: bool = True) -> Dict[str, torch.nn.Module]: + """Load model given a pretrained name or path, then build models for ONNX conversion. + + Args: + model_name_or_path (str): pretrained model name or path + cache_dir (str): cache directory + device (torch.device): device to run the model + merge_encoder_and_decoder_init (bool, optional): Whether merge encoder and decoder initialization into one ONNX model. Defaults to True. + + Returns: + Dict[str, torch.nn.Module]: mapping from name to modules for ONNX conversion. + """ + model = T5ForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir) + + decoder = T5Decoder(model.decoder, model.lm_head, model.config) + decoder.eval().to(device) + + if merge_encoder_and_decoder_init: + encoder_decoder_init = T5EncoderDecoderInit(model.encoder, + model.decoder, + model.lm_head, + model.config, + decoder_start_token_id=None) + return {"encoder_decoder_init": encoder_decoder_init, "decoder": decoder} + else: + encoder = T5Encoder(model.encoder, model.config) + encoder.eval().to(device) + decoder_init = T5DecoderInit(model.decoder, model.lm_head, model.config) + decoder_init.eval().to(device) + return {"encoder": encoder, "decoder": decoder, "decoder_init": decoder_init} + + @staticmethod + def export_onnx(model: Union[T5Encoder, T5Decoder, T5DecoderInit, T5EncoderDecoderInit], + device: torch.device, + onnx_model_path: str, + verbose: bool = True, + use_external_data_format: bool = False, + use_decoder_input_ids: bool = True): + if isinstance(model, T5Encoder): + T5EncoderHelper.export_onnx(model, device, onnx_model_path, verbose, use_external_data_format) + elif isinstance(model, T5EncoderDecoderInit): + T5EncoderDecoderInitHelper.export_onnx(model, device, onnx_model_path, use_decoder_input_ids, verbose, + use_external_data_format) + else: + T5DecoderHelper.export_onnx(model, device, onnx_model_path, verbose, use_external_data_format) + + @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): + """ 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, + hidden_size=hidden_size, + opt_level=0, + optimization_options=None, + use_gpu=False) + if is_float16: + m.convert_model_float32_to_float16(cast_input_output=False) + + m.save_model_to_file(optimized_model_path, use_external_data_format) + + @staticmethod + def verify_onnx(model: Union[T5Encoder, T5Decoder, T5DecoderInit, T5EncoderDecoderInit], + ort_session: InferenceSession, device: torch.device): + """ Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good. + """ + if isinstance(model, T5Encoder): + return T5EncoderHelper.verify_onnx(model, ort_session, device) + elif isinstance(model, T5EncoderDecoderInit): + return T5EncoderDecoderInitHelper.verify_onnx(model, ort_session, device) + else: + return T5DecoderHelper.verify_onnx(model, ort_session, device)