saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-19 19:00:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

auto mixed precision for t5 (#11252)

Browse source
This commit is contained in:
Tianlei Wu 2022-04-19 12:42:11 -07:00 committed by GitHub
parent 5ee8e2e491
commit bab9b80f1f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 128 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py
View file

@ -76,6 +76,12 @@ def parse_arguments():
parser.add_argument('-w', '--overwrite', required=False, action='store_true', help="overwrite existing ONNX model")
parser.set_defaults(overwrite=False)
parser.add_argument('--disable_auto_mixed_precision',
required=False,
action='store_true',
help="use pure fp16 instead of mixed precision")
parser.set_defaults(disable_auto_mixed_precision=False)
args = parser.parse_args()
return args
@ -91,7 +97,8 @@ def export_onnx_models(model_name_or_path,
verbose,
use_decoder_start_token: bool = True,
merge_encoder_and_decoder_init: bool = True,
overwrite: bool = False):
overwrite: bool = False,
disable_auto_mixed_precision: bool = False):
device = torch.device("cuda:0" if use_gpu else "cpu")
models = T5Helper.load_model(model_name_or_path, cache_dir, device, merge_encoder_and_decoder_init)
@ -102,6 +109,7 @@ def export_onnx_models(model_name_or_path,
output_paths = []
for name, model in models.items():
model.to(device)
filename_suffix = "_" + name
onnx_path = T5Helper.get_onnx_path(output_dir,
@ -112,7 +120,8 @@ def export_onnx_models(model_name_or_path,
if overwrite or not os.path.exists(onnx_path):
logger.info(f"Exporting ONNX model to {onnx_path}")
# We have to clone model before exporting onnx, otherwise verify_onnx will report large difference.
T5Helper.export_onnx(copy.deepcopy(model),
cloned_model = copy.deepcopy(model).to(device)
T5Helper.export_onnx(cloned_model,
device,
onnx_path,
verbose,
@ -130,8 +139,13 @@ def export_onnx_models(model_name_or_path,
if overwrite or not os.path.exists(output_path):
logger.info(f"Optimizing model to {output_path}")
T5Helper.optimize_onnx(onnx_path, output_path, precision == Precision.FLOAT16, config.num_heads,
config.hidden_size, use_external_data_format)
T5Helper.optimize_onnx(onnx_path,
output_path,
precision == Precision.FLOAT16,
config.num_heads,
config.hidden_size,
use_external_data_format,
auto_mixed_precision=not disable_auto_mixed_precision)
else:
logger.info(f"Skip optimizing: existed ONNX model {onnx_path}")
else:
@ -175,7 +189,7 @@ def main():
output_paths = export_onnx_models(args.model_name_or_path, cache_dir, output_dir, args.use_gpu,
args.use_external_data_format, args.optimize_onnx, args.precision,
args.verbose, args.use_decoder_start_token, merge_encoder_and_decoder_init,
args.overwrite)
args.overwrite, args.disable_auto_mixed_precision)
logger.info(f"Done! Outputs: {output_paths}")

36
onnxruntime/python/tools/transformers/models/t5/t5_decoder.py
View file

@ -18,6 +18,7 @@ from past_helper import PastKeyValuesHelper
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from torch_onnx_export_helper import torch_onnx_export
from io_binding_helper import TypeHelper
logger = logging.getLogger(__name__)
@ -26,6 +27,7 @@ class T5DecoderInit(torch.nn.Module):
""" A T5 decoder with LM head to create initial past key values.
This model is only called once during starting decoding.
"""
def __init__(self,
decoder: torch.nn.Module,
lm_head: torch.nn.Module,
@ -62,6 +64,7 @@ class T5DecoderInit(torch.nn.Module):
class T5Decoder(torch.nn.Module):
""" A T5 decoder with LM head and past key values"""
def __init__(self, decoder, lm_head, config):
super().__init__()
self.decoder = decoder
@ -92,6 +95,7 @@ class T5Decoder(torch.nn.Module):
class T5DecoderInputs:
def __init__(self, decoder_input_ids, encoder_attention_mask, encoder_hidden_states, past_key_values=None):
self.decoder_input_ids: torch.LongTensor = decoder_input_ids
self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask
@ -99,8 +103,12 @@ class T5DecoderInputs:
self.past_key_values: Union[List[torch.FloatTensor], List[torch.HalfTensor], None] = past_key_values
@staticmethod
def create_dummy(config: T5Config, batch_size: int, encode_sequence_length: int, past_decode_sequence_length: int,
device: torch.device): # -> T5DecoderInputs:
def create_dummy(config: T5Config,
batch_size: int,
encode_sequence_length: int,
past_decode_sequence_length: int,
device: torch.device,
float16: bool = False): # -> T5DecoderInputs:
""" Create dummy inputs for T5Decoder.
Args:
@ -109,6 +117,7 @@ class T5DecoderInputs:
encode_sequence_length (int): sequence length of input_ids for encoder
past_decode_sequence_length (int): past sequence length of input_ids for decoder
device (torch.device): device of output tensors
float16 (bool): whether the model uses float32 or float16 in input
Returns:
T5DecoderInputs: dummy inputs for decoder
@ -127,10 +136,11 @@ class T5DecoderInputs:
encoder_inputs = T5EncoderInputs.create_dummy(batch_size, encode_sequence_length, vocab_size, device)
float_type = torch.float16 if float16 else torch.float32
encoder_hidden_state = torch.rand(batch_size,
encode_sequence_length,
hidden_size,
dtype=torch.float32,
dtype=float_type,
device=device)
if past_decode_sequence_length > 0:
@ -145,10 +155,10 @@ class T5DecoderInputs:
past = []
for _ in range(2 * num_layers):
past.append(torch.rand(self_attention_past_shape, dtype=torch.float32, device=device))
past.append(torch.rand(self_attention_past_shape, dtype=float_type, device=device))
for _ in range(2 * num_layers):
past.append(torch.rand(cross_attention_past_shape, dtype=torch.float32, device=device))
past.append(torch.rand(cross_attention_past_shape, dtype=float_type, device=device))
else:
past = None
@ -160,8 +170,15 @@ class T5DecoderInputs:
input_list.extend(self.past_key_values)
return input_list
def to_fp32(self):
encoder_hidden_state = self.encoder_hidden_states.to(dtype=torch.float32)
past = [p.to(dtype=torch.float32) for p in self.past_key_values]
return T5DecoderInputs(self.decoder_input_ids.clone(), self.encoder_attention_mask.clone(),
encoder_hidden_state, past)
class T5DecoderHelper:
@staticmethod
def export_onnx(decoder: Union[T5Decoder, T5DecoderInit],
device: torch.device,
@ -292,6 +309,8 @@ class T5DecoderHelper:
max_cases=4):
""" Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good.
"""
float16: bool = (TypeHelper.get_input_type(ort_session, "encoder_hidden_states") == "tensor(float16)")
test_cases = [(4, 11, 3), (1, 2, 5), (3, 1, 1), (8, 5, 2)]
test_cases_max_diff = []
for (batch_size, encode_sequence_length, past_decode_sequence_length) in test_cases[:max_cases]:
@ -302,8 +321,11 @@ class T5DecoderHelper:
batch_size,
encode_sequence_length,
past_decode_sequence_length,
device=device)
input_list = inputs.to_list()
device=device,
float16=float16)
# We use fp32 PyTroch model as baseline even when ONNX model is fp16
input_list = inputs.to_fp32().to_list()
# Run inference of PyTorch model
with torch.no_grad():

85
onnxruntime/python/tools/transformers/models/t5/t5_helper.py
View file

@ -7,7 +7,7 @@
import os
import sys
from pathlib import Path
from typing import Union, Dict
from typing import Union, Dict, List
import logging
import torch
from transformers import T5ForConditionalGeneration
@ -16,12 +16,19 @@ from t5_encoder import T5Encoder, T5EncoderHelper
from t5_decoder import T5DecoderInit, T5Decoder, T5DecoderHelper
from t5_encoder_decoder_init import T5EncoderDecoderInit, T5EncoderDecoderInitHelper
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from onnx_model import OnnxModel
from float16 import float_to_float16_max_diff
from fusion_utils import FusionUtils
from optimizer import optimize_model
logger = logging.getLogger(__name__)
PRETRAINED_T5_MODELS = ["t5-small", "t5-base", "t5-large", "t5-3B", "t5-11B"]
class T5Helper:
@staticmethod
def get_onnx_path(output_dir: str, model_name_or_path: str, suffix: str = "", new_folder: bool = False) -> str:
"""Build onnx path
@ -96,17 +103,82 @@ class T5Helper:
else:
T5DecoderHelper.export_onnx(model, device, onnx_model_path, verbose, use_external_data_format)
@staticmethod
def auto_mixed_precision(
onnx_model: OnnxModel,
op_block_list: List[str] = ["Pow", "ReduceMean", "Add", "Sqrt", "Div", "Mul", "Softmax", "Relu"]):
"""Convert model to mixed precision.
It detects whether original model has fp16 precision weights, and set parameters for float16 conversion automatically.
Args:
onnx_model (OnnxModel): optimized ONNX model
op_block_list (List[str], optional): . Defaults to ["Pow", "ReduceMean", "Add", "Sqrt", "Div", "Mul", "Softmax", "Relu"]
Returns:
parameters(dict): a dictionary of parameters used in float16 conversion
"""
op_full_set = set([node.op_type for node in onnx_model.nodes()])
fp32_op_set = set(op_block_list)
fp16_op_set = op_full_set.difference(fp32_op_set)
logger.info(f"fp32 op: {fp32_op_set} fp16 op: {fp16_op_set}")
# logits is the first output
logits_output_name = onnx_model.graph().output[0].name
# We use the weight in last MatMul node to detect whether the model is stored with float16 weights from training.
is_weight_fp16_precision = False
output_name_to_node = onnx_model.output_name_to_node()
assert logits_output_name in output_name_to_node
node = output_name_to_node[logits_output_name]
last_matmul_node = None
if node.op_type == "MatMul":
last_matmul_node = node
logger.info(f"Found last MatMul node for logits: {node.name}")
initializer = None
for input in node.input:
initializer = onnx_model.get_initializer(input)
if initializer is not None:
break
# when the max difference of value after converting float to float16 is lower than a threshold (1e-6),
# we can deduce that the weights are stored in float16 precision.
max_diff = float_to_float16_max_diff(initializer)
logger.debug(f"max diff of converting weights in last MatMul node {node.name}: {max_diff}")
is_weight_fp16_precision = (max_diff < 1E-6)
else:
logger.warning(f"Failed to find MatMul node for logits. Found {node.op_type} of node {node.name}")
keep_io_types = []
node_block_list = []
if (not is_weight_fp16_precision) and (last_matmul_node is not None):
# When original weight is float32 precision, keep logits and last MatMul in float32 could get better precision.
keep_io_types = [logits_output_name]
node_block_list = [last_matmul_node.name]
parameters = {
"keep_io_types": keep_io_types,
"op_block_list": op_block_list,
"node_block_list": node_block_list,
"force_fp16_initializers": is_weight_fp16_precision
}
logger.info(f"auto_mixed_precision parameters: {parameters}")
onnx_model.convert_float_to_float16(use_symbolic_shape_infer=True, **parameters)
fusion_utils = FusionUtils(onnx_model)
fusion_utils.remove_cascaded_cast_nodes()
fusion_utils.remove_useless_cast_nodes()
return parameters
@staticmethod
def optimize_onnx(onnx_model_path: str,
optimized_model_path: str,
is_float16: bool,
num_attention_heads: int,
hidden_size: int,
use_external_data_format: bool = False):
use_external_data_format: bool = False,
auto_mixed_precision: bool = True):
""" Optimize ONNX model with an option to convert it to use mixed precision.
"""
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from optimizer import optimize_model
m = optimize_model(onnx_model_path,
model_type='bert',
num_heads=num_attention_heads,
@ -115,7 +187,10 @@ class T5Helper:
optimization_options=None,
use_gpu=False)
if is_float16:
m.convert_model_float32_to_float16(cast_input_output=False)
if auto_mixed_precision:
T5Helper.auto_mixed_precision(m)
else:
m.convert_model_float32_to_float16(cast_input_output=False)
m.save_model_to_file(optimized_model_path, use_external_data_format)