# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # ------------------------------------------------------------------------- import logging import os from pathlib import Path import torch from float16 import float_to_float16_max_diff from onnx_model import OnnxModel from optimizer import optimize_model from t5_decoder import T5Decoder, T5DecoderHelper from t5_encoder_decoder_init import T5EncoderDecoderInit, T5EncoderDecoderInitHelper from transformers import MT5ForConditionalGeneration, T5ForConditionalGeneration from onnxruntime import InferenceSession logger = logging.getLogger(__name__) PRETRAINED_T5_MODELS = ["t5-small", "t5-base", "t5-large", "t5-3b", "t5-11b"] PRETRAINED_MT5_MODELS = [ "google/mt5-small", "google/mt5-base", "google/mt5-large", "google/mt5-xl", "google/mt5-xxl", ] 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 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, model_type: str = "t5", state_dict_path: str = "", encoder_decoder_init: bool = False, ) -> dict[str, T5EncoderDecoderInit | T5Decoder]: """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 model_type (str, optional): model type "t5" or "mt5" state_dict_path(str, optional): state dictionary path encoder_decoder_init (bool, optional): combine encoder and decoder kv cache initialization into one model. Returns: Dict[str, torch.nn.Module]: mapping from name to modules for ONNX conversion. """ if model_type == "t5": model = T5ForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir) elif model_type == "mt5": model = MT5ForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir) else: raise ValueError("only support mode_type=t5 or mt5") if state_dict_path: model.load_state_dict(torch.load(state_dict_path)) decoder = T5Decoder(model.decoder, model.lm_head, model.config) decoder.eval().to(device) encoder = T5EncoderDecoderInit( model.encoder, model.decoder, model.lm_head, model.config, decoder_start_token_id=None, output_cross_only=not encoder_decoder_init, ) encoder_name = "encoder_decoder_init" if encoder_decoder_init else "encoder" return {encoder_name: encoder, "decoder": decoder} @staticmethod def export_onnx( model: T5Decoder | T5EncoderDecoderInit, 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, T5EncoderDecoderInit): T5EncoderDecoderInitHelper.export_onnx( model, device, onnx_model_path, use_decoder_input_ids, verbose, use_external_data_format, use_int32_inputs, ) else: T5DecoderHelper.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: list[str] | None = None, force_fp16_logits: bool = False, use_symbolic_shape_infer: bool = True, ): """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): operators need to run in fp32. force_fp16_logits (bool, optional): force logits and last MatMul node to be in float16. Defaults to False. use_symbolic_shape_infer (bool, optional): use symbolic shape inference to convert float to float16. Defaults to True. Returns: parameters(dict): a dictionary of parameters used in float16 conversion """ if op_block_list is None: op_block_list = [ "SimplifiedLayerNormalization", "SkipSimplifiedLayerNormalization", "Relu", "Add", ] op_full_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) and not force_fp16_logits: # 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] if "Add" not in op_block_list: input_name_to_nodes = onnx_model.input_name_to_nodes() fp32_add = 0 changed = True add_nodes = onnx_model.get_nodes_by_op_type("Add") while changed: changed = False for node in add_nodes: if node.name not in node_block_list: parents = onnx_model.get_parents(node, output_name_to_node) children = onnx_model.get_children(node, input_name_to_nodes) blocked_children = [ child for child in children if child.op_type in op_block_list or child in node_block_list ] blocked_parents = [ parent for parent in parents if parent.op_type in op_block_list or parent in node_block_list ] # If any child or parent is in fp32, we place the Add node to fp32. if (len(blocked_children) + len(blocked_parents)) > 0: node_block_list.append(node.name) fp32_add += 1 changed = True fp16_add = len(add_nodes) - fp32_add logger.info(f"node counter of Add operator: fp32={fp32_add} fp16={fp16_add}") logger.info(f"node_block_list: {node_block_list}") parameters = { "keep_io_types": keep_io_types, "op_block_list": op_block_list, "node_block_list": node_block_list, "force_fp16_initializers": is_weight_fp16_precision, } logger.info(f"auto_mixed_precision parameters: {parameters}") if use_symbolic_shape_infer: onnx_model.convert_float_to_float16(use_symbolic_shape_infer=True, **parameters) else: # Workaround when symbolic shape inference fails. # Need enable shape_infer_before_optimization in convert_to_onnx.py as well. from float16 import convert_float_to_float16 convert_float_to_float16( onnx_model.model, disable_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, force_fp16_io: 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("t5") # SkipLayerNormalization is faster but might bring accuracy drop since it uses fp16 accumulation. optimization_options.enable_skip_layer_norm = not auto_mixed_precision m = optimize_model( onnx_model_path, model_type="t5", num_heads=num_attention_heads, hidden_size=hidden_size, opt_level=0, optimization_options=optimization_options, use_gpu=use_gpu, ) if is_float16: if auto_mixed_precision: T5Helper.auto_mixed_precision(m, force_fp16_logits=force_fp16_io) else: m.convert_model_float32_to_float16(cast_input_output=force_fp16_io) m.save_model_to_file(optimized_model_path, use_external_data_format, all_tensors_to_one_file=True) @staticmethod def verify_onnx( model: T5Decoder | T5EncoderDecoderInit, ort_session: InferenceSession, device: torch.device, use_int32_inputs: bool, ): """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good.""" if isinstance(model, T5EncoderDecoderInit): return T5EncoderDecoderInitHelper.verify_onnx(model, ort_session, device, use_int32_inputs) return T5DecoderHelper.verify_onnx(model, ort_session, device, use_int32_inputs)