diff --git a/onnxruntime/python/tools/transformers/models/t5/past_helper.py b/onnxruntime/python/tools/transformers/models/t5/past_helper.py index 60fe27174c..915b09da79 100644 --- a/onnxruntime/python/tools/transformers/models/t5/past_helper.py +++ b/onnxruntime/python/tools/transformers/models/t5/past_helper.py @@ -5,6 +5,9 @@ # -------------------------------------------------------------------------- import logging +from typing import List, Tuple + +import torch logger = logging.getLogger(__name__) @@ -66,3 +69,82 @@ class PastKeyValuesHelper: ] for i in range(num_layers) ) + + @staticmethod + def back_group_by_layer(past_key_values: Tuple[Tuple[torch.Tensor]]): + """Categorize present_key_values from self and cross attention to layer by layer. + + 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), + + Args: + present_key_values: From past_key_values of a model (group by self and cross attention) + + Returns: + past_tuples: present key and values grouped by layer. + """ + past_tuples = () + half_idx = len(past_key_values) // 2 + for i in range(len(past_key_values) // 4): + idx = 2 * i + past_tuples += ( + ( + past_key_values[idx], + past_key_values[idx + 1], + past_key_values[half_idx + idx], + past_key_values[half_idx + idx + 1], + ), + ) + return past_tuples + + @staticmethod + def group_by_self_and_cross(present_key_values: Tuple[torch.Tensor], concat: bool = False): + """Categorize present_key_values into self and cross attention. + + 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, ...) + + Args: + present_key_values: From past_key_values of a model (group by layer) + concat: If concat self attention with cross attention key/value to return + + Returns: + present_self (Tuple[torch.Tensor]): present key and values from self attention + present_cross (Tuple[torch.Tensor]): present key and values from cross attention + """ + present_self: List[torch.Tensor] = [] + present_cross: List[torch.Tensor] = [] + for _, 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]) + if concat: + return present_self + present_cross + else: + return present_self, present_cross + + @staticmethod + def get_input_names(past_key_values: Tuple[Tuple[torch.Tensor]], encoder=True): + """Process input names of model wrapper. + + Args: + past_key_values: Consider `self` and `cross` past_key_values + + Returns: + names (List[string]): input names + """ + names = [] + num_layers = len(past_key_values) // 4 if encoder else len(past_key_values) + prefix = "past_" if not encoder else "present_" + for i in range(num_layers): + names.extend([prefix + s for s in [f"key_self_{i}", f"value_self_{i}"]]) + for i in range(num_layers): + names.extend([prefix + s for s in [f"key_cross_{i}", f"value_cross_{i}"]]) + return names diff --git a/onnxruntime/python/tools/transformers/models/whisper/__init__.py b/onnxruntime/python/tools/transformers/models/whisper/__init__.py new file mode 100644 index 0000000000..cc667396a2 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/whisper/__init__.py @@ -0,0 +1,4 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- diff --git a/onnxruntime/python/tools/transformers/models/whisper/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/whisper/convert_to_onnx.py new file mode 100644 index 0000000000..ae1111f245 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/whisper/convert_to_onnx.py @@ -0,0 +1,294 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import argparse +import copy +import logging +import os +import sys + +import torch +from whisper_chain import chain_model +from whisper_helper import PRETRAINED_WHISPER_MODELS, WhisperHelper + +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) +from benchmark_helper import Precision, create_onnxruntime_session, prepare_environment, setup_logger # noqa: E402 + +logger = logging.getLogger("") + + +def parse_arguments(): + parser = argparse.ArgumentParser() + + pretrained_models = PRETRAINED_WHISPER_MODELS + parser.add_argument( + "-m", + "--model_name_or_path", + required=False, + default=PRETRAINED_WHISPER_MODELS[0], + type=str, + help="Model path, or pretrained model name in the list: " + ", ".join(pretrained_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. 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) + + 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) + + parser.add_argument( + "--separate_encoder_and_decoder_init", + required=False, + action="store_true", + help="Do not merge encode and decoder init. Output 3 instead of 2 onnx models.", + ) + parser.set_defaults(separate_encoder_and_decoder_init=False) + + parser.add_argument( + "--use_int64_inputs", + required=False, + action="store_true", + help="Use int64 instead of int32 for input_ids, position_ids and attention_mask.", + ) + parser.set_defaults(use_int64_inputs=False) + + parser.add_argument( + "--chain_model", + required=False, + action="store_true", + help="Produce beam search model with chained encdecinit and decoder.", + ) + parser.set_defaults(chain_model=True) + + parser.add_argument( + "--beam_output_model", + type=str, + default="whisper_beamsearch.onnx", + help="default name is whisper_beamsearch.onnx.", + ) + + parser.add_argument("--no_repeat_ngram_size", type=int, default=3, help="default to 3") + + 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 = False, + merge_encoder_and_decoder_init: bool = True, + overwrite: bool = False, + disable_auto_mixed_precision: bool = False, + use_int32_inputs: bool = True, +): + device = torch.device("cuda:0" if use_gpu else "cpu") + + models = WhisperHelper.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("Try use_external_data_format when model size > 2GB") + + output_paths = [] + for name, model in models.items(): + model.to(device) + filename_suffix = "_" + name + + onnx_path = WhisperHelper.get_onnx_path( + output_dir, + model_name_or_path, + suffix=filename_suffix, + new_folder=False, + ) + + 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. + cloned_model = copy.deepcopy(model).to(device) + WhisperHelper.export_onnx( + cloned_model, + device, + onnx_path, + verbose, + use_external_data_format, + use_decoder_input_ids=not use_decoder_start_token, + use_int32_inputs=use_int32_inputs, + ) + else: + logger.info(f"Skip exporting: existed ONNX model {onnx_path}") + + # Optimize ONNX graph. Note that we have not implemented graph optimization for Whisper yet. + if optimize_onnx or precision != Precision.FLOAT32: + output_path = WhisperHelper.get_onnx_path( + output_dir, + model_name_or_path, + suffix=filename_suffix + "_" + str(precision), + new_folder=False, + ) + + if overwrite or not os.path.exists(output_path): + logger.info(f"Optimizing model to {output_path}") + WhisperHelper.optimize_onnx( + onnx_path, + output_path, + precision == Precision.FLOAT16, + config.encoder_attention_heads, + config.d_model, + use_external_data_format, + auto_mixed_precision=not disable_auto_mixed_precision, + use_gpu=use_gpu, + ) + 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"], + ) + + with torch.no_grad(): + max_diff = WhisperHelper.verify_onnx(model, ort_session, device, use_int32_inputs) + logger.info(f"PyTorch and OnnxRuntime results max difference = {max_diff}") + if max_diff > 1e-4: + logger.warning("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.warning("Graph optimization for Whisper is not implemented yet.") + + 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, + not args.separate_encoder_and_decoder_init, + args.overwrite, + args.disable_auto_mixed_precision, + not args.use_int64_inputs, + ) + + if args.chain_model: + logger.info("Chaining model ... :") + args.beam_model_output_dir = WhisperHelper.get_onnx_path( + output_dir, + args.model_name_or_path, + suffix="_beamsearch", + new_folder=False, + ) + for path in output_paths: + if "encoder_decoder" in path: + args.encoder_path = path + elif "decoder" in path: + args.decoder_path = path + chain_model(args) + output_paths.append(args.beam_model_output_dir) + + logger.info(f"Done! Outputs: {output_paths}") + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py new file mode 100644 index 0000000000..b9e17d6ef1 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py @@ -0,0 +1,101 @@ +import onnx +from onnx import TensorProto, helper +from transformers import WhisperConfig + + +def add_attention_mask(model): + # Add attention mask - required by BeamSearch but unused in Pytorch + mask = helper.make_tensor_value_info( + "encoder_attention_mask", TensorProto.INT32, shape=["batch", "feature_size", "sequence"] + ) + model.graph.input.insert(1, mask) + + +def chain_model(args): + # Load encoder/decoder and insert necessary (but unused) graph inputs expected by BeamSearch op + encoder_model = onnx.load(args.encoder_path, load_external_data=True) + encoder_model.graph.name = "encoderdecoderinit subgraph" + add_attention_mask(encoder_model) + + decoder_model = onnx.load(args.decoder_path, load_external_data=True) + decoder_model.graph.name = "decoder subgraph" + add_attention_mask(decoder_model) + + config = WhisperConfig.from_pretrained(args.model_name_or_path) + pad_token_id = config.pad_token_id + + beam_inputs = [ + "input_features", + "max_length", + "min_length", + "num_beams", + "num_return_sequences", + "length_penalty", + "repetition_penalty", + "", + "", + "attention_mask", + ] + beam_outputs = ["sequences"] + + node = helper.make_node("BeamSearch", inputs=beam_inputs, outputs=beam_outputs, name="BeamSearch_zcode") + node.domain = "com.microsoft" + node.attribute.extend( + [ + helper.make_attribute("eos_token_id", 50256), + helper.make_attribute("pad_token_id", pad_token_id), + helper.make_attribute("decoder_start_token_id", 50257), + helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size), + helper.make_attribute("early_stopping", True), + helper.make_attribute("model_type", 2), + helper.make_attribute("decoder", decoder_model.graph), + helper.make_attribute("encoder", encoder_model.graph), + ] + ) + + # beam graph inputs + input_features = helper.make_tensor_value_info( + "input_features", TensorProto.FLOAT, ["batch_size", "feature_size", "sequence_length"] + ) + max_length = helper.make_tensor_value_info("max_length", TensorProto.INT32, [1]) + min_length = helper.make_tensor_value_info("min_length", TensorProto.INT32, [1]) + num_beams = helper.make_tensor_value_info("num_beams", TensorProto.INT32, [1]) + num_return_sequences = helper.make_tensor_value_info("num_return_sequences", TensorProto.INT32, [1]) + length_penalty = helper.make_tensor_value_info("length_penalty", TensorProto.FLOAT, [1]) + repetition_penalty = helper.make_tensor_value_info("repetition_penalty", TensorProto.FLOAT, [1]) + attention_mask = helper.make_tensor_value_info( + "attention_mask", TensorProto.INT32, ["batch_size", "feature_size", "sequence_length"] + ) + + graph_inputs = [ + input_features, + max_length, + min_length, + num_beams, + num_return_sequences, + length_penalty, + repetition_penalty, + attention_mask, + ] + + # graph outputs + sequences = helper.make_tensor_value_info( + "sequences", TensorProto.INT32, ["batch_size", "num_return_sequences", "max_length"] + ) + graph_outputs = [sequences] + + # Initializers/opsets + initializers = [] + opset_import = [helper.make_opsetid(domain="com.microsoft", version=1), helper.make_opsetid(domain="", version=17)] + + beam_graph = helper.make_graph([node], "beam-search-test", graph_inputs, graph_outputs, initializers) + beam_model = helper.make_model(beam_graph, producer_name="pytorch", opset_imports=opset_import) + + onnx.save( + beam_model, + args.beam_model_output_dir, + save_as_external_data=True, + all_tensors_to_one_file=True, + convert_attribute=True, + ) + onnx.checker.check_model(args.beam_model_output_dir, full_check=True) diff --git a/onnxruntime/python/tools/transformers/models/whisper/whisper_decoder.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_decoder.py new file mode 100644 index 0000000000..856809a020 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_decoder.py @@ -0,0 +1,416 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import logging +import os +import sys +import tempfile +from pathlib import Path +from typing import List, Union + +import numpy +import onnx +import torch +from transformers import WhisperConfig, file_utils +from whisper_encoder import WhisperEncoderInputs + +from onnxruntime import InferenceSession + +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) +from io_binding_helper import TypeHelper # noqa: E402 +from models.t5.past_helper import PastKeyValuesHelper # noqa: E402 +from onnx_model import OnnxModel # noqa: E402 +from torch_onnx_export_helper import torch_onnx_export # noqa: E402 + +logger = logging.getLogger(__name__) + + +class WhisperDecoderInit(torch.nn.Module): + """A Whisper 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: WhisperConfig, + 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, + ): + encoder_outputs = file_utils.ModelOutput() + encoder_outputs["last_hidden_state"] = encoder_hidden_states + encoder_outputs["hidden_states"] = None + encoder_outputs["attentions"] = None + + out = self.decoder.model( + None, + encoder_outputs=encoder_outputs, + decoder_input_ids=decoder_input_ids, + head_mask=encoder_attention_mask, + past_key_values=None, + use_cache=True, + return_dict=True, + ) + logits = self.decoder.proj_out(out[0]) + return logits, out.past_key_values, out.encoder_last_hidden_state + + +class WhisperDecoder(torch.nn.Module): + """A Whisper 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, *past): + encoder_outputs = file_utils.ModelOutput() + dummy_encoder_hidden_states = torch.randn((decoder_input_ids.shape[0], 3000, int(self.config.d_model))) + encoder_outputs["last_hidden_state"] = dummy_encoder_hidden_states + encoder_outputs["hidden_states"] = dummy_encoder_hidden_states + encoder_outputs["attentions"] = None + if len(past) == 0: + past_key_values = None + else: + past_key_values = PastKeyValuesHelper.back_group_by_layer(past) + + decoder_out = self.decoder( + None, + encoder_outputs=encoder_outputs, + decoder_input_ids=decoder_input_ids, + # decoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=True, + return_dict=True, + ) + logits = decoder_out[0] + present_self, _ = PastKeyValuesHelper.group_by_self_and_cross(decoder_out.past_key_values) + return logits, present_self + + +class WhisperDecoderInputs: + def __init__( + self, + decoder_input_ids, + encoder_attention_mask=None, + past_key_values=None, + ): + self.decoder_input_ids: torch.LongTensor = decoder_input_ids + self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask + self.past_key_values: Union[List[torch.FloatTensor], List[torch.HalfTensor], None] = past_key_values + + @staticmethod + def create_dummy( + config: WhisperConfig, + batch_size: int, + encode_sequence_length: int, + past_decode_sequence_length: int, + device: torch.device, + float16: bool = False, + use_int32_inputs: bool = False, + ): # -> WhisperDecoderInputs: + """Create dummy inputs for WhisperDecoder. + + 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 + float16 (bool): whether the model uses float32 or float16 in input + use_int32_inputs(bool): whether use int32 instead of int64 for some inputs + + Returns: + WhisperDecoderInputs: dummy inputs for decoder + """ + num_attention_heads: int = config.encoder_attention_heads + num_layers: int = config.decoder_layers # + config.encoder_layers + vocab_size: int = config.vocab_size + + # Use head_size, use hidden_size / num_attention_heads here. + # For example, whisper-large, d_model=1280 and num_heads=20 + head_size: int = config.d_model // config.encoder_attention_heads + + 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.int32 if use_int32_inputs else torch.int64), + device=device, + ) + + encoder_inputs = WhisperEncoderInputs.create_dummy( + batch_size, + encode_sequence_length, + vocab_size, + device, + use_int32_inputs=use_int32_inputs, + ) + + float_type = torch.float16 if float16 else torch.float32 + + if past_decode_sequence_length > 0: + self_attention_past_shape = [ + batch_size, + num_attention_heads, + past_decode_sequence_length, + head_size, + ] + cross_attention_past_shape = [ + batch_size, + num_attention_heads, + encode_sequence_length, + head_size, + ] + + past = [] + for _ in range(2 * num_layers): + 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=float_type, device=device)) + else: + past = None + + encoder_attention_mask = torch.zeros( + (encoder_inputs.input_ids.shape[0], 1, encoder_inputs.input_ids.shape[1], encoder_inputs.input_ids.shape[1]) + ).type(torch.int8) + return WhisperDecoderInputs(decoder_input_ids, encoder_attention_mask, past) + + def to_list(self) -> List: + input_list = [ + self.decoder_input_ids, + self.encoder_attention_mask, + ] + if self.past_key_values: + input_list.extend(self.past_key_values) + return input_list + + def to_fp32(self): + past = [p.to(dtype=torch.float32) for p in self.past_key_values] if self.past_key_values else None + return WhisperDecoderInputs( + self.decoder_input_ids.clone(), + self.encoder_attention_mask.clone(), + past, + ) + + +class WhisperDecoderHelper: + @staticmethod + def export_onnx( + decoder: WhisperDecoder, + device: torch.device, + onnx_model_path: str, + verbose: bool = True, + use_external_data_format: bool = False, + use_int32_inputs: bool = False, + ): + """Export decoder to ONNX + + Args: + decoder (Union[WhisperDecoder, WhisperDecoderNoPastState]): 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. + use_int32_inputs (bool, optional): use int32 inputs + """ + assert isinstance(decoder, (WhisperDecoder, WhisperDecoderInit)) + + inputs = WhisperDecoderInputs.create_dummy( + decoder.config, + batch_size=2, + encode_sequence_length=3000, + past_decode_sequence_length=5 if isinstance(decoder, WhisperDecoder) else 0, + device=device, + use_int32_inputs=use_int32_inputs, + ) + input_list = inputs.to_list() + + # Fix past disappearing bug - duplicate first past entry + # input_list.insert(2, input_list[2]) + + past_names = PastKeyValuesHelper.get_past_names(decoder.config.decoder_layers, present=False) + present_names = PastKeyValuesHelper.get_past_names(decoder.config.decoder_layers, present=True) + present_self_names = present_names[: 2 * decoder.config.decoder_layers] + + input_past_names = past_names if isinstance(decoder, WhisperDecoder) else [] + output_present_names = present_self_names if isinstance(decoder, WhisperDecoder) 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) + # 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 = ["input_ids"] + input_names.append("encoder_attention_mask") + input_names.extend(input_past_names) + + dynamic_axes = { + "input_ids": {0: "batch_size"}, + "encoder_attention_mask": {0: "batch_size", 1: "encode_sequence_length"}, + "encoder_hidden_states": {0: "batch_size", 1: "encode_sequence_length / 2"}, + "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, WhisperDecoder): + 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) + + with tempfile.TemporaryDirectory() as tmp_dir_name: + temp_onnx_model_path = os.path.join(tmp_dir_name, "decoder.onnx") + Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + torch_onnx_export( + decoder, + args=tuple(input_list), + f=temp_onnx_model_path if use_external_data_format else onnx_model_path, + export_params=True, + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + opset_version=17, + do_constant_folding=True, + use_external_data_format=use_external_data_format, + verbose=verbose, + ) + + if use_external_data_format: + model = onnx.load_model(temp_onnx_model_path, load_external_data=True) + OnnxModel.save( + model, + onnx_model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + ) + + @staticmethod + def onnxruntime_inference(ort_session, inputs: WhisperDecoderInputs): + """Run inference of ONNX model.""" + logger.debug("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()), + } + + 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[WhisperDecoder, WhisperDecoderInit], + 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.""" + float16: bool = TypeHelper.get_input_type(ort_session, "past_key_self_0") == "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]: + if isinstance(model, WhisperDecoderInit): + dec_seq_len = 0 + else: + dec_seq_len = past_decode_sequence_length + + inputs = WhisperDecoderInputs.create_dummy( + model.config, + batch_size, + encode_sequence_length, + dec_seq_len, + device=device, + float16=float16, + use_int32_inputs=use_int32_inputs, + ) + + # 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(): + torch_outputs = model(*input_list) + + ort_outputs = WhisperDecoderHelper.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, WhisperDecoderInit): + 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}, " + + f"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/whisper/whisper_encoder.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder.py new file mode 100644 index 0000000000..fcf3ce0904 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder.py @@ -0,0 +1,162 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import logging +import os +import sys +import tempfile +from pathlib import Path +from typing import List + +import numpy +import onnx +import torch +from transformers import WhisperConfig + +from onnxruntime import InferenceSession + +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) +from onnx_model import OnnxModel # noqa: E402 +from torch_onnx_export_helper import torch_onnx_export # noqa: E402 + +logger = logging.getLogger(__name__) + + +class WhisperEncoder(torch.nn.Module): + """Whisper encoder outputs only the last hidden state""" + + def __init__(self, encoder, config: WhisperConfig): + super().__init__() + self.encoder = encoder + self.config = config + + def forward(self, input_features, attention_mask): + return self.encoder.model.encoder(input_features)[0] + + +class WhisperEncoderInputs: + def __init__(self, input_features, attention_mask): + self.input_ids: torch.LongTensor = input_features + # HF Whisper model doesn't support Attention Mask functionality + + @staticmethod + def create_dummy( + batch_size: int, sequence_length: int, feature_size: int, device: torch.device, use_int32_inputs: bool + ): + """Create dummy inputs for Whisper encoder. + + Args: + batch_size (int): batch size + sequence_length (int): sequence length + feature_size (int): feature size for spectrogram input + device (torch.device): device of output tensors + + Returns: + WhisperEncoderInputs: dummy inputs for encoder + """ + dtype = torch.float32 + + input_features = torch.randn( + size=(batch_size, feature_size, sequence_length), + device=device, + ) + attention_mask = torch.ones([batch_size, feature_size, sequence_length], dtype=dtype, device=device) + return WhisperEncoderInputs(input_features, attention_mask) + + def to_list(self) -> List: + if self.input_features is None: + return [] + return [self.input_features] + + +class WhisperEncoderHelper: + @staticmethod + def export_onnx( + encoder, + device: torch.device, + onnx_model_path: str, + verbose: bool = True, + use_external_data_format: bool = False, + ): + """Export encoder to ONNX + + Args: + encoder (WhisperEncoder): 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 = WhisperEncoderInputs.create_dummy( + batch_size=2, + sequence_length=3000, + feature_size=config.num_mel_bins, + device=device, + ) + + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory() as tmp_dir_name: + temp_onnx_model_path = os.path.join(tmp_dir_name, "encoder.onnx") + Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + torch_onnx_export( + encoder, + args=tuple(encoder_inputs.to_list()), + f=temp_onnx_model_path if use_external_data_format else onnx_model_path, + export_params=True, + input_names=["input_features"], + output_names=["hidden_states"], + dynamic_axes={ + "input_ids": {0: "batch_size", 1: "feature_size", 2: "sequence_length"}, + "hidden_states": {0: "batch_size", 1: "sequence_length"}, + }, + opset_version=17, + do_constant_folding=True, + use_external_data_format=use_external_data_format, + verbose=verbose, + ) + + if use_external_data_format: + model = onnx.load_model(temp_onnx_model_path, load_external_data=True) + OnnxModel.save( + model, + onnx_model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + ) + + @staticmethod + def onnxruntime_inference(ort_session, inputs: WhisperEncoderInputs): + """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: WhisperEncoder, ort_session: InferenceSession, device: torch.device, use_int32_inputs: bool = False + ): + """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good.""" + inputs = WhisperEncoderInputs.create_dummy( + batch_size=4, + sequence_length=11, + device=device, + use_int32_inputs=use_int32_inputs, + ) + input_list = inputs.to_list() + torch_outputs = model(*input_list) + + ort_outputs = WhisperEncoderHelper.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/whisper/whisper_encoder_decoder_init.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder_decoder_init.py new file mode 100644 index 0000000000..ce8b72d89a --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder_decoder_init.py @@ -0,0 +1,307 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import logging +import os +import sys +import tempfile +from pathlib import Path +from typing import List, Optional + +import numpy +import onnx +import torch +from transformers import WhisperConfig +from whisper_decoder import WhisperDecoderInit +from whisper_encoder import WhisperEncoder, WhisperEncoderInputs + +from onnxruntime import InferenceSession + +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) +from models.t5.past_helper import PastKeyValuesHelper # noqa: E402 +from onnx_model import OnnxModel # noqa: E402 +from torch_onnx_export_helper import torch_onnx_export # noqa: E402 + +logger = logging.getLogger(__name__) + + +class WhisperEncoderDecoderInit(torch.nn.Module): + """A combination of WhisperEncoder and WhisperDecoderInit.""" + + def __init__( + self, + encoder: torch.nn.Module, + decoder: torch.nn.Module, + lm_head: torch.nn.Module, + config: WhisperConfig, + decoder_start_token_id: Optional[int] = None, + ): + super().__init__() + self.config = config + self.whisper_encoder = WhisperEncoder(encoder, config) + self.whisper_decoder_init = WhisperDecoderInit(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.whisper_encoder(encoder_input_ids, None) + # Decoder out: (logits, past_key_values, encoder_hidden_state) + decinit_out = self.whisper_decoder_init(decoder_input_ids, encoder_attention_mask, encoder_hidden_states) + present_self, present_cross = PastKeyValuesHelper.group_by_self_and_cross(decinit_out[1]) + present = present_self + present_cross + return decinit_out[0], encoder_hidden_states, present + + +class WhisperEncoderDecoderInitInputs: + 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: WhisperConfig, + batch_size: int, + encode_sequence_length: int, + use_decoder_input_ids: int, + device: torch.device, + use_int32_inputs: bool = False, + ): # -> WhisperEncoderDecoderInitInputs: + encoder_inputs: WhisperEncoderInputs = WhisperEncoderInputs.create_dummy( + batch_size, + sequence_length=3000, + feature_size=config.num_mel_bins, + device=device, + use_int32_inputs=use_int32_inputs, + ) + decoder_input_ids = None + encoder_attention_mask = torch.zeros( + (encoder_inputs.input_ids.shape[0], 1, encoder_inputs.input_ids.shape[1], encoder_inputs.input_ids.shape[1]) + ).type(torch.int8) + 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 WhisperEncoderDecoderInitInputs(encoder_inputs.input_ids, encoder_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 WhisperEncoderDecoderInitHelper: + @staticmethod + def export_onnx( + model: WhisperEncoderDecoderInit, + 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 (WhisperEncoderDecoderInit): 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, WhisperEncoderDecoderInit) + + inputs = WhisperEncoderDecoderInitInputs.create_dummy( + model.config, + batch_size=2, + encode_sequence_length=3000, + use_decoder_input_ids=use_decoder_input_ids, + device=device, + use_int32_inputs=use_int32_inputs, + ) + input_list = inputs.to_list() + + out = model(inputs.encoder_input_ids, inputs.encoder_attention_mask, inputs.decoder_input_ids) + present = out[2] + # pdb.set_trace() + present_names = PastKeyValuesHelper.get_input_names(present, encoder=True) + # 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, 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 'Transposepresent_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.encoder_attention_heads) + hidden_size = str(model.config.d_model) + head_size = str(model.config.d_model // model.config.encoder_attention_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, + } + + 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=17, + do_constant_folding=True, + use_external_data_format=use_external_data_format, + verbose=verbose, + ) + + # Workaround as mentioned earlier: change numeric dim_param to dim_value + model = onnx.load(temp_onnx_model_path) + for tensor in 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 + + OnnxModel.save( + 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: WhisperEncoderDecoderInitInputs): + """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: WhisperEncoderDecoderInit, + 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 = WhisperEncoderDecoderInitInputs.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 = WhisperEncoderDecoderInitHelper.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/whisper/whisper_helper.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_helper.py new file mode 100644 index 0000000000..6d444bdd87 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_helper.py @@ -0,0 +1,265 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import logging +import os +import sys +from pathlib import Path +from typing import Dict, Tuple, Union + +import torch +from transformers import WhisperForConditionalGeneration +from whisper_decoder import WhisperDecoder, WhisperDecoderHelper, WhisperDecoderInit +from whisper_encoder import WhisperEncoder, WhisperEncoderHelper +from whisper_encoder_decoder_init import WhisperEncoderDecoderInit, WhisperEncoderDecoderInitHelper + +from onnxruntime import InferenceSession + +sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) +from float16 import float_to_float16_max_diff # noqa: E402 +from onnx_model import OnnxModel # noqa: E402 +from optimizer import optimize_model # noqa: E402 + +logger = logging.getLogger(__name__) + +PRETRAINED_WHISPER_MODELS = [ + "whisper-tiny", + "whisper-tiny.en", + "whisper-small", + "whisper-small.en", + "whisper-medium", + "whisper-medium.en", + "whisper-base", + "whisper-base.en", + "whisper-large", + "whisper-large-v2", +] + + +class WhisperHelper: + @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 + + directory = os.path.join(output_dir, model_name) if new_folder else output_dir + return os.path.join(directory, 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 = WhisperForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir) + + decoder = WhisperDecoder(model, None, model.config) + decoder.eval().to(device) + + if merge_encoder_and_decoder_init: + encoder_decoder_init = WhisperEncoderDecoderInit( + model, + model, + None, + model.config, + decoder_start_token_id=None, + ) + return {"encoder_decoder_init": encoder_decoder_init, "decoder": decoder} + else: + encoder = WhisperEncoder(model.model.encoder, model.config) + encoder.eval().to(device) + decoder_init = WhisperDecoderInit(model.decoder, None, model.config) + decoder_init.eval().to(device) + return { + "encoder": encoder, + "decoder": decoder, + "decoder_init": decoder_init, + } + + @staticmethod + def export_onnx( + model: Union[WhisperEncoder, WhisperDecoder, WhisperDecoderInit, WhisperEncoderDecoderInit], + device: torch.device, + onnx_model_path: str, + verbose: bool = True, + use_external_data_format: bool = False, + use_decoder_input_ids: bool = True, + use_int32_inputs: bool = False, + ): + if isinstance(model, WhisperEncoder): + WhisperEncoderHelper.export_onnx( + model, + device, + onnx_model_path, + verbose, + use_external_data_format, + ) + elif isinstance(model, WhisperEncoderDecoderInit): + WhisperEncoderDecoderInitHelper.export_onnx( + model, + device, + onnx_model_path, + use_decoder_input_ids, + verbose, + use_external_data_format, + use_int32_inputs, + ) + else: + WhisperDecoderHelper.export_onnx( + model, + device, + onnx_model_path, + verbose, + use_external_data_format, + use_int32_inputs, + ) + + @staticmethod + def auto_mixed_precision( + onnx_model: OnnxModel, + op_block_list: Tuple[str] = ( + "SimplifiedLayerNormalization", + "SkipSimplifiedLayerNormalization", + "Relu", + "Add", + ), + ): + """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 ["SimplifiedLayerNormalization", "SkipSimplifiedLayerNormalization", "Relu", "Add"] + 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": 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) + + 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, + auto_mixed_precision: bool = True, + use_gpu: bool = False, + ): + """Optimize ONNX model with an option to convert it to use mixed precision.""" + + from fusion_options import FusionOptions + + optimization_options = None + if is_float16: + optimization_options = FusionOptions("bart") + optimization_options.enable_skip_layer_norm = False + + m = optimize_model( + onnx_model_path, + model_type="bart", + num_heads=num_attention_heads, + hidden_size=hidden_size, + opt_level=2 if not use_external_data_format else None, + optimization_options=optimization_options, + use_gpu=False, + only_onnxruntime=False, + ) + + if is_float16: + if auto_mixed_precision: + WhisperHelper.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, all_tensors_to_one_file=True) + + @staticmethod + def verify_onnx( + model: Union[WhisperEncoder, WhisperDecoder, WhisperDecoderInit, WhisperEncoderDecoderInit], + ort_session: InferenceSession, + device: torch.device, + use_int32_inputs: bool, + ): + """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good.""" + # Not implemented for Whisper currently + return True