diff --git a/onnxruntime/python/tools/transformers/benchmark_helper.py b/onnxruntime/python/tools/transformers/benchmark_helper.py index 8d22e0862a..64efeb22f4 100644 --- a/onnxruntime/python/tools/transformers/benchmark_helper.py +++ b/onnxruntime/python/tools/transformers/benchmark_helper.py @@ -537,3 +537,17 @@ def measure_memory(is_gpu, func): print(f"CPU memory usage: before={memory_before_test:.1f} MB, peak={max_usage:.1f} MB") return max_usage - memory_before_test + + +def get_ort_environment_variables(): + # Environment variables might impact ORT performance on transformer models. Note that they are for testing only. + env_names = ["ORT_DISABLE_FUSED_ATTENTION", "ORT_TRANSFORMER_OPTIONS", "ORT_CUDA_GEMM_OPTIONS"] + env = "" + for name in env_names: + value = os.getenv(name) + if value is None: + continue + if env: + env += "," + env += f"{name}={value}" + return env diff --git a/onnxruntime/python/tools/transformers/convert_generation.py b/onnxruntime/python/tools/transformers/convert_generation.py index b45d5e63f2..0524b9baf2 100644 --- a/onnxruntime/python/tools/transformers/convert_generation.py +++ b/onnxruntime/python/tools/transformers/convert_generation.py @@ -367,7 +367,6 @@ def gpt2_to_onnx(args: argparse.Namespace): "1", "--test_cases", "10", - "--use_int32_inputs", # BeamSearch requires to use int32 for input_ids, position_ids and attention_mask "--overwrite", # Overwrite onnx file if existed ] if args.use_gpu: diff --git a/onnxruntime/python/tools/transformers/fusion_embedlayer.py b/onnxruntime/python/tools/transformers/fusion_embedlayer.py index 0be0f81c41..e0cfc15a33 100644 --- a/onnxruntime/python/tools/transformers/fusion_embedlayer.py +++ b/onnxruntime/python/tools/transformers/fusion_embedlayer.py @@ -736,6 +736,12 @@ class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): mask_int32 = self.attention.input[3] children_nodes = input_name_to_nodes[mask_int32] + if self.model.find_graph_input(mask_int32): + attention_nodes = [node for node in children_nodes if node.op_type == "Attention"] + self.replace_mask(mask_int32, attention_nodes) + self.increase_counter("EmbedLayerNormalization(with mask)") + return + if mask_int32 not in output_name_to_node: logger.debug("EmbedLayerNormalization will not have mask since %s is not a node output", mask_int32) self.increase_counter("EmbedLayerNormalization(no mask)") diff --git a/onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py b/onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py index f8cbdaaa32..4041b285ea 100644 --- a/onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/benchmark_gpt2.py @@ -22,7 +22,13 @@ from transformers import AutoConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) -from benchmark_helper import Precision, create_onnxruntime_session, prepare_environment, setup_logger +from benchmark_helper import ( + Precision, + create_onnxruntime_session, + get_ort_environment_variables, + prepare_environment, + setup_logger, +) from quantize_helper import QuantizeHelper logger = logging.getLogger("") @@ -89,6 +95,19 @@ def parse_arguments(argv=None): ) parser.set_defaults(optimize_onnx=False) + parser.add_argument( + "--stage", + type=int, + default=0, + required=False, + choices=[0, 1, 2], + help="Stage in generation: 1 (initial decoder), 2 (decoder), 0 (both). " + "1 - decode the first token when past_sequence_length is zero; " + "2 - decode the remaining tokens when past_sequence_length is not zero; " + "0 - one onnx model for both stages 1 and 2. " + "Note that we will optimize 1 and 2 differently for best performance.", + ) + parser.add_argument("--use_gpu", required=False, action="store_true", help="use GPU for inference") parser.set_defaults(use_gpu=False) @@ -145,6 +164,12 @@ def parse_arguments(argv=None): parser.add_argument("--verbose", required=False, action="store_true") parser.set_defaults(verbose=False) + parser.add_argument("--output_torch_latency", required=False, action="store_true") + parser.set_defaults(output_torch_latency=False) + + parser.add_argument("--disable_io_binding", required=False, action="store_true") + parser.set_defaults(disable_io_binding=False) + search_option_group = parser.add_argument_group("configurable one step search options") search_option_group.add_argument( @@ -213,6 +238,9 @@ def main(args): if args.precision == Precision.INT8: assert not args.use_gpu, "quantization only supports CPU" + if args.stage == 1: + assert args.past_sequence_lengths == [0], "past_sequence_lengths shall be 0 for stage==1 (init decoder)" + torch.set_num_threads(psutil.cpu_count(logical=True) if args.thread_num <= 0 else args.thread_num) print(torch.__config__.parallel_info()) @@ -294,6 +322,7 @@ def main(args): model.config.hidden_size, use_external_data_format, auto_mixed_precision=True, + stage=args.stage, ) if args.precision == Precision.INT8: @@ -352,6 +381,8 @@ def main(args): column_names = [ "model_name", "model_class", + "stage", + "environment_variables", "gpu", "precision", "optimizer", @@ -359,9 +390,9 @@ def main(args): "batch_size", "sequence_length", "past_sequence_length", + "disable_io_binding", "torch_latency", "onnxruntime_latency", - "onnxruntime_io_binding_latency", ] csv_writer = csv.DictWriter(csv_file, fieldnames=column_names) csv_writer.writeheader() @@ -420,45 +451,43 @@ def main(args): ) try: - outputs, torch_latency = gpt2helper.pytorch_inference(model, dummy_inputs, args.test_times) + if args.validate_onnx or args.output_torch_latency: + 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}") + # 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}") + else: + outputs = None + torch_latency = None - 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, - output_buffers, - output_shapes, - args.test_times, - return_numpy=False, - include_copy_output_latency=args.include_copy_output_latency, - ) + if args.disable_io_binding: + ort_outputs, ort_latency = gpt2helper.onnxruntime_inference( + session, dummy_inputs, args.test_times + ) + else: + ort_outputs, ort_latency = gpt2helper.onnxruntime_inference_with_binded_io( + session, + dummy_inputs, + output_buffers, + output_shapes, + args.test_times, + return_numpy=False, + include_copy_output_latency=args.include_copy_output_latency, + ) if args.validate_onnx: - if gpt2helper.compare_outputs( - outputs, - ort_outputs, - model_class=args.model_class, - rtol=DEFAULT_TOLERANCE[args.precision], - atol=DEFAULT_TOLERANCE[args.precision], - ): - logger.info( - f"Pytorch and ONNX Runtime outputs are all close (tolerance={DEFAULT_TOLERANCE[args.precision]})." - ) - - # Results of IO binding might be in GPU. Copy outputs to CPU for comparison. - copy_outputs = [] - for output in ort_io_outputs: - copy_outputs.append(output.cpu().numpy()) + copy_outputs = ort_outputs + if not args.disable_io_binding: + # Results of IO binding might be in GPU. Copy outputs to CPU for comparison. + copy_outputs = [] + for output in ort_outputs: + copy_outputs.append(output.cpu().numpy()) if gpt2helper.compare_outputs( outputs, @@ -468,16 +497,24 @@ def main(args): atol=DEFAULT_TOLERANCE[args.precision], ): logger.info( - f"Pytorch and ONNX Runtime IO Binding outputs are all close (tolerance={DEFAULT_TOLERANCE[args.precision]})." + f"Pytorch and ONNX Runtime outputs are all close (tolerance={DEFAULT_TOLERANCE[args.precision]})." ) logger.info( - f"batch_size={batch_size}, sequence_length={sequence_length}, past_sequence_length={past_sequence_length}, torch_latency={torch_latency:.2f}, onnxruntime_latency={ort_latency:.2f}, onnxruntime_io_binding_latency={ort_io_latency:.2f}" + "batch_size=%d, sequence_length=%d, past_sequence_length=%d, onnxruntime_latency=%.2f %s %s", + batch_size, + sequence_length, + past_sequence_length, + ort_latency, + "(disable_io_binding)" if args.disable_io_binding else "", + ", torch_latency={torch_latency}" if torch_latency else "", ) row = { "model_name": args.model_name_or_path, "model_class": args.model_class, + "stage": args.stage, + "environment_variables": get_ort_environment_variables(), "gpu": args.use_gpu, "precision": args.precision, "optimizer": args.optimize_onnx, @@ -485,9 +522,9 @@ def main(args): "batch_size": batch_size, "sequence_length": sequence_length, "past_sequence_length": past_sequence_length, - "torch_latency": f"{torch_latency:.2f}", + "disable_io_binding": args.disable_io_binding, + "torch_latency": f"{torch_latency:.2f}" if torch_latency else "None", "onnxruntime_latency": f"{ort_latency:.2f}", - "onnxruntime_io_binding_latency": f"{ort_io_latency:.2f}", } csv_writer.writerow(row) except: diff --git a/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py index 8087f97f1a..6a03a091ed 100644 --- a/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py @@ -31,7 +31,13 @@ from transformers import AutoConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) -from benchmark_helper import Precision, create_onnxruntime_session, prepare_environment, setup_logger +from benchmark_helper import ( + Precision, + create_onnxruntime_session, + get_ort_environment_variables, + prepare_environment, + setup_logger, +) from quantize_helper import QuantizeHelper logger = logging.getLogger("") @@ -147,12 +153,26 @@ def parse_arguments(argv=None): parser.set_defaults(overwrite=False) parser.add_argument( - "--use_int32_inputs", + "--use_int64_inputs", required=False, action="store_true", help="Use int32 instead of int64 for input_ids, position_ids and attention_mask.", ) - parser.set_defaults(use_int32_inputs=False) + parser.set_defaults(use_int64_inputs=False) + + parser.add_argument( + "-s", + "--stage", + type=int, + default=0, + required=False, + choices=[0, 1, 2], + help="Stage in generation: 1 (initial decoder), 2 (decoder), 0 (both). " + "1 - decode the first token when past_sequence_length is zero; " + "2 - decode the remaining tokens when past_sequence_length is not zero; " + "0 - one onnx model for both stages 1 and 2. " + "Note that we will optimize 1 and 2 differently for best performance.", + ) parser.add_argument( "--beam_size", @@ -241,7 +261,8 @@ def parse_arguments(argv=None): "--op_block_list", nargs="+", default=[], - help="List of operators (like Attention Gather Add LayerNormalization FastGelu MatMul) to compute in float32 instead of float16.", + help="List of operators (like Add LayerNormalization SkipLayerNormalization EmbedLayerNormalization FastGelu) " + "to compute in float32 instead of float16.", ) fp16_option_group.add_argument( @@ -271,11 +292,11 @@ def get_onnx_model_size(onnx_path: str, use_external_data_format: bool): return sum([f.stat().st_size for f in Path(onnx_path).parent.rglob("*")]) -def get_latency_name(): - return "average_latency(batch_size=8,sequence_length=1,past_sequence_length=32)" +def get_latency_name(batch_size, sequence_length, past_sequence_length): + return f"average_latency(batch_size={batch_size},sequence_length={sequence_length},past_sequence_length={past_sequence_length})" -def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_results.csv"): +def main(argv=None, experiment_name: str = "", run_id: str = "0", csv_filename: str = "gpt2_parity_results.csv"): result = {} from transformers import __version__ as transformers_version @@ -366,6 +387,8 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu raw_onnx_model = onnx_model_paths["raw"] + int_data_type = torch.int64 if args.use_int64_inputs else torch.int32 + if os.path.exists(raw_onnx_model) and not args.overwrite: logger.warning(f"Skip exporting ONNX model since it existed: {raw_onnx_model}") else: @@ -378,9 +401,9 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu args.use_external_data_format, has_position_ids=use_padding, has_attention_mask=use_padding, - input_ids_dtype=torch.int32 if args.use_int32_inputs else torch.int64, - position_ids_dtype=torch.int32 if args.use_int32_inputs else torch.int64, - attention_mask_dtype=torch.int32 if args.use_int32_inputs else torch.int64, + input_ids_dtype=int_data_type, + position_ids_dtype=int_data_type, + attention_mask_dtype=int_data_type, ) fp16_params = {"keep_io_types": args.keep_io_types} @@ -395,11 +418,13 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu is_io_float16 = args.precision == Precision.FLOAT16 and not args.keep_io_types + optimized_ops = "" + all_ops = "" if args.optimize_onnx or args.precision != Precision.FLOAT32: output_path = onnx_model_paths[str(args.precision) if args.precision != Precision.INT8 else "fp32"] logger.info(f"Optimizing model to {output_path}") - gpt2helper.optimize_onnx( + m = gpt2helper.optimize_onnx( raw_onnx_model, output_path, args.precision == Precision.FLOAT16, @@ -407,8 +432,18 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu model.config.hidden_size, args.use_external_data_format, auto_mixed_precision=args.auto_mixed_precision, + stage=args.stage, **fp16_params, ) + + nodes = m.nodes() + op_list = set([node.op_type for node in nodes]) + all_ops = ",".join(op_list) + + # print optimized operators + optimized_op_counter = m.get_fused_operator_statistics() + if optimized_op_counter: + optimized_ops = ",".join([key for key in optimized_op_counter if optimized_op_counter[key] > 0]) else: output_path = raw_onnx_model @@ -442,14 +477,20 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu model_class=args.model_class, has_position_ids=use_padding, has_attention_mask=use_padding, - input_ids_dtype=torch.int32 if args.use_int32_inputs else torch.int64, - position_ids_dtype=torch.int32 if args.use_int32_inputs else torch.int64, - attention_mask_dtype=torch.int32 if args.use_int32_inputs else torch.int64, + input_ids_dtype=int_data_type, + position_ids_dtype=int_data_type, + attention_mask_dtype=int_data_type, test_cases_per_run=args.test_cases, total_runs=args.test_runs, + stage=args.stage, verbose=args.verbose, ) + # An example configuration for testing performance + batch_size = 8 + sequence_length = 32 if args.stage == 1 else 1 + past_sequence_length = 0 if args.stage == 1 else 32 + latency = gpt2helper.test_performance( session, model, @@ -460,12 +501,12 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu model_class=args.model_class, has_position_ids=use_padding, has_attention_mask=use_padding, - input_ids_dtype=torch.int32 if args.use_int32_inputs else torch.int64, - position_ids_dtype=torch.int32 if args.use_int32_inputs else torch.int64, - attention_mask_dtype=torch.int32 if args.use_int32_inputs else torch.int64, - batch_size=8, - sequence_length=1, - past_sequence_length=32, + input_ids_dtype=int_data_type, + position_ids_dtype=int_data_type, + attention_mask_dtype=int_data_type, + batch_size=batch_size, + sequence_length=sequence_length, + past_sequence_length=past_sequence_length, ) if args.precision == Precision.FLOAT16: @@ -476,7 +517,7 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu from onnxruntime import __version__ as ort_version - latency_name = get_latency_name() + latency_name = get_latency_name(batch_size, sequence_length, past_sequence_length) csv_file_existed = os.path.exists(csv_filename) with open(csv_filename, mode="a", newline="") as csv_file: column_names = [ @@ -484,6 +525,7 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu "run_id", "model_name", "model_class", + "stage", "gpu", "precision", "optimizer", @@ -495,8 +537,9 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu "node_block_list", "force_fp16_initializers", "auto_mixed_precision", - "ORT_TRANSFORMER_OPTIONS", - "ORT_CUDA_GEMM_OPTIONS", + "optimized_operators", + "operators", + "environment_variables", "onnxruntime", latency_name, "top1_match_rate", @@ -517,6 +560,7 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu "run_id": run_id, "model_name": args.model_name_or_path, "model_class": args.model_class, + "stage": args.stage, "gpu": args.use_gpu, "precision": args.precision, "optimizer": args.optimize_onnx, @@ -528,8 +572,9 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu "node_block_list": args.node_block_list, "force_fp16_initializers": args.force_fp16_initializers, "auto_mixed_precision": args.auto_mixed_precision, - "ORT_TRANSFORMER_OPTIONS": os.getenv("ORT_TRANSFORMER_OPTIONS"), - "ORT_CUDA_GEMM_OPTIONS": os.getenv("ORT_CUDA_GEMM_OPTIONS"), + "optimized_operators": optimized_ops, + "operators": all_ops, + "environment_variables": get_ort_environment_variables(), "onnxruntime": ort_version, latency_name: f"{latency:.2f}", "diff_50_percentile": parity_result["max_diff_percentile_50"], @@ -576,12 +621,12 @@ def main(argv=None, experiment_name="", run_id=0, csv_filename="gpt2_parity_resu position_ids.masked_fill_(position_ids < 0, 0) inputs = { - "input_ids": input_ids.to(torch.int32) if args.use_int32_inputs else input_ids, - "position_ids": position_ids.to(torch.int32) if args.use_int32_inputs else position_ids, - "attention_mask": attention_mask.to(torch.int32) if args.use_int32_inputs else attention_mask, + "input_ids": input_ids.to(int_data_type), + "position_ids": position_ids.to(int_data_type), + "attention_mask": attention_mask.to(int_data_type), } else: - inputs = {"input_ids": input_ids.to(torch.int32) if args.use_int32_inputs else input_ids} + inputs = {"input_ids": input_ids.to(int_data_type)} if model_type == "beam_search_step" or model_type == "configurable_one_step_search": beam_select_idx = torch.zeros([1, input_ids.shape[0]]).long() diff --git a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py index a154a52858..f8112013d4 100644 --- a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py @@ -24,6 +24,7 @@ sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) from benchmark_helper import Precision from float16 import float_to_float16_max_diff +from fusion_options import AttentionMaskFormat from io_binding_helper import IOBindingHelper from onnx_model import OnnxModel from torch_onnx_export_helper import torch_onnx_export @@ -507,6 +508,7 @@ class Gpt2Helper: hidden_size, use_external_data_format=False, auto_mixed_precision=False, + stage=0, **kwargs, ): """Optimize ONNX model with an option to convert it to use mixed precision.""" @@ -514,9 +516,12 @@ class Gpt2Helper: from optimizer import optimize_model optimization_options = FusionOptions("gpt2") - # optimization_options.enable_gelu = False - # optimization_options.enable_layer_norm = False - # optimization_options.enable_attention = False + + if is_float16 and stage == 1: + # For init_decoder, enable mask index to use fused causal cuda kernel. + # Potentially, we can add other optimization like unpad for effective transformer + optimization_options.attention_mask_format = AttentionMaskFormat.MaskIndexEnd + m = optimize_model( onnx_model_path, model_type="gpt2", @@ -536,17 +541,25 @@ class Gpt2Helper: m.convert_float_to_float16(use_symbolic_shape_infer=True, **kwargs) m.save_model_to_file(optimized_model_path, use_external_data_format) + return m @staticmethod def auto_mixed_precision( onnx_model: OnnxModel, - op_block_list: List[str] = ["Add", "LayerNormalization", "FastGelu"], + op_block_list: List[str] = [ + "Add", + "LayerNormalization", + "SkipLayerNormalization", + "FastGelu", + "EmbedLayerNormalization", + ], ): """Convert GPT-2 model to mixed precision. - It detects whether original model has fp16 precision weights, and set parameters for float16 conversion automatically. + It detects whether original model has fp16 weights, and set parameters for float16 conversion automatically. Args: onnx_model (OnnxModel): optimized ONNX model - op_block_list (List[str], optional): . Defaults to ['Add', 'LayerNormalization', 'FastGelu'] + op_block_list (List[str], optional): operators to compute in fp32. Defaults to ["Add", "LayerNormalization", + "SkipLayerNormalization", "FastGelu", "EmbedLayerNormalization"] Returns: parameters(dict): a dictionary of parameters used in float16 conversion """ @@ -770,6 +783,7 @@ class Gpt2Helper: input_ids_dtype=torch.int32, position_ids_dtype=torch.int32, attention_mask_dtype=torch.int32, + stage=0, verbose=False, enable_pickle_output=False, ): @@ -801,7 +815,7 @@ class Gpt2Helper: for i in range(total_test_cases): run_id = int(i / test_cases_per_run) sequence_length = random.randint(1, max_seq_len) - past_sequence_length = random.randint(0, max_past_seq_len) + past_sequence_length = 0 if (stage == 1) else random.randint(0, max_past_seq_len) batch_size = random.randint(1, max_batch_size) logger.debug( diff --git a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_parity.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_parity.py index 4c52aaaf2e..e48fcc1cfc 100644 --- a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_parity.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_parity.py @@ -13,20 +13,20 @@ import argparse import csv import datetime +import json import logging import os import sys import onnx import scipy.stats -import torch -from convert_to_onnx import get_latency_name, main +from convert_to_onnx import main from gpt2_helper import PRETRAINED_GPT2_MODELS, Gpt2Helper from onnx_model import OnnxModel sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) -from benchmark_helper import setup_logger +from benchmark_helper import get_ort_environment_variables, setup_logger logger = logging.getLogger("") @@ -85,6 +85,14 @@ def parse_arguments(argv=None): ) parser.set_defaults(skip_test=False) + parser.add_argument( + "--overwrite", + required=False, + action="store_true", + help="Overwrite existing csv file", + ) + parser.set_defaults(overwrite=False) + args = parser.parse_args(argv) return args @@ -110,11 +118,13 @@ class ParityTask: run_id=run_id, csv_filename=self.csv_path, ) + if result: + self.results.append(result) except: logger.exception(f"Failed to run experiment {experiment_name}") + result = None - if result: - self.results.append(result) + return result def load_results_from_csv(csv_path): @@ -128,9 +138,17 @@ def load_results_from_csv(csv_path): return rows +def get_latency(row): + for name in row: + if name.startswith("average_latency(batch_size="): + return float(row[name]) + + raise RuntimeError("Failed to get average_latency from output") + + def score(row): """Scoring function based on 3 metrics. The larger score is better.""" - latency_in_ms = float(row[get_latency_name()]) + latency_in_ms = get_latency(row) top1_match_rate = float(row["top1_match_rate"]) onnx_size_in_MB = float(row["onnx_size_in_MB"]) # A simple scoring function: cost of 0.1ms latency ~ 0.1% match rate ~ 100MB size @@ -167,17 +185,15 @@ def print_wins(wins, rows, test_name): for row in rows: if row["run_id"] == key: logger.info( - "{:02d}: WINs={:02d}, run_id={}, latency={:5.2f} top1_match={:.4f} size={}_MB experiment={} {}".format( + "{:02d}: WINs={:02d}, run_id={}, latency={:5.2f}, top1_match={:.4f}, size={}_MB, experiment={}, {}".format( rank, value, key, - float(row[get_latency_name()]), + get_latency(row), float(row["top1_match_rate"]), row["onnx_size_in_MB"], row["experiment"], - " (Half2 Disabled)" - if (row["ORT_CUDA_GEMM_OPTIONS"] == "4" and "Half2" not in row["experiment"]) - else "", + get_ort_environment_variables(), ) ) break @@ -215,6 +231,11 @@ def run_significance_test(rows, output_csv_path): for i in range(num_results - 1): result1 = rows[i] + if isinstance(result1["top1_match_rate_per_run"], str): + a = json.loads(result1["top1_match_rate_per_run"]) + else: + a = result1["top1_match_rate_per_run"] + for j in range(i + 1, num_results, 1): result2 = rows[j] @@ -226,13 +247,9 @@ def run_significance_test(rows, output_csv_path): if not all_matched: continue - if isinstance(result1["top1_match_rate_per_run"], str): - import json - - a = json.loads(result1["top1_match_rate_per_run"]) + if isinstance(result2["top1_match_rate_per_run"], str): b = json.loads(result2["top1_match_rate_per_run"]) else: - a = result1["top1_match_rate_per_run"] b = result2["top1_match_rate_per_run"] try: @@ -244,7 +261,7 @@ def run_significance_test(rows, output_csv_path): utest_pvalue = None ttest_statistic, ttest_pvalue = scipy.stats.ttest_ind(a, b, axis=None, equal_var=True) - if utest_pvalue < 0.05: + if utest_pvalue is not None and utest_pvalue < 0.05: if float(result1["top1_match_rate"]) > float(result2["top1_match_rate"]): utest_wins[result1["run_id"]] += 1 else: @@ -317,11 +334,16 @@ def run_candidate( ): parameters = get_mixed_precision_parameters(args, last_matmul_node_name, op_block_list) op_block_list_str = ",".join(sorted(op_block_list)) - name_suffix = " (Half2 Disabled)" if os.getenv("ORT_CUDA_GEMM_OPTIONS") == "4" else "" + if op_block_list: - name = f"Mixed precision baseline + {op_block_list_str} in FP32{name_suffix}" + name = f"Mixed precision baseline + {op_block_list_str} in FP32" else: - name = f"Mixed precision baseline (logits output and last MatMul node {last_matmul_node_name} in FP32){name_suffix}" + name = f"Mixed precision baseline (logits output and last MatMul node {last_matmul_node_name} in FP32)" + + env_vars = get_ort_environment_variables() + if env_vars: + name = name + f" ({env_vars})" + task.run(parameters, name) @@ -340,12 +362,7 @@ def get_baselines(args): return fp32_baseline, fp16_baseline -def get_all_operators(): - """All operators in the optimized model""" - return "Attention Gather Add LayerNormalization FastGelu MatMul".split() - - -def run_tuning_step0(task, fp16_baseline): +def run_tuning_step0(task, fp16_baseline, all_ops, optimized_ops): """Step 0 is to check which operator in FP16 causes most loss""" fp32_logits = ["--io_block_list", "logits"] task.run(fp16_baseline + fp32_logits, "FP16 except logits") @@ -353,23 +370,28 @@ def run_tuning_step0(task, fp16_baseline): fp32_io = ["--keep_io_types"] task.run(fp16_baseline + fp32_io, "Graph I/O FP32, Other FP16") - op_list = get_all_operators() - # task.run(fp16_baseline + fp32_io + ["--op_block_list"] + [o for o in op_list], "Everthing in FP32") - # Only weights in FP16 task.run( - fp16_baseline + fp32_io + ["--op_block_list"] + [o for o in op_list] + ["--force_fp16_initializers"], + fp16_baseline + fp32_io + ["--op_block_list"] + [o for o in all_ops] + ["--force_fp16_initializers"], "FP32 except weights in FP16", ) + optimized_ops_results = [] + op_list = optimized_ops for op in op_list: op_block_list = ["--op_block_list"] + [o for o in op_list if o != op] - task.run(fp16_baseline + fp32_io + op_block_list, f"FP32 except {op} in FP16") + result = task.run(fp16_baseline + fp32_io + op_block_list, f"FP32 except {op} in FP16") + if result: + optimized_ops_results.append(result) + + # Check which optimized operator causes the most loss in precision + min_result = min(optimized_ops_results, key=lambda y: y["top1_match_rate"]) + print("step 0: optimized operator causes the most loss in precision", min_result) -def run_tuning_step1(task, mixed_precision_baseline): - """Step 1 is to figure out which operator in FP32 could benefit most""" - for op in get_all_operators(): +def run_tuning_step1(task, mixed_precision_baseline, optimized_ops): + """Step 1 is to figure out which optimized operator in FP32 could benefit most""" + for op in optimized_ops: op_block_list = ["--op_block_list", op] task.run( mixed_precision_baseline + op_block_list, @@ -377,32 +399,21 @@ def run_tuning_step1(task, mixed_precision_baseline): ) -def run_tuning_step2(task, mixed_precision_baseline): - """Assumed that you have run step 1 to figure out that Logits FP32 and Add FP32 is important, - Step 2 is to figure out a combination of two operators (one is Add from step one) to get better result +def run_tuning_step2(task, mixed_precision_baseline, optimized_ops): + """Assumed that you have run step 0 and 1 to figure out that Logits FP32 and some operators shall be in FP32, + This step will try add one more operator. """ - for op in get_all_operators(): - if op not in ["Add"]: - op_block_list = ["--op_block_list", "Add", op] + candidate_fp32_ops = ["FastGelu", "LayerNormalization", "SkipLayerNormalization"] + fp32_ops = [x for x in candidate_fp32_ops if x in optimized_ops] + for op in optimized_ops: + if op not in fp32_ops: + op_block_list = fp32_ops + [op] task.run( - mixed_precision_baseline + op_block_list, - f"Mixed precision baseline + Add,{op} in FP32", + mixed_precision_baseline + ["--op_block_list"] + op_block_list, + "Mixed precision baseline + {},{} in FP32".format(",".join(fp32_ops), op), ) -def run_parity_disable_half2(task: ParityTask, args): - onnx_model_paths = Gpt2Helper.get_onnx_paths( - "onnx_models", - args.model_name_or_path, - new_folder=args.use_external_data_format, - remove_existing=[], - ) - last_matmul_node_name = get_last_matmul_node_name(onnx_model_paths["raw"]) - run_candidate(task, args, last_matmul_node_name, op_block_list=[]) - run_candidate(task, args, last_matmul_node_name, op_block_list=["Add"]) - run_candidate(task, args, last_matmul_node_name, op_block_list=["LayerNormalization", "Add"]) - - def run_parity(task: ParityTask, args): onnx_model_paths = Gpt2Helper.get_onnx_paths( "onnx_models", @@ -413,7 +424,19 @@ def run_parity(task: ParityTask, args): fp32_baseline, fp16_baseline = get_baselines(args) - task.run(fp32_baseline, "FP32 baseline") + result = task.run(fp32_baseline, "FP32 baseline") + + optimized_ops = [] + if result and ("optimized_operators" in result) and result["optimized_operators"]: + optimized_ops = result["optimized_operators"].split(",") + else: + raise RuntimeError("Failed to get optimized operators") + + all_ops = [] + if result and ("operators" in result) and result["operators"]: + all_ops = result["operators"].split(",") + else: + raise RuntimeError("Failed to get operators") # The following tests for fp16 requires GPU if not args.use_gpu: @@ -427,41 +450,36 @@ def run_parity(task: ParityTask, args): # Mixed precision baseline run_candidate(task, args, last_matmul_node_name, op_block_list=[]) - # Result from tuning step 1 - run_candidate(task, args, last_matmul_node_name, op_block_list=["Add"]) + get_fp32_ops = lambda x: [op for op in x if op in all_ops] if args.all: - run_tuning_step0(task, fp16_baseline) + run_tuning_step0(task, fp16_baseline, all_ops, optimized_ops) mixed_precision_baseline = get_mixed_precision_parameters(args, last_matmul_node_name, op_block_list=[]) - run_tuning_step1(task, mixed_precision_baseline) - run_tuning_step2(task, mixed_precision_baseline) + run_tuning_step1(task, mixed_precision_baseline, optimized_ops) + run_tuning_step2(task, mixed_precision_baseline, optimized_ops) else: run_candidate( task, args, last_matmul_node_name, - op_block_list=["LayerNormalization", "Add"], + op_block_list=get_fp32_ops(["SkipLayerNormalization", "LayerNormalization", "Add"]), ) - run_candidate(task, args, last_matmul_node_name, op_block_list=["FastGelu", "Add"]) + run_candidate(task, args, last_matmul_node_name, op_block_list=["FastGelu"]) # Run a few good candidates run_candidate( task, args, last_matmul_node_name, - op_block_list=["FastGelu", "LayerNormalization", "Add"], + op_block_list=get_fp32_ops(["FastGelu", "SkipLayerNormalization", "LayerNormalization", "Add"]), ) run_candidate( task, args, last_matmul_node_name, - op_block_list=["FastGelu", "LayerNormalization", "Add", "Gather"], - ) - run_candidate( - task, - args, - last_matmul_node_name, - op_block_list=["FastGelu", "LayerNormalization", "Add", "Gather", "MatMul"], + op_block_list=get_fp32_ops( + ["FastGelu", "EmbedLayerNormalization", "SkipLayerNormalization", "LayerNormalization", "Add"] + ), ) @@ -471,19 +489,23 @@ if __name__ == "__main__": if args.test_cases < 100 or args.runs < 20 or args.test_cases * args.runs < 10000: logger.warning( - "Not enough test cases or runs to get stable results or test significance. Recommend test_cases >= 100, runs >= 20, test_cases * runs >= 10000." + "Not enough test cases or runs to get stable results or test significance. " + "Recommend test_cases >= 100, runs >= 20, test_cases * runs >= 10000." ) + if os.path.exists(args.csv) and not args.skip_test: + if not args.overwrite: + raise RuntimeError( + f"Output file {args.csv} existed. Please remove the file, or use either --skip_test or --overwrite." + ) + else: + logger.info("Remove existing file %s since --overwrite is specified", args.csv) + os.remove(args.csv) + task = ParityTask(args.test_cases, args.runs, args.csv) if not args.skip_test: - if os.getenv("ORT_CUDA_GEMM_OPTIONS") == "4" and args.use_gpu: - assert ( - torch.cuda.get_device_capability()[0] >= 7 - ), "half2 kernel is not avaiable in current GPU device. Please set environment variable ORT_CUDA_GEMM_OPTIONS=0 or use supported GPU like V100 or T4" - run_parity_disable_half2(task, args) - else: - run_parity(task, args) + run_parity(task, args) try: rows = load_results_from_csv(task.csv_path) diff --git a/onnxruntime/test/python/transformers/test_gpt2.py b/onnxruntime/test/python/transformers/test_gpt2_benchmark.py similarity index 87% rename from onnxruntime/test/python/transformers/test_gpt2.py rename to onnxruntime/test/python/transformers/test_gpt2_benchmark.py index 974c1f5a97..48df008060 100644 --- a/onnxruntime/test/python/transformers/test_gpt2.py +++ b/onnxruntime/test/python/transformers/test_gpt2_benchmark.py @@ -32,6 +32,10 @@ class TestGpt2(unittest.TestCase): self.assertIsNotNone(csv_filename) self.assertTrue(os.path.exists(csv_filename)) + @pytest.mark.slow + def test_gpt2_stage1(self): + self.run_benchmark_gpt2("-m gpt2 --precision fp32 -v -b 1 --sequence_lengths 8 -s 0 --stage 1") + @pytest.mark.slow def test_gpt2_fp32(self): self.run_benchmark_gpt2("-m gpt2 --precision fp32 -v -b 1 --sequence_lengths 2 -s 3") @@ -69,12 +73,12 @@ class TestGpt2(unittest.TestCase): "-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 --sequence_lengths 5 -s 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 --sequence_lengths 5 -s 3 --use_gpu" + # ) @pytest.mark.slow def test_gpt2_configurable_one_step_search_int8(self): diff --git a/onnxruntime/test/python/transformers/test_gpt2_to_onnx.py b/onnxruntime/test/python/transformers/test_gpt2_to_onnx.py new file mode 100644 index 0000000000..bc07b7a6b4 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_gpt2_to_onnx.py @@ -0,0 +1,63 @@ +# ------------------------------------------------------------------------- +# 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 unittest + +import coloredlogs +import pytest +from parity_utilities import find_transformers_source + +from onnxruntime import get_available_providers + +if find_transformers_source(sub_dir_paths=["models", "gpt2"]): + from convert_to_onnx import main as gpt2_to_onnx +else: + from onnxruntime.transformers.models.gpt2.convert_to_onnx import main as gpt2_to_onnx + + +class TestGpt2ConvertToOnnx(unittest.TestCase): + """Test convert_to_onnx.py of GPT-2 model.""" + + def setUp(self): + self.test_cuda = "CUDAExecutionProvider" in get_available_providers() + + def run_gpt2_to_onnx(self, arguments: str, stage: int): + result = gpt2_to_onnx(arguments.split()) + self.assertIsNotNone(result) + self.assertTrue(isinstance(result, dict)) + self.assertIn("top1_match_rate", result) + self.assertGreater(result["top1_match_rate"], 0.98) + self.assertIn("stage", result) + self.assertEqual(result["stage"], stage) + if stage == 1: + self.assertIn("average_latency(batch_size=8,sequence_length=32,past_sequence_length=0)", result) + else: + self.assertIn("average_latency(batch_size=8,sequence_length=1,past_sequence_length=32)", result) + + @pytest.mark.slow + def test_stage1(self): + """Test stage 1 onnx model""" + if self.test_cuda: + self.run_gpt2_to_onnx("-m distilgpt2 -p fp32 --use_gpu -s 1 -t 100 --overwrite", 1) + + @pytest.mark.slow + def test_stage2(self): + """Test stage 2 onnx model""" + if self.test_cuda: + self.run_gpt2_to_onnx("-m distilgpt2 -p fp32 --use_gpu -s 2 -t 100 --overwrite", 2) + + @pytest.mark.slow + def test_auto_mixed_precision(self): + """Test mixed preicison onnx model""" + if self.test_cuda: + self.run_gpt2_to_onnx("-m distilgpt2 -p fp32 --use_gpu -p fp16 -o -a -t 100 --overwrite", 0) + + +if __name__ == "__main__": + coloredlogs.install(fmt="%(message)s") + logging.getLogger("transformers").setLevel(logging.ERROR) + unittest.main()