diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 4cd0d676c7..230f5ebb8b 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -364,6 +364,9 @@ file(GLOB onnxruntime_python_quantization_cal_table_flatbuffers_src CONFIGURE_DE file(GLOB onnxruntime_python_transformers_src CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/python/tools/transformers/*.py" ) +file(GLOB onnxruntime_python_transformers_models_gpt2_src CONFIGURE_DEPENDS + "${ONNXRUNTIME_ROOT}/python/tools/transformers/models/gpt2/*.py" +) file(GLOB onnxruntime_python_transformers_models_longformer_src CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/python/tools/transformers/models/longformer/*.py" ) @@ -420,6 +423,7 @@ add_custom_command( COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/tools/ort_format_model/ort_flatbuffers_py COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/transformers COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/transformers/models + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/transformers/models/gpt2 COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/transformers/models/longformer COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/transformers/models/t5 COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/quantization @@ -501,6 +505,9 @@ add_custom_command( COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_python_transformers_src} $/onnxruntime/transformers/ + COMMAND ${CMAKE_COMMAND} -E copy + ${onnxruntime_python_transformers_models_gpt2_src} + $/onnxruntime/transformers/models/gpt2/ COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_python_transformers_models_longformer_src} $/onnxruntime/transformers/models/longformer/ diff --git a/onnxruntime/python/tools/transformers/README.md b/onnxruntime/python/tools/transformers/README.md index 5316aa2a61..3a9cb9cb8a 100644 --- a/onnxruntime/python/tools/transformers/README.md +++ b/onnxruntime/python/tools/transformers/README.md @@ -66,9 +66,9 @@ In your python code, you can use the optimizer like the following: ```python from onnxruntime.transformers import optimizer -optimized_model = optimizer.optimize_model("gpt2.onnx", model_type='gpt2', num_heads=12, hidden_size=768) +optimized_model = optimizer.optimize_model("bert.onnx", model_type='bert', num_heads=12, hidden_size=768) optimized_model.convert_float_to_float16() -optimized_model.save_model_to_file("gpt2_fp16.onnx") +optimized_model.save_model_to_file("bert_fp16.onnx") ``` You can also use command line. Example of optimizing a BERT-large model to use mixed precision (float16): @@ -78,7 +78,7 @@ python -m onnxruntime.transformers.optimizer --input bert_large.onnx --output be You can also download the latest script files from [here](https://github.com/microsoft/onnxruntime/tree/master/onnxruntime/python/tools/transformers/). Then run it like the following: ```console -python optimizer.py --input gpt2.onnx --output gpt2_opt.onnx --model_type gpt2 +python optimizer.py --input bert.onnx --output bert_opt.onnx --model_type bert ``` ### Optimizer Options @@ -164,11 +164,11 @@ The model has 12 layers and 768 hidden, with input_ids, position_ids, attention_ Since past state is used, sequence length in input_ids is 1. For example, s=4 means the past sequence length is 4 and the total sequence length is 5. -[benchmark_gpt2.py](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/transformers/benchmark_gpt2.py) is used to get the results like the following commands: +[benchmark_gpt2.py](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py) is used to get the results like the following commands: ```console -python -m onnxruntime.transformers.benchmark_gpt2 --use_gpu -m gpt2 -o -v -b 1 8 32 128 -s 4 8 32 128 -p fp32 -python -m onnxruntime.transformers.benchmark_gpt2 --use_gpu -m gpt2 -o -v -b 1 8 32 128 -s 4 8 32 128 -p fp16 +python -m onnxruntime.transformers.models.gpt2.benchmark_gpt2 --use_gpu -m gpt2 -o -v -b 1 8 32 128 -s 4 8 32 128 -p fp32 +python -m onnxruntime.transformers.models.gpt2.benchmark_gpt2 --use_gpu -m gpt2 -o -v -b 1 8 32 128 -s 4 8 32 128 -p fp16 ``` ### Benchmark.py diff --git a/onnxruntime/python/tools/transformers/__init__.py b/onnxruntime/python/tools/transformers/__init__.py index ad5632855c..8302e1b7c4 100644 --- a/onnxruntime/python/tools/transformers/__init__.py +++ b/onnxruntime/python/tools/transformers/__init__.py @@ -1,4 +1,13 @@ -import os -import sys +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +#-------------------------------------------------------------------------- -sys.path.append(os.path.dirname(__file__)) +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), 'models', 'gpt2')) + +# added for backward compatible +import gpt2_helper +import convert_to_onnx diff --git a/onnxruntime/python/tools/transformers/convert_beam_search.py b/onnxruntime/python/tools/transformers/convert_beam_search.py index 8efef8c15f..2e6f5746de 100644 --- a/onnxruntime/python/tools/transformers/convert_beam_search.py +++ b/onnxruntime/python/tools/transformers/convert_beam_search.py @@ -2,7 +2,6 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. #------------------------------------------------------------------------- - """ This converts GPT2 or T5 model to onnx with beam search operator. @@ -26,11 +25,15 @@ from typing import List, Union import torch from packaging import version from transformers import GPT2Config, T5Config -from gpt2_helper import PRETRAINED_GPT2_MODELS -from convert_to_onnx import main as convert_gpt2_to_onnx from benchmark_helper import Precision from onnx import onnx_pb as onnx_proto +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), 'models', 'gpt2')) +from gpt2_helper import PRETRAINED_GPT2_MODELS +from convert_to_onnx import main as convert_gpt2_to_onnx config: Union[GPT2Config, T5Config] = None @@ -295,11 +298,11 @@ def convert_model(args): # TODO: fix shape inference for T5. Currently symbolic shape inference on T5 is broken. enable_shape_inference = args.model_type == "gpt2" - + if enable_shape_inference: print(f"Run symbolic shape inference on {args.decoder_onnx}. The file will be overwritten.") shape_inference(args.decoder_onnx) - + global config if args.model_type == "gpt2": config = GPT2Config.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir) diff --git a/onnxruntime/python/tools/transformers/io_binding_helper.py b/onnxruntime/python/tools/transformers/io_binding_helper.py index eeab7dd1ce..dbed946192 100644 --- a/onnxruntime/python/tools/transformers/io_binding_helper.py +++ b/onnxruntime/python/tools/transformers/io_binding_helper.py @@ -28,9 +28,10 @@ class TypeHelper: def ort_type_to_numpy_type(ort_type: str): ort_type_to_numpy_type_map = { "tensor(int64)": numpy.longlong, - "tensor(int32)": numpy.int32, #numpy.intc? + "tensor(int32)": numpy.intc, "tensor(float)": numpy.float32, "tensor(float16)": numpy.float16, + "tensor(bool)": numpy.bool, } if ort_type not in ort_type_to_numpy_type_map: raise ValueError(f"{ort_type} not found in map") @@ -44,6 +45,7 @@ class TypeHelper: "tensor(int32)": torch.int32, "tensor(float)": torch.float32, "tensor(float16)": torch.float16, + "tensor(bool)": torch.bool, } if ort_type not in ort_type_to_torch_type_map: raise ValueError(f"{ort_type} not found in map") @@ -54,9 +56,11 @@ class TypeHelper: def numpy_type_to_torch_type(numpy_type: numpy.dtype): numpy_type_to_torch_type_map = { numpy.longlong: torch.int64, + numpy.intc: torch.int32, numpy.int32: torch.int32, numpy.float32: torch.float32, numpy.float16: torch.float16, + numpy.bool: torch.bool, } if numpy_type not in numpy_type_to_torch_type_map: raise ValueError(f"{numpy_type} not found in map") @@ -67,9 +71,10 @@ class TypeHelper: def torch_type_to_numpy_type(torch_type: torch.dtype): torch_type_to_numpy_type_map = { torch.int64: numpy.longlong, - torch.int32: numpy.int32, + torch.int32: numpy.intc, torch.float32: numpy.float32, torch.float16: numpy.float16, + torch.bool: numpy.bool, } if torch_type not in torch_type_to_numpy_type_map: raise ValueError(f"{torch_type} not found in map") diff --git a/onnxruntime/python/tools/transformers/models/__init__.py b/onnxruntime/python/tools/transformers/models/__init__.py index d31e7ad45e..7c2a88f4d9 100644 --- a/onnxruntime/python/tools/transformers/models/__init__.py +++ b/onnxruntime/python/tools/transformers/models/__init__.py @@ -1,8 +1,4 @@ #------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. +# Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -#------------------------------------------------------------------------- -import os -import sys - -sys.path.append(os.path.dirname(__file__)) +#-------------------------------------------------------------------------- diff --git a/onnxruntime/python/tools/transformers/models/gpt2/__init__.py b/onnxruntime/python/tools/transformers/models/gpt2/__init__.py new file mode 100644 index 0000000000..7c2a88f4d9 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/gpt2/__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/benchmark_gpt2.py b/onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py similarity index 91% rename from onnxruntime/python/tools/transformers/benchmark_gpt2.py rename to onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py index b318755d39..f1600017ca 100644 --- a/onnxruntime/python/tools/transformers/benchmark_gpt2.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py @@ -6,20 +6,23 @@ # This script benchmarks gpt2 model with past state. # For gpt2 model without past state, use benchmark.py to measure performance. -import os -import sys -import numpy import csv from datetime import datetime import psutil import argparse import logging import torch -import onnx from packaging import version from transformers import AutoConfig + from gpt2_helper import Gpt2Helper, DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS from gpt2_beamsearch_helper import Gpt2HelperFactory, MODEL_CLASSES + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) + from quantize_helper import QuantizeHelper from benchmark_helper import create_onnxruntime_session, setup_logger, prepare_environment, Precision @@ -227,8 +230,13 @@ def main(args): if args.optimize_onnx or args.precision != Precision.FLOAT32: onnx_model_path = onnx_model_paths[str(args.precision) if args.precision != Precision.INT8 else 'fp32'] - gpt2helper.optimize_onnx(onnx_model_paths["raw"], onnx_model_path, args.precision == Precision.FLOAT16, - model.config.num_attention_heads, model.config.hidden_size, use_external_data_format) + gpt2helper.optimize_onnx(onnx_model_paths["raw"], + onnx_model_path, + args.precision == Precision.FLOAT16, + model.config.num_attention_heads, + model.config.hidden_size, + use_external_data_format, + auto_mixed_precision=True) if args.precision == Precision.INT8: logger.info("quantizing model...") @@ -254,9 +262,15 @@ def main(args): # Allocate output buffers for IO Binding if model_type == 'beam_search_step' or model_type == 'configurable_one_step_search': - max_output_shapes = gpt2helper.get_output_shapes(max(args.batch_sizes), max(args.past_sequence_lengths), - max(args.past_sequence_lengths), max(args.sequence_lengths), 4, - 0, config, args.model_class) + max_output_shapes = gpt2helper.get_output_shapes(max(args.batch_sizes), + context_len=max(args.past_sequence_lengths), + past_sequence_length=max(args.past_sequence_lengths), + sequence_length=max(args.sequence_lengths), + beam_size=args.beam_size, + step=0, + config=config, + model_class=args.model_class) + output_buffers = gpt2helper.get_output_buffers(max_output_shapes, device, args.precision == Precision.FLOAT16) else: @@ -294,8 +308,8 @@ def main(args): has_position_ids=use_padding, has_attention_mask=use_padding) output_shapes = gpt2helper.get_output_shapes(batch_size, past_sequence_length, - past_sequence_length, sequence_length, 4, 0, - config, args.model_class) + past_sequence_length, sequence_length, + args.beam_size, 0, config, args.model_class) else: dummy_inputs = gpt2helper.get_dummy_inputs(batch_size, past_sequence_length, @@ -312,9 +326,18 @@ def main(args): config, args.model_class) try: - outputs, torch_latency = Gpt2Helper.pytorch_inference(model, dummy_inputs, args.test_times) - ort_outputs, ort_latency = Gpt2Helper.onnxruntime_inference(session, dummy_inputs, + outputs, torch_latency = gpt2helper.pytorch_inference(model, dummy_inputs, args.test_times) + + # Dump Torch output shape + for i, value in enumerate(outputs): + if isinstance(value, tuple): + logger.debug(f"torch output {i} is tuple of size {len(value)}, shape {value[0].shape}") + else: + logger.debug(f"torch output {i} shape {value.shape}") + + ort_outputs, ort_latency = gpt2helper.onnxruntime_inference(session, dummy_inputs, args.test_times) + ort_io_outputs, ort_io_latency = gpt2helper.onnxruntime_inference_with_binded_io( session, dummy_inputs, @@ -327,7 +350,7 @@ def main(args): if args.validate_onnx: if gpt2helper.compare_outputs(outputs, ort_outputs, - model_class, + model_class=args.model_class, rtol=DEFAULT_TOLERANCE[args.precision], atol=DEFAULT_TOLERANCE[args.precision]): logger.info( @@ -341,7 +364,7 @@ def main(args): if gpt2helper.compare_outputs(outputs, copy_outputs, - model_class, + model_class=args.model_class, rtol=DEFAULT_TOLERANCE[args.precision], atol=DEFAULT_TOLERANCE[args.precision]): logger.info( @@ -369,6 +392,7 @@ def main(args): csv_writer.writerow(row) except: logger.error(f"Exception", exc_info=True) + return None logger.info(f"Results are saved to file {csv_filename}") return csv_filename diff --git a/onnxruntime/python/tools/transformers/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py similarity index 99% rename from onnxruntime/python/tools/transformers/convert_to_onnx.py rename to onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py index 9d7191202e..88999b0c38 100644 --- a/onnxruntime/python/tools/transformers/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py @@ -15,7 +15,6 @@ This converts GPT2 model to onnx. Examples: """ -import os import argparse import logging import torch @@ -27,6 +26,12 @@ from transformers import AutoConfig from gpt2_helper import DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS from gpt2_beamsearch_helper import Gpt2HelperFactory, MODEL_CLASSES from gpt2_beamsearch_tester import Gpt2TesterFactory + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) + from quantize_helper import QuantizeHelper from benchmark_helper import create_onnxruntime_session, setup_logger, prepare_environment, Precision diff --git a/onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_beamsearch_helper.py similarity index 84% rename from onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py rename to onnxruntime/python/tools/transformers/models/gpt2/gpt2_beamsearch_helper.py index d1cd7e233e..463b2e7829 100644 --- a/onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_beamsearch_helper.py @@ -4,19 +4,21 @@ # license information. # -------------------------------------------------------------------------- # This script helps onnx conversion and validation for GPT2 model with past state. -import os import logging import torch -import onnx import random import numpy import time -import re from pathlib import Path -from typing import List, Dict, Tuple, Union +from typing import List, Dict, Union from transformers import GPT2LMHeadModel, GPT2Config -from benchmark_helper import Precision -from gpt2_helper import Gpt2Helper, Gpt2Inputs, GPT2ModelNoPastState, MyGPT2Model, MyGPT2LMHeadModel, MyGPT2LMHeadModel_NoPadding +from gpt2_helper import Gpt2Helper, Gpt2Inputs, MyGPT2Model, MyGPT2LMHeadModel, MyGPT2LMHeadModel_NoPadding + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) +from io_binding_helper import TypeHelper from torch_onnx_export_helper import torch_onnx_export logger = logging.getLogger(__name__) @@ -25,6 +27,7 @@ BIG_NEG = -1e4 class Gpt2HelperFactory: + @staticmethod def create_helper(helper_type="default"): helpers = { @@ -39,6 +42,7 @@ class Gpt2HelperFactory: class GPT2LMHeadModel_BeamSearchStep(GPT2LMHeadModel): """Here we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state and one step beam search.""" + def __init__(self, config, batch_size, beam_size): super().__init__(config) self.config.batch_size = batch_size @@ -89,7 +93,9 @@ class GPT2LMHeadModel_BeamSearchStep(GPT2LMHeadModel): sorted=True) # output shape=(batch, beam_size) # select the correspondent sentences/next tokens - selected_input_seq = selected_index_flat // self.config.beam_size + selected_input_seq = torch.div(selected_index_flat, self.config.beam_size, + rounding_mode='trunc') # selected_index_flat // self.config.beam_size + next_token_ids = next_token_ids.view(self.config.batch_size, -1).gather(-1, selected_index_flat) prev_step_results = prev_step_results.view(self.config.batch_size, -1, prev_step_results.size(-1)) @@ -98,7 +104,9 @@ class GPT2LMHeadModel_BeamSearchStep(GPT2LMHeadModel): selected_input_seq.unsqueeze(-1).repeat(1, 1, prev_step_results.size(-1))) output_unfinished_sents = input_unfinished_sents.gather(1, selected_input_seq) - output_unfinished_sents = (output_unfinished_sents & next_token_ids.ne(self.config.eos_token_id)) + # Add ones_like to walkaround error like Shape mismatch attempting to re-use buffer. {1,1} != {1,4} + output_unfinished_sents = (output_unfinished_sents & next_token_ids.ne( + torch.ones_like(next_token_ids, dtype=torch.int) * self.config.eos_token_id)) # get the next full input_ids current_step_results = torch.cat([prev_step_results, next_token_ids.unsqueeze(-1)], dim=-1).contiguous() @@ -123,6 +131,7 @@ class GPT2LMHeadModel_BeamSearchStep(GPT2LMHeadModel): class GPT2LMHeadModel_ConfigurableOneStepSearch(GPT2LMHeadModel): """Here we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state and one step beam search with configuration support.""" + def __init__(self, config, batch_size, @@ -190,6 +199,7 @@ class GPT2LMHeadModel_ConfigurableOneStepSearch(GPT2LMHeadModel): if past: last_seq_len = past[0].size(-2) + # input_ids and position_ids contains past sequence input_ids_unfinished_flat = input_ids_unfinished_flat[:, last_seq_len:] position_ids = position_ids[:, last_seq_len:] @@ -264,7 +274,9 @@ class GPT2LMHeadModel_ConfigurableOneStepSearch(GPT2LMHeadModel): sorted=True) # output shape=(batch, beam_size) # select the correspondent sentences/next tokens - selected_input_seq = selected_index_flat // self.config.beam_size + selected_input_seq = torch.div(selected_index_flat, self.config.beam_size, + rounding_mode='trunc') # selected_index_flat // self.config.beam_size + next_token_ids = next_token_ids.view(self.config.batch_size, -1).gather(-1, selected_index_flat) prev_step_results = input_ids.view(self.config.batch_size, -1, input_ids.size(-1)).contiguous() @@ -302,14 +314,13 @@ MODEL_CLASSES = { 'GPT2LMHeadModel': (MyGPT2LMHeadModel, 'logits', True), 'GPT2LMHeadModel_NoPadding': (MyGPT2LMHeadModel_NoPadding, 'logits', False), 'GPT2Model': (MyGPT2Model, 'last_state', True), - "GPT2LMHeadModel_BeamSearchStep": - (GPT2LMHeadModel_BeamSearchStep, "last_state", True), # defined in gpt2_beamsearch_helper.py - "GPT2LMHeadModel_ConfigurableOneStepSearch": - (GPT2LMHeadModel_ConfigurableOneStepSearch, "last_state", False), # defined in gpt2_beamsearch_helper.py + "GPT2LMHeadModel_BeamSearchStep": (GPT2LMHeadModel_BeamSearchStep, "last_state", True), + "GPT2LMHeadModel_ConfigurableOneStepSearch": (GPT2LMHeadModel_ConfigurableOneStepSearch, "last_state", False), } class Gpt2BeamSearchInputs(Gpt2Inputs): + def __init__( self, input_ids, @@ -362,6 +373,7 @@ class Gpt2BeamSearchInputs(Gpt2Inputs): class Gpt2BeamSearchHelper(Gpt2Helper): """A helper class for Gpt2 model conversion, inference and verification.""" + @staticmethod def get_dummy_inputs(batch_size: int, past_sequence_length: int, @@ -373,13 +385,26 @@ class Gpt2BeamSearchHelper(Gpt2Helper): device: torch.device, float16: bool = False, has_position_ids: bool = True, - has_attention_mask: bool = True) -> Gpt2BeamSearchInputs: - """Create random inputs for GPT2 model. - Returns torch tensors of input_ids, position_ids, attention_mask and a list of past state tensors. + has_attention_mask: bool = True, + input_ids_dtype: torch.dtype = torch.int64, + position_ids_dtype: torch.dtype = torch.int64, + attention_mask_dtype: torch.dtype = torch.int64) -> Gpt2BeamSearchInputs: + """Create random inputs for GPT2 beam search. """ - gpt2_dummy_inputs = Gpt2Helper.get_dummy_inputs(batch_size, past_sequence_length, sequence_length, - num_attention_heads, hidden_size, num_layer, vocab_size, device, - float16, has_position_ids, has_attention_mask) + gpt2_dummy_inputs = Gpt2Helper.get_dummy_inputs(batch_size, + past_sequence_length, + sequence_length, + num_attention_heads, + hidden_size, + num_layer, + vocab_size, + device, + float16, + has_position_ids, + has_attention_mask, + input_ids_dtype=input_ids_dtype, + position_ids_dtype=position_ids_dtype, + attention_mask_dtype=attention_mask_dtype) float_type = torch.float16 if float16 else torch.float32 beam_select_idx = torch.zeros([1, batch_size], device=device).long() @@ -418,7 +443,7 @@ class Gpt2BeamSearchHelper(Gpt2Helper): beam_size: int, step: int, config: GPT2Config, - model_class: str = "GPT2LMHeadModel", + model_class: str = "GPT2LMHeadModel_BeamSearchStep", num_seq: int = 0) -> Dict[str, List[int]]: """Returns a dictionary with output name as key, and shape as value.""" num_attention_heads = config.num_attention_heads @@ -431,25 +456,34 @@ class Gpt2BeamSearchHelper(Gpt2Helper): if model_class == "GPT2LMHeadModel_BeamSearchStep": last_state_shape = [batch_size, beam_size] else: - last_state_shape = [batch_size * beam_size, past_sequence_length + sequence_length + 1] + last_state_shape = [batch_size * beam_size, past_sequence_length - context_len + sequence_length + 1] - if step == 0: + if model_class == "GPT2LMHeadModel_BeamSearchStep": + if step == 0: + present_state_shape = [ + 2, + batch_size, + num_attention_heads, + past_sequence_length + sequence_length, + int(hidden_size / num_attention_heads), + ] + else: + if num_seq == 0: + num_seq = beam_size + + present_state_shape = [ + 2, + batch_size * num_seq, + num_attention_heads, + past_sequence_length + sequence_length, + int(hidden_size / num_attention_heads), + ] + else: present_state_shape = [ 2, batch_size, num_attention_heads, - past_sequence_length + sequence_length, - int(hidden_size / num_attention_heads), - ] - else: - if num_seq == 0: - num_seq = beam_size - - present_state_shape = [ - 2, - batch_size * num_seq, - num_attention_heads, - past_sequence_length + sequence_length, + past_sequence_length - context_len + sequence_length, int(hidden_size / num_attention_heads), ] @@ -457,14 +491,16 @@ class Gpt2BeamSearchHelper(Gpt2Helper): for i in range(num_layer): output_shapes["present_" + str(i)] = present_state_shape + # TODO: reshape output_selected_indices as [batch_size, beam_size] output_shapes["output_selected_indices"] = [1, batch_size * beam_size] output_shapes["output_log_probs"] = [batch_size, beam_size] output_shapes["output_unfinished_sents"] = [batch_size, beam_size] if model_class == "GPT2LMHeadModel_BeamSearchStep": - output_shapes["current_step_results"] = [batch_size * beam_size, past_sequence_length + sequence_length + 1] - output_shapes["current_step_scores"] = [ - batch_size * beam_size, past_sequence_length + sequence_length - context_len + 2 - ] + output_shapes["current_step_results"] = [ + batch_size * beam_size, past_sequence_length - context_len + sequence_length + 1 + ] + output_shapes["current_step_scores"] = [batch_size * beam_size, past_sequence_length - context_len + 2] + print("output_shapes", output_shapes) return output_shapes @staticmethod @@ -528,6 +564,8 @@ class Gpt2BeamSearchHelper(Gpt2Helper): has_position_ids: bool = True, has_attention_mask: bool = True): """Export GPT-2 model with past state to ONNX model.""" + assert isinstance(model, (GPT2LMHeadModel_BeamSearchStep, GPT2LMHeadModel_ConfigurableOneStepSearch)) + config: GPT2Config = model.config num_layer = config.n_layer dummy_inputs = Gpt2BeamSearchHelper.get_dummy_inputs(batch_size=1, @@ -544,7 +582,6 @@ class Gpt2BeamSearchHelper(Gpt2Helper): input_list = dummy_inputs.to_list() with torch.no_grad(): - # outputs = model(input_ids, position_id, attention_mask, beam_select_idx, past) outputs = model(*input_list) past_names = [f"past_{i}" for i in range(num_layer)] @@ -568,14 +605,6 @@ class Gpt2BeamSearchHelper(Gpt2Helper): "current_step_scores", ] - # Shape of input tensors: - # input_ids: (batch_size, seq_len) - # past_{i}: (2, batch_size, num_heads, past_seq_len, hidden_size/num_heads) - # attention_mask: (batch_size, past_seq_len + seq_len) - # Shape of output tensors: - # last_state: (batch_size, seq_len, hidden_size) - # or logits: (batch_size, seq_len, vocab_size) - # present_{i}: (2, batch_size, num_heads, past_seq_len + seq_len, hidden_size/num_heads) dynamic_axes = { "input_ids": { 0: "batch_size", @@ -613,15 +642,23 @@ class Gpt2BeamSearchHelper(Gpt2Helper): # add dynamic output axes present_axes = {1: 'batch_size', 3: 'cur_seq_len'} - dynamic_axes["last_state"] = {0: 'batch_size', 1: 'beam_size'} + + if isinstance(model, GPT2LMHeadModel_BeamSearchStep): + dynamic_axes["last_state"] = {0: "batch_size", 1: "beam_size"} + else: + dynamic_axes["last_state"] = {0: "batch_size * beam_size", 1: "total_seq_len"} + for i in range(num_layer): dynamic_axes["present_" + str(i)] = present_axes - dynamic_axes["output_selected_indices"] = {0: "batch_size", 1: "'beam_size_or_1'"} - dynamic_axes["output_log_probs"] = {0: "batch_size", 1: "'beam_size'"} - dynamic_axes["output_unfinished_sents"] = {0: "batch_size", 1: "'beam_size'"} - dynamic_axes["current_step_results"] = {0: "beam_size_or_1", 1: "total_seq_len"} - dynamic_axes["current_step_scores"] = {0: "beam_size_or_1", 1: "total_seq_len"} + dynamic_axes["output_selected_indices"] = {1: "batch_size * beam_size"} + dynamic_axes["output_log_probs"] = {0: "batch_size", 1: "beam_size"} + dynamic_axes["output_unfinished_sents"] = {0: "batch_size", 1: "beam_size"} + + if "current_step_results" in output_names: + dynamic_axes["current_step_results"] = {0: "batch_size * beam_size", 1: "total_seq_len"} + + dynamic_axes["current_step_scores"] = {0: "batch_size * beam_size"} logger.info( f"Shapes: input_ids={dummy_inputs.input_ids.shape} past={dummy_inputs.past[0].shape} output={outputs[0].shape} present={outputs[1][0].shape}" @@ -636,7 +673,7 @@ class Gpt2BeamSearchHelper(Gpt2Helper): input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, - opset_version=12, + opset_version=14, do_constant_folding=True, use_external_data_format=use_external_data_format, verbose=verbose, @@ -697,7 +734,7 @@ class Gpt2BeamSearchHelper(Gpt2Helper): prev_step_scores=None): """Returnas IO binding object for a session.""" - # Bind inputs and outputs to onnxruntime session + # Bind (input_ids, position_ids, attention_mask and past_*) and all outputs io_binding = Gpt2Helper.prepare_io_binding(ort_session, input_ids, position_ids, @@ -706,114 +743,25 @@ class Gpt2BeamSearchHelper(Gpt2Helper): output_buffers=output_buffers, output_shapes=output_shapes) - # Bind inputs - data_type = output_buffers[ort_session.get_outputs()[1].name].dtype - float_type = numpy.float16 if data_type == torch.float16 else numpy.float32 - - if past is not None: - for i, past_i in enumerate(past): - assert past_i.is_contiguous() - - data_ptr = past_i.data_ptr() - if data_ptr == 0: - # When past_sequence_length is 0, its data_ptr will be zero. IO Binding asserts that data_ptr shall not be zero. - # Here we workaround and pass data pointer of input_ids. Actual data is not used for past so it does not matter. - data_ptr = input_ids.data_ptr() - - io_binding.bind_input(f'past_{i}', past_i.device.type, 0, float_type, list(past_i.size()), data_ptr) - - if attention_mask is not None: - assert attention_mask.is_contiguous() - io_binding.bind_input('attention_mask', attention_mask.device.type, 0, float_type, - list(attention_mask.size()), attention_mask.data_ptr()) - - if beam_select_idx is not None: - assert beam_select_idx.is_contiguous() - io_binding.bind_input( - "beam_select_idx", - beam_select_idx.device.type, - 0, - numpy.longlong, - list(beam_select_idx.size()), - beam_select_idx.data_ptr(), - ) - - if input_log_probs is not None: - assert input_log_probs.is_contiguous() - io_binding.bind_input( - "input_log_probs", - input_log_probs.device.type, - 0, - float_type, - list(input_log_probs.size()), - input_log_probs.data_ptr(), - ) - - if input_unfinished_sents is not None: - assert input_unfinished_sents.is_contiguous() - io_binding.bind_input( - "input_unfinished_sents", - input_unfinished_sents.device.type, - 0, - numpy.bool, - list(input_unfinished_sents.size()), - input_unfinished_sents.data_ptr(), - ) - - if prev_step_results is not None: - assert prev_step_results.is_contiguous() - io_binding.bind_input( - "prev_step_results", - prev_step_results.device.type, - 0, - numpy.longlong, - list(prev_step_results.size()), - prev_step_results.data_ptr(), - ) - - if prev_step_scores is not None: - assert prev_step_scores.is_contiguous() - io_binding.bind_input( - "prev_step_scores", - prev_step_scores.device.type, - 0, - float_type, - list(prev_step_scores.size()), - prev_step_scores.data_ptr(), - ) - - # Bind outputs - for output in ort_session.get_outputs(): - output_name = output.name - output_buffer = output_buffers[output_name] - logger.debug(f"{output_name} device type={output_buffer.device.type} shape={list(output_buffer.size())}") - if (output_name == "output_selected_indices" or output_name == "last_state" - or output_name == "current_step_results"): - io_binding.bind_output( - output_name, - output_buffer.device.type, + # Bind the remaining inputs + other_inputs = { + "beam_select_idx": beam_select_idx, + "input_log_probs": input_log_probs, + "input_unfinished_sents": input_unfinished_sents, + "prev_step_results": prev_step_results, + "prev_step_scores": prev_step_scores, + } + name_to_np_type = TypeHelper.get_io_numpy_type_map(ort_session) + for name, tensor in other_inputs.items(): + if tensor is not None: + assert tensor.is_contiguous() + io_binding.bind_input( + name, + tensor.device.type, 0, - numpy.longlong, - output_shapes[output_name], - output_buffer.data_ptr(), - ) - elif output_name == "output_unfinished_sents": - io_binding.bind_output( - output_name, - output_buffer.device.type, - 0, - numpy.bool, - output_shapes[output_name], - output_buffer.data_ptr(), - ) - else: - io_binding.bind_output( - output_name, - output_buffer.device.type, - 0, - float_type, - output_shapes[output_name], - output_buffer.data_ptr(), + name_to_np_type[name], + list(tensor.size()), + tensor.data_ptr(), ) return io_binding diff --git a/onnxruntime/python/tools/transformers/gpt2_beamsearch_tester.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_beamsearch_tester.py similarity index 99% rename from onnxruntime/python/tools/transformers/gpt2_beamsearch_tester.py rename to onnxruntime/python/tools/transformers/models/gpt2/gpt2_beamsearch_tester.py index a8f7047bde..540b96189c 100644 --- a/onnxruntime/python/tools/transformers/gpt2_beamsearch_tester.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_beamsearch_tester.py @@ -4,24 +4,25 @@ # license information. # -------------------------------------------------------------------------- # This script helps evaluation of GPT-2 model. -import os import logging import torch -import random import numpy -import time import timeit -import math -import statistics -from pathlib import Path from gpt2_tester import Gpt2Tester, Gpt2Metric from gpt2_beamsearch_helper import Gpt2BeamSearchHelper, Gpt2BeamSearchInputs + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) + from benchmark_helper import Precision logger = logging.getLogger(__name__) class Gpt2TesterFactory: + @staticmethod def create_tester(tester_type="default"): testers = { @@ -34,6 +35,7 @@ class Gpt2TesterFactory: class Gpt2BeamSearchTester(Gpt2Tester): + def __init__( self, input_ids, diff --git a/onnxruntime/python/tools/transformers/gpt2_helper.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py similarity index 98% rename from onnxruntime/python/tools/transformers/gpt2_helper.py rename to onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py index ff619082ba..a7a42d7da8 100644 --- a/onnxruntime/python/tools/transformers/gpt2_helper.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py @@ -4,18 +4,22 @@ # license information. # -------------------------------------------------------------------------- # This script helps onnx conversion and validation for GPT2 model with past state. -import os import logging import torch import shutil import random import numpy import time -import re import pickle from pathlib import Path from typing import List, Dict, Tuple, Union from transformers import GPT2Model, GPT2LMHeadModel, GPT2Config, TFGPT2Model + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) + from float16 import float_to_float16_max_diff from onnx_model import OnnxModel from fusion_utils import FusionUtils @@ -261,8 +265,9 @@ class Gpt2Helper: return numpy.amax(diff) @staticmethod - def compare_outputs(torch_outputs, ort_outputs, rtol=1e-03, atol=1e-03): + def compare_outputs(torch_outputs, ort_outputs, rtol=1e-03, atol=1e-03, **kwargs): """ Returns True if torch and ORT outputs are close for given thresholds, and False otherwise. + Note: need kwargs since Gpt2BeamSearchHelper.compare_outputs has an extra parameter model_class """ is_close = numpy.allclose(ort_outputs[0], torch_outputs[0].cpu().numpy(), rtol=rtol, atol=atol) logger.debug(f'PyTorch and OnnxRuntime output 0 (last_state) are close: {is_close}') @@ -444,6 +449,8 @@ class Gpt2Helper: if auto_mixed_precision: Gpt2Helper.auto_mixed_precision(m) else: + if "keep_io_types" not in kwargs: + kwargs["keep_io_types"] = False m.convert_float_to_float16(use_symbolic_shape_infer=True, **kwargs) m.save_model_to_file(optimized_model_path, use_external_data_format) @@ -490,10 +497,9 @@ class Gpt2Helper: else: logger.warning(f"Failed to find MatMul node for logits. Found {node.op_type} of node {node.name}") - if is_weight_fp16_precision: - keep_io_types = [] - node_block_list = [] - else: + 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] diff --git a/onnxruntime/python/tools/transformers/gpt2_parity.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_parity.py similarity index 99% rename from onnxruntime/python/tools/transformers/gpt2_parity.py rename to onnxruntime/python/tools/transformers/models/gpt2/gpt2_parity.py index d74544772c..7147e36d6e 100644 --- a/onnxruntime/python/tools/transformers/gpt2_parity.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_parity.py @@ -10,12 +10,8 @@ # pvalue < 0.05 means two experiments have significant difference on top 1 match rate. # User could use this script to select the best mixed precision model according to these metrics. -from convert_to_onnx import main, get_latency_name -import os import argparse import logging -from gpt2_helper import PRETRAINED_GPT2_MODELS, Gpt2Helper -from benchmark_helper import setup_logger from onnx_model import OnnxModel import onnx import csv @@ -23,6 +19,16 @@ import datetime import scipy.stats import torch +from gpt2_helper import PRETRAINED_GPT2_MODELS, Gpt2Helper +from convert_to_onnx import main, get_latency_name + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) + +from benchmark_helper import setup_logger + logger = logging.getLogger('') @@ -69,6 +75,7 @@ def parse_arguments(argv=None): class ParityTask: + def __init__(self, test_cases, total_runs, csv_path): self.total_runs = total_runs self.test_cases = test_cases diff --git a/onnxruntime/python/tools/transformers/gpt2_tester.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_tester.py similarity index 99% rename from onnxruntime/python/tools/transformers/gpt2_tester.py rename to onnxruntime/python/tools/transformers/models/gpt2/gpt2_tester.py index 5f65a0a1c2..b2b7532786 100644 --- a/onnxruntime/python/tools/transformers/gpt2_tester.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_tester.py @@ -4,22 +4,26 @@ # license information. # -------------------------------------------------------------------------- # This script helps evaluation of GPT-2 model. -import os import logging import torch -import random import numpy -import time import timeit import math import statistics from gpt2_helper import Gpt2Helper, Gpt2Inputs + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) + from benchmark_helper import Precision logger = logging.getLogger(__name__) class Gpt2Metric: + def __init__(self, treatment_name, baseline_name='Torch', top_k=20): assert top_k > 1 and top_k <= 100 self.baseline = baseline_name @@ -112,6 +116,7 @@ class Gpt2Metric: class Gpt2Tester: + def __init__(self, input_ids, position_ids, diff --git a/onnxruntime/python/tools/transformers/parity_check_helper.py b/onnxruntime/python/tools/transformers/models/gpt2/parity_check_helper.py similarity index 98% rename from onnxruntime/python/tools/transformers/parity_check_helper.py rename to onnxruntime/python/tools/transformers/models/gpt2/parity_check_helper.py index a444478a80..2f878fd3ce 100644 --- a/onnxruntime/python/tools/transformers/parity_check_helper.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/parity_check_helper.py @@ -14,6 +14,12 @@ import torch from pathlib import Path from onnx import numpy_helper, TensorProto from gpt2_helper import Gpt2Helper + +import sys +import os + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) + from benchmark_helper import create_onnxruntime_session NON_ZERO_VALUE = str(1) diff --git a/onnxruntime/python/tools/transformers/models/longformer/__init__.py b/onnxruntime/python/tools/transformers/models/longformer/__init__.py index 864a0f3b5b..7c2a88f4d9 100644 --- a/onnxruntime/python/tools/transformers/models/longformer/__init__.py +++ b/onnxruntime/python/tools/transformers/models/longformer/__init__.py @@ -1,10 +1,4 @@ -# ------------------------------------------------------------------------- +#------------------------------------------------------------------------- # 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 - -sys.path.append(os.path.dirname(__file__)) +# Licensed under the MIT License. +#-------------------------------------------------------------------------- diff --git a/onnxruntime/python/tools/transformers/models/t5/__init__.py b/onnxruntime/python/tools/transformers/models/t5/__init__.py index d31e7ad45e..7c2a88f4d9 100644 --- a/onnxruntime/python/tools/transformers/models/t5/__init__.py +++ b/onnxruntime/python/tools/transformers/models/t5/__init__.py @@ -1,8 +1,4 @@ #------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. +# Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -#------------------------------------------------------------------------- -import os -import sys - -sys.path.append(os.path.dirname(__file__)) +#-------------------------------------------------------------------------- diff --git a/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2-OneStepSearch_OnnxRuntime_CPU.ipynb b/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2-OneStepSearch_OnnxRuntime_CPU.ipynb index 92c3dbb47b..a108024b2b 100644 --- a/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2-OneStepSearch_OnnxRuntime_CPU.ipynb +++ b/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2-OneStepSearch_OnnxRuntime_CPU.ipynb @@ -128,7 +128,13 @@ } ], "source": [ - "from onnxruntime.transformers.gpt2_beamsearch_helper import Gpt2BeamSearchHelper, GPT2LMHeadModel_BeamSearchStep\n", + "from packaging import version\n", + "from onnxruntime import __version__ as ort_verison\n", + "if version.parse(ort_verison) >= version.parse('1.12.0'):\n", + " from onnxruntime.transformers.models.gpt2.gpt2_beamsearch_helper import Gpt2BeamSearchHelper, GPT2LMHeadModel_BeamSearchStep\n", + "else:\n", + " from onnxruntime.transformers.gpt2_beamsearch_helper import Gpt2BeamSearchHelper, GPT2LMHeadModel_BeamSearchStep\n", + "\n", "from transformers import AutoConfig\n", "import torch\n", "\n", diff --git a/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2_with_OnnxRuntime_on_CPU.ipynb b/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2_with_OnnxRuntime_on_CPU.ipynb index 6428b2fed1..62c5f079ce 100644 --- a/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2_with_OnnxRuntime_on_CPU.ipynb +++ b/onnxruntime/python/tools/transformers/notebooks/Inference_GPT2_with_OnnxRuntime_on_CPU.ipynb @@ -183,7 +183,13 @@ } ], "source": [ - "from onnxruntime.transformers.gpt2_helper import Gpt2Helper, MyGPT2LMHeadModel\n", + "from packaging import version\n", + "from onnxruntime import __version__ as ort_verison\n", + "if version.parse(ort_verison) >= version.parse('1.12.0'):\n", + " from onnxruntime.transformers.models.gpt2.gpt2_helper import Gpt2Helper, MyGPT2LMHeadModel\n", + "else:\n", + " from onnxruntime.transformers.gpt2_helper import Gpt2Helper, MyGPT2LMHeadModel\n", + "\n", "from transformers import AutoConfig\n", "import torch\n", "\n", diff --git a/onnxruntime/test/python/transformers/parity_utilities.py b/onnxruntime/test/python/transformers/parity_utilities.py index a4cc62bdb4..d450700581 100644 --- a/onnxruntime/test/python/transformers/parity_utilities.py +++ b/onnxruntime/test/python/transformers/parity_utilities.py @@ -10,13 +10,13 @@ import numpy import torch -def find_transformers_source(): - source_dir = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'python', 'tools', 'transformers') +def find_transformers_source(sub_dir_paths=[]): + source_dir = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'python', 'tools', 'transformers', *sub_dir_paths) if (os.path.exists(source_dir)): if source_dir not in sys.path: sys.path.append(source_dir) return True - return False + return False def create_inputs(batch_size=1, sequence_length=1, hidden_size=768, float16=False, device=torch.device('cuda')): diff --git a/onnxruntime/test/python/transformers/test_gpt2.py b/onnxruntime/test/python/transformers/test_gpt2.py index 2264a97d4b..449393c5e9 100644 --- a/onnxruntime/test/python/transformers/test_gpt2.py +++ b/onnxruntime/test/python/transformers/test_gpt2.py @@ -13,13 +13,15 @@ import coloredlogs import pytest from parity_utilities import find_transformers_source -if find_transformers_source(): + +if find_transformers_source(sub_dir_paths=['models', 'gpt2']): from benchmark_gpt2 import parse_arguments, main else: - from onnxruntime.transformers.benchmark_gpt2 import parse_arguments, main + from onnxruntime.transformers.models.gpt2.benchmark_gpt2 import parse_arguments, main class TestGpt2(unittest.TestCase): + def setUp(self): from onnxruntime import get_available_providers self.test_cuda = 'CUDAExecutionProvider' in get_available_providers() @@ -27,51 +29,56 @@ class TestGpt2(unittest.TestCase): def run_benchmark_gpt2(self, arguments: str): args = parse_arguments(arguments.split()) csv_filename = main(args) + self.assertIsNotNone(csv_filename) self.assertTrue(os.path.exists(csv_filename)) @pytest.mark.slow def test_gpt2_fp32(self): - self.run_benchmark_gpt2('-m gpt2 --precision fp32 -v -b 1 -s 128') + self.run_benchmark_gpt2('-m gpt2 --precision fp32 -v -b 1 --sequence_lengths 2 -s 3') @pytest.mark.slow def test_gpt2_fp16(self): if self.test_cuda: - self.run_benchmark_gpt2('-m gpt2 --precision fp16 -o -b 1 -s 128 --use_gpu') + self.run_benchmark_gpt2('-m gpt2 --precision fp16 -o -b 1 --sequence_lengths 2 -s 3 --use_gpu') @pytest.mark.slow def test_gpt2_int8(self): - self.run_benchmark_gpt2('-m gpt2 --precision int8 -o -b 1 -s 128') + self.run_benchmark_gpt2('-m gpt2 --precision int8 -o -b 1 --sequence_lengths 2 -s 3') @pytest.mark.slow def test_gpt2_beam_search_step_fp32(self): - self.run_benchmark_gpt2('-m gpt2 --model_class=GPT2LMHeadModel_BeamSearchStep --precision fp32 -v -b 1 -s 128') + self.run_benchmark_gpt2( + '-m gpt2 --model_class=GPT2LMHeadModel_BeamSearchStep --precision fp32 -v -b 1 --sequence_lengths 5 -s 3') - @pytest.mark.slow - def test_gpt2_beam_search_step_fp16(self): - if self.test_cuda: - self.run_benchmark_gpt2( - '-m gpt2 --model_class=GPT2LMHeadModel_BeamSearchStep --precision fp16 -o -b 1 -s 128 --use_gpu') + # @pytest.mark.slow + # def test_gpt2_beam_search_step_fp16(self): + # if self.test_cuda: + # self.run_benchmark_gpt2( + # '-m gpt2 --model_class=GPT2LMHeadModel_BeamSearchStep --precision fp16 -o -b 1 --sequence_lengths 5 -s 3 --use_gpu') @pytest.mark.slow def test_gpt2_beam_search_step_int8(self): - self.run_benchmark_gpt2('-m gpt2 --model_class=GPT2LMHeadModel_BeamSearchStep --precision int8 -o -b 1 -s 128') + self.run_benchmark_gpt2( + '-m gpt2 --model_class=GPT2LMHeadModel_BeamSearchStep --precision int8 -o -b 1 --sequence_lengths 5 -s 3') @pytest.mark.slow def test_gpt2_configurable_one_step_search_fp32(self): self.run_benchmark_gpt2( - '-m gpt2 --model_class=GPT2LMHeadModel_ConfigurableOneStepSearch --precision fp32 -v -b 1 -s 128') + '-m gpt2 --model_class=GPT2LMHeadModel_ConfigurableOneStepSearch --precision fp32 -v -b 1 --sequence_lengths 5 --past_sequence_lengths 3 --use_gpu' + ) @pytest.mark.slow def test_gpt2_configurable_one_step_search_fp16(self): if self.test_cuda: self.run_benchmark_gpt2( - '-m gpt2 --model_class=GPT2LMHeadModel_ConfigurableOneStepSearch --precision fp16 -o -b 1 -s 128 --use_gpu' + '-m gpt2 --model_class=GPT2LMHeadModel_ConfigurableOneStepSearch --precision fp16 -o -b 1 --sequence_lengths 5 -s 3 --use_gpu' ) @pytest.mark.slow def test_gpt2_configurable_one_step_search_int8(self): self.run_benchmark_gpt2( - '-m gpt2 --model_class=GPT2LMHeadModel_ConfigurableOneStepSearch --precision int8 -o -b 1 -s 128') + '-m gpt2 --model_class=GPT2LMHeadModel_ConfigurableOneStepSearch --precision int8 -o -b 1 --sequence_lengths 5 -s 3' + ) if __name__ == '__main__': diff --git a/setup.py b/setup.py index a9749a9317..262b7e27ef 100644 --- a/setup.py +++ b/setup.py @@ -321,8 +321,9 @@ packages = [ 'onnxruntime.quantization.operators', 'onnxruntime.quantization.CalTableFlatBuffers', 'onnxruntime.transformers', - 'onnxruntime.transformers.models.t5', + 'onnxruntime.transformers.models.gpt2', 'onnxruntime.transformers.models.longformer', + 'onnxruntime.transformers.models.t5', ] requirements_file = "requirements.txt"