# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # ------------------------------------------------------------------------- import logging import os import tempfile from pathlib import Path import numpy import onnx import torch from onnx_model import OnnxModel from past_helper import PastKeyValuesHelper from t5_decoder import T5DecoderInit from t5_encoder import T5Encoder, T5EncoderInputs from torch_onnx_export_helper import torch_onnx_export from transformers import MT5Config, T5Config from onnxruntime import InferenceSession 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.Linear, config: T5Config | MT5Config, decoder_start_token_id: int | None = None, output_cross_only: bool = False, ): super().__init__() self.config: T5Config | MT5Config = config self.t5_encoder = T5Encoder(encoder, config) self.t5_decoder_init = T5DecoderInit(decoder, lm_head, config, decoder_start_token_id) self.output_cross_only = output_cross_only def forward( self, encoder_input_ids: torch.Tensor, encoder_attention_mask: torch.Tensor, decoder_input_ids: torch.Tensor | None = 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 ) if self.output_cross_only: return past_cross else: 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 | None = decoder_input_ids @staticmethod def create_dummy( config: T5Config | MT5Config, batch_size: int, encode_sequence_length: int, use_decoder_input_ids: int, device: torch.device, use_int32_inputs: bool = False, ): # -> T5EncoderDecoderInitInputs: encoder_inputs: T5EncoderInputs = T5EncoderInputs.create_dummy( batch_size, encode_sequence_length, config.vocab_size, device, use_int32_inputs=use_int32_inputs, ) decoder_input_ids = None if use_decoder_input_ids: dtype = torch.int32 if use_int32_inputs else torch.int64 decoder_input_ids = torch.ones((batch_size, 1), dtype=dtype, 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, use_int32_inputs: 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. use_int32_inputs (bool, optional): use int32 instead of int64 for integer inputs. Defaults to False. """ assert isinstance(model, T5EncoderDecoderInit) # Do not exclude decoder in torch onnx export so that cross can show up. output_cross_only = model.output_cross_only model.output_cross_only = False inputs = T5EncoderDecoderInitInputs.create_dummy( model.config, batch_size=2, encode_sequence_length=3, use_decoder_input_ids=use_decoder_input_ids, device=device, use_int32_inputs=use_int32_inputs, ) input_list = inputs.to_list() present_names = PastKeyValuesHelper.get_past_names(model.config.num_decoder_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, head_size) # past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size) # Shape of output tensors: # logits: (batch_size, sequence_length, vocab_size) # past_self_*: (batch_size, num_heads, past_decode_sequence_length + sequence_length, head_size) # past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size) input_names = ["encoder_input_ids", "encoder_attention_mask"] # ONNX exporter might mark dimension like 'present_value_self_1_dim_2' in shape inference. # We use a workaround here: first use dim_param "1" for sequence_length, and later change to dim_value. sequence_length = "1" num_heads = str(model.config.num_heads) hidden_size = str(model.config.d_model) head_size = str(model.config.d_kv) 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, } with tempfile.TemporaryDirectory() as tmp_dir_name: temp_onnx_model_path = os.path.join(tmp_dir_name, "encoder_decoder_init.onnx") Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True) torch_onnx_export( model, args=tuple(input_list), f=temp_onnx_model_path, export_params=True, input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, opset_version=12, do_constant_folding=True, use_external_data_format=use_external_data_format, verbose=verbose, ) # Restore output_cross_only setting. model.output_cross_only = output_cross_only # Workaround as mentioned earlier: change numeric dim_param to dim_value exported_model: onnx.ModelProto = onnx.load(temp_onnx_model_path) for tensor in exported_model.graph.output: for dim_proto in tensor.type.tensor_type.shape.dim: if dim_proto.HasField("dim_param") and dim_proto.dim_param in [ sequence_length, num_heads, hidden_size, head_size, ]: dim_value = int(dim_proto.dim_param) dim_proto.Clear() dim_proto.dim_value = dim_value if output_cross_only: # Rewrite onnx graph to only keep present_[key|value]_cross_* outputs. onnx_model = OnnxModel(exported_model) output_name_to_node = onnx_model.output_name_to_node() for output in exported_model.graph.output: if "cross" in output.name: assert output.name in output_name_to_node transpose_node = output_name_to_node[output.name] assert transpose_node and transpose_node.op_type == "Transpose" permutation = OnnxModel.get_node_attribute(transpose_node, "perm") assert isinstance(permutation, list) assert permutation == [0, 2, 1, 3] matched_nodes = onnx_model.match_parent_path( transpose_node, ["Reshape", "MatMul"], [0, 0], output_name_to_node, ) assert matched_nodes is not None reshape_node, matmul_node = matched_nodes assert "encoder_hidden_states" in matmul_node.input if not onnx_model.get_initializer("cross_reshape_shape"): shape_tensor = onnx.helper.make_tensor( name="cross_reshape_shape", data_type=onnx.TensorProto.INT64, dims=[4], vals=[0, 0, int(num_heads), int(head_size)], raw=False, ) onnx_model.add_initializer(shape_tensor) reshape_node.input[1] = "cross_reshape_shape" cross_outputs = [output.name for output in exported_model.graph.output if "cross" in output.name] onnx_model.prune_graph(cross_outputs, allow_remove_graph_inputs=True) OnnxModel.save( exported_model, onnx_model_path, save_as_external_data=use_external_data_format, all_tensors_to_one_file=True, ) @staticmethod def onnxruntime_inference(ort_session, inputs: T5EncoderDecoderInitInputs): """Run inference of ONNX model.""" logger.debug("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, use_int32_inputs: bool, max_cases: int = 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, use_int32_inputs=use_int32_inputs, ) ort_outputs = T5EncoderDecoderInitHelper.onnxruntime_inference(ort_session, inputs) # Run inference of PyTorch model input_list = inputs.to_list() torch_outputs = model(*input_list) num_decoder_layers = model.config.num_decoder_layers if not model.output_cross_only: 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 * num_decoder_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 * num_decoder_layers): max_diff = numpy.amax( numpy.abs(torch_outputs[3][i].cpu().numpy() - ort_outputs[2 + 2 * num_decoder_layers + i]) ) logger.debug(f"cross attention past state {i} max_diff={max_diff}") max_diff_all = max(max_diff_all, max_diff) else: max_diff_all = -float("inf") for i in range(2 * num_decoder_layers): max_diff = numpy.amax(numpy.abs(torch_outputs[i].cpu().numpy() - ort_outputs[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)