diff --git a/onnxruntime/python/tools/transformers/benchmark.py b/onnxruntime/python/tools/transformers/benchmark.py index f9a0f5ad8d..2b3c433808 100644 --- a/onnxruntime/python/tools/transformers/benchmark.py +++ b/onnxruntime/python/tools/transformers/benchmark.py @@ -235,11 +235,13 @@ def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwri def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_attention_heads, hidden_size, use_gpu, - fp16, overwrite): + fp16, use_raw_attention_mask, overwrite): if overwrite or not os.path.exists(optimized_model_path): from optimizer import optimize_model from onnx_model_bert import BertOptimizationOptions optimization_options = BertOptimizationOptions(model_type) + if use_raw_attention_mask: + optimization_options.use_raw_attention_mask() if fp16: optimization_options.enable_gelu_approximation = True @@ -264,7 +266,7 @@ def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_a def export_onnx_model(model_name, cache_dir, onnx_dir, input_names, use_gpu, fp16, optimize_onnx, validate_onnx, - overwrite): + use_raw_attention_mask, overwrite): config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir) model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir) model.cpu() @@ -310,7 +312,7 @@ def export_onnx_model(model_name, cache_dir, onnx_dir, input_names, use_gpu, fp1 model_type = MODELS[model_name][3] optimized_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), True, use_gpu, fp16, False) optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, config.num_attention_heads, - config.hidden_size, use_gpu, fp16, overwrite) + config.hidden_size, use_gpu, fp16, use_raw_attention_mask, overwrite) onnx_model_path = optimized_model_path if validate_onnx: @@ -390,7 +392,8 @@ def allocateOutputBuffers(output_buffers, max_last_state_size, max_pooler_size, def run_onnxruntime(use_gpu, model_names, fp16, batch_sizes, sequence_lengths, repeat_times, input_counts, - optimize_onnx, validate_onnx, cache_dir, onnx_dir, verbose, overwrite, disable_ort_io_binding): + optimize_onnx, validate_onnx, cache_dir, onnx_dir, verbose, overwrite, disable_ort_io_binding, + use_raw_attention_mask): import onnxruntime results = [] @@ -414,7 +417,7 @@ def run_onnxruntime(use_gpu, model_names, fp16, batch_sizes, sequence_lengths, r with torch.no_grad(): onnx_model_file, is_valid_onnx_model, vocab_size, max_sequence_length = export_onnx_model( model_name, cache_dir, onnx_dir, input_names, use_gpu, fp16, optimize_onnx, validate_onnx, - overwrite) + use_raw_attention_mask, overwrite) if not is_valid_onnx_model: continue @@ -691,6 +694,12 @@ def parse_arguments(): help='Disable running ONNX Runtime with binded inputs and outputs. ') parser.set_defaults(disable_ort_io_binding=False) + parser.add_argument('--use_raw_attention_mask', + required=False, + action='store_true', + help='Use raw attention mask in Attention operator for Bert models.') + parser.set_defaults(use_raw_attention_mask=False) + args = parser.parse_args() return args @@ -742,7 +751,7 @@ def main(): results += run_onnxruntime(args.use_gpu, args.models, args.fp16, args.batch_sizes, args.sequence_lengths, args.test_times, args.input_counts, args.optimize_onnx, args.validate_onnx, args.cache_dir, args.onnx_dir, args.verbose, args.overwrite, - args.disable_ort_io_binding) + args.disable_ort_io_binding, args.use_raw_attention_mask) except: logger.error(f"Exception", exc_info=True) diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index 4118d1348b..f5e87c08f0 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -4,6 +4,7 @@ #-------------------------------------------------------------------------- import numpy as np from logging import getLogger +from enum import Enum from onnx import helper, numpy_helper, TensorProto from onnx_model import OnnxModel from fusion_base import Fusion @@ -12,6 +13,12 @@ from fusion_utils import FusionUtils logger = getLogger(__name__) +class AttentionMaskFormat: + MaskIndexEnd = 0 + MaskIndexEndAndStart = 1 + AttentionMask = 2 + + class AttentionMask(): """ Fuse Attention subgraph into one Attention node. @@ -23,6 +30,10 @@ class AttentionMask(): # A lookup table with mask input as key, and cast (to int32) output as value self.mask_casted = {} self.utils = FusionUtils(model) + self.mask_format = AttentionMaskFormat.MaskIndexEnd + + def set_mask_format(self, mask_format: AttentionMaskFormat): + self.mask_format = mask_format def set_mask_indice(self, mask, mask_index): if mask in self.mask_indice: @@ -47,7 +58,12 @@ class AttentionMask(): if casted: self.mask_casted[input] = input_name - # Add a mask processing node + # Attention supports int32 attention mask (2D) since 1.4.0 + if self.mask_format == AttentionMaskFormat.AttentionMask: + self.mask_indice[input] = input_name + return input_name + + # Add a mask processing node to convert attention mask to mask index (1D) output_name = self.model.create_node_name('mask_index') mask_index_node = helper.make_node('ReduceSum', inputs=[input_name], diff --git a/onnxruntime/python/tools/transformers/fusion_embedlayer.py b/onnxruntime/python/tools/transformers/fusion_embedlayer.py index d42e01bb5a..1ccf82ad8d 100644 --- a/onnxruntime/python/tools/transformers/fusion_embedlayer.py +++ b/onnxruntime/python/tools/transformers/fusion_embedlayer.py @@ -41,7 +41,7 @@ class FusionEmbedLayerNoMask(Fusion): self.utils = FusionUtils(model) def fuse(self, node, input_name_to_nodes, output_name_to_node): - # already fused. Assumes that only one mebedding layer in a transformer model. + # Embed layer has been fused. Assumes that only one mebedding layer in a transformer model. if self.nodes_to_add: return @@ -61,7 +61,7 @@ class FusionEmbedLayerNoMask(Fusion): normalize_node = node word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 0]) if word_embedding_path is None: - logger.info("Failed to find word embedding") + logger.info("Word embedding path is not found. Embed layer cannot be fused.") return add_node, word_embedding_gather = word_embedding_path input_ids = word_embedding_gather.input[1] @@ -87,7 +87,7 @@ class FusionEmbedLayerNoMask(Fusion): if position_embedding_path is not None: position_embedding_weight_node, position_embedding_expand = position_embedding_path else: - logger.info("Failed to find position embedding") + logger.info("Position embedding path is not found. Embed layer cannot be fused.") return if position_embedding_shape is not None and position_embedding_shape.input[0] != input_ids: @@ -98,7 +98,7 @@ class FusionEmbedLayerNoMask(Fusion): if segment_embedding_path is None: segment_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 1]) if segment_embedding_path is None: - logger.info("Failed to find segment embedding") + logger.info("Segment embedding is not found. Embed layer cannot be fused.") return _, segment_embedding_gather = segment_embedding_path else: @@ -206,19 +206,23 @@ class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): else: embed_node = self.nodes_to_add.pop() mask_input_name = next(iter(self.mask_indice)) - mask_output_name = self.mask_indice[mask_input_name] - mask_node = output_name_to_node[mask_output_name] - - self.nodes_to_remove.extend([mask_node]) # store inputs for further processing self.mask_input_name = mask_input_name - # When mask has been casted to int32, use that casted one as input of embed layer norm. - if mask_input_name in self.mask_casted: - mask_input_name = self.mask_casted[mask_input_name] + mask_output_name = self.mask_indice[mask_input_name] + mask_node = output_name_to_node[mask_output_name] + # Remove mask processing node (like ReduceSum), but keep Cast node for 2D attention mask. + if mask_node.op_type == "ReduceSum" or mask_node.op_type == "MaskIndex": + self.nodes_to_remove.extend([mask_node]) + + # When mask has been casted to int32, use that casted one as input of embed layer norm. + if mask_input_name in self.mask_casted: + mask_input_name = self.mask_casted[mask_input_name] + + embed_node.input.append(mask_input_name) + + embed_node.output[1] = mask_output_name - embed_node.input.append(mask_input_name) - embed_node.output[1] = mask_output_name self.nodes_to_add.append(embed_node) self.prune_graph = True diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert.py b/onnxruntime/python/tools/transformers/onnx_model_bert.py index ebf6ed994e..f27cbb12ac 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert.py @@ -10,7 +10,7 @@ from fusion_reshape import FusionReshape from fusion_layernorm import FusionLayerNormalization, FusionLayerNormalizationTF from fusion_skiplayernorm import FusionSkipLayerNormalization, FusionBiasSkipLayerNormalization from fusion_embedlayer import FusionEmbedLayerNormalization -from fusion_attention import FusionAttention, AttentionMask +from fusion_attention import FusionAttention, AttentionMask, AttentionMaskFormat from fusion_gelu import FusionGelu from fusion_fastgelu import FusionFastGelu from fusion_biasgelu import FusionBiasGelu @@ -30,9 +30,14 @@ class BertOptimizationOptions: self.enable_bias_skip_layer_norm = True self.enable_bias_gelu = True self.enable_gelu_approximation = False + self.attention_mask_format = AttentionMaskFormat.MaskIndexEnd if model_type == 'gpt2': self.enable_skip_layer_norm = False + self.attention_mask_format = AttentionMaskFormat.AttentionMask + + def use_raw_attention_mask(self): + self.attention_mask_format = AttentionMaskFormat.AttentionMask class BertOnnxModel(OnnxModel): @@ -158,8 +163,9 @@ class BertOnnxModel(OnnxModel): # After: # input_ids --> Shape --> ConstantOfShape -->Cast --> EmbedLayerNormaliation/ReduceSum # TODO: merge ConstantOfShape -->Cast to ConstantOfShape (need update the data type of value) - if node.op_type == 'EmbedLayerNormalization' or node.op_type == 'ReduceSum': - i = 1 if node.op_type == 'EmbedLayerNormalization' else 0 + op_input_id = {"EmbedLayerNormalization": 1, "ReduceSum": 0, "Attention": 3} + if node.op_type in op_input_id: + i = op_input_id[node.op_type] parent_nodes = self.match_parent_path( node, ['Cast', 'ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [i, 0, 0, 0, 0, 0], output_name_to_node) @@ -208,6 +214,8 @@ class BertOnnxModel(OnnxModel): self.fuse_skip_layer_norm() if (options is None) or options.enable_attention: + if options is not None: + self.attention_mask.set_mask_format(options.attention_mask_format) self.fuse_attention() if (options is None) or options.enable_embed_layer_norm: diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py index e211593aba..d6e4fd5ea8 100644 --- a/onnxruntime/python/tools/transformers/optimizer.py +++ b/onnxruntime/python/tools/transformers/optimizer.py @@ -87,7 +87,7 @@ def optimize_by_onnxruntime(onnx_model_path: str, assert 'CUDAExecutionProvider' in session.get_providers() # Make sure there is GPU assert os.path.exists(optimized_model_path) and os.path.isfile(optimized_model_path) - logger.info("Save optimized model by onnxruntime to {}".format(optimized_model_path)) + logger.debug("Save optimized model by onnxruntime to {}".format(optimized_model_path)) return optimized_model_path @@ -186,6 +186,12 @@ def _parse_arguments(): help="enable Gelu/BiasGelu to FastGelu conversion") parser.set_defaults(enable_gelu_approximation=False) + parser.add_argument('--use_raw_attention_mask', + required=False, + action='store_true', + help="use raw attention mask instead of mask index in attention operator") + parser.set_defaults(use_raw_attention_mask=False) + parser.add_argument('--verbose', required=False, action='store_true') parser.set_defaults(verbose=False) @@ -225,6 +231,9 @@ def _get_optimization_options(args): optimization_options.enable_bias_gelu = False if args.enable_gelu_approximation: optimization_options.enable_gelu_approximation = True + if args.use_raw_attention_mask: + optimization_options.use_raw_attention_mask() + return optimization_options @@ -258,16 +267,14 @@ def optimize_model(input, """ (optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type] - input_model_path = input - if opt_level > 1: # Optimization specified for an execution provider. - input_model_path = optimize_by_onnxruntime(input_model_path, use_gpu=use_gpu, opt_level=opt_level) + temp_model_path = optimize_by_onnxruntime(input, use_gpu=use_gpu, opt_level=opt_level) elif run_onnxruntime: # Use Onnxruntime to do optimizations (like constant folding and cast elimation) that is not specified to exection provider. # CPU provider is used here so that there is no extra node for GPU memory copy. - input_model_path = optimize_by_onnxruntime(input_model_path, use_gpu=False, opt_level=1) + temp_model_path = optimize_by_onnxruntime(input, use_gpu=False, opt_level=1) - model = load_model(input_model_path, format=None, load_external_data=True) + model = load_model(temp_model_path, format=None, load_external_data=True) if model.producer_name and producer != model.producer_name: logger.warning( @@ -282,6 +289,10 @@ def optimize_model(input, if not only_onnxruntime: optimizer.optimize(optimization_options) + # Remove the temporary model. + os.remove(temp_model_path) + logger.debug("Remove tempoary model: {}".format(temp_model_path)) + return optimizer diff --git a/onnxruntime/python/tools/transformers/run_benchmark.sh b/onnxruntime/python/tools/transformers/run_benchmark.sh index 7d0557dbb5..b8c9f83a5d 100644 --- a/onnxruntime/python/tools/transformers/run_benchmark.sh +++ b/onnxruntime/python/tools/transformers/run_benchmark.sh @@ -27,7 +27,7 @@ if [ "$run_cpu" = true ] ; then average_over=100 fi -# enable optimizer (use script instead of OnnxRuntime for graph optimization) +# Enable optimizer (use script instead of OnnxRuntime for graph optimization) use_optimizer=true # Batch Sizes and Sequence Lengths @@ -45,12 +45,21 @@ models_to_test="bert-base-cased roberta-base gpt2" # If you have mutliple GPUs, you can choose one GPU for test. Here is an example to use the second GPU: # export CUDA_VISIBLE_DEVICES=1 -#This script will generate a logs file with a list of commands used in tests. +# This script will generate a logs file with a list of commands used in tests. echo echo "ort=$run_ort torch=$run_torch torchscript=$run_torchscript gpu_fp32=$run_gpu_fp32 gpu_fp16=$run_gpu_fp16 cpu=$run_cpu optimizer=$use_optimizer batch=$batch_sizes sequence=$sequence_length models=$models_to_test" >> benchmark.log -#Set it to false to skip testing. You can use it to dry run this script with the log file. +# Set it to false to skip testing. You can use it to dry run this script with the log file. run_tests=true +# Directory for downloading pretrained models. +cache_dir="./cache_models" + +# Directory for ONNX models +onnx_dir="./onnx_models" + +# Use raw attention mask in Attention operator or not. +use_raw_attention_mask=false + # ------------------------------------------- if [ "$run_cpu" = true ] ; then if [ "$run_gpu_fp32" = true ] ; then @@ -79,41 +88,46 @@ fi if [ "$run_cli" = true ] ; then echo "Use onnxruntime_tools.transformers.benchmark" - optimizer_script="-m onnxruntime_tools.transformers.benchmark" + benchmark_script="-m onnxruntime_tools.transformers.benchmark" else - optimizer_script="benchmark.py" + benchmark_script="benchmark.py" fi -onnx_export_options="-i $input_counts -v -b 0 --overwrite -f fusion.csv" -benchmark_options="-b $batch_sizes -s $sequence_lengths -t $average_over -f fusion.csv -r result.csv -d detail.csv" +onnx_export_options="-i $input_counts -v -b 0 --overwrite -f fusion.csv -c $cache_dir --onnx_dir $onnx_dir" +benchmark_options="-b $batch_sizes -s $sequence_lengths -t $average_over -f fusion.csv -r result.csv -d detail.csv -c $cache_dir --onnx_dir $onnx_dir" if [ "$use_optimizer" = true ] ; then onnx_export_options="$onnx_export_options -o" benchmark_options="$benchmark_options -o" fi +if [ "$use_raw_attention_mask" = true ] ; then + onnx_export_options="$onnx_export_options --use_raw_attention_mask" + benchmark_options="$benchmark_options --use_raw_attention_mask" +fi + # ------------------------------------------- run_one_test() { if [ "$run_ort" = true ] ; then - echo python $optimizer_script -m $1 $onnx_export_options $2 $3 >> benchmark.log - echo python $optimizer_script -m $1 $benchmark_options $2 $3 -i $input_counts >> benchmark.log + echo python $benchmark_script -m $1 $onnx_export_options $2 $3 >> benchmark.log + echo python $benchmark_script -m $1 $benchmark_options $2 $3 -i $input_counts >> benchmark.log if [ "$run_tests" = true ] ; then - python $optimizer_script -m $1 $onnx_export_options $2 $3 - python $optimizer_script -m $1 $benchmark_options $2 $3 -i $input_counts + python $benchmark_script -m $1 $onnx_export_options $2 $3 + python $benchmark_script -m $1 $benchmark_options $2 $3 -i $input_counts fi fi if [ "$run_torch" = true ] ; then - echo python $optimizer_script -e torch -m $1 $benchmark_options $2 $3 >> benchmark.log + echo python $benchmark_script -e torch -m $1 $benchmark_options $2 $3 >> benchmark.log if [ "$run_tests" = true ] ; then - python $optimizer_script -e torch -m $1 $benchmark_options $2 $3 + python $benchmark_script -e torch -m $1 $benchmark_options $2 $3 fi fi if [ "$run_torchscript" = true ] ; then - echo python $optimizer_script -e torchscript -m $1 $benchmark_options $2 $3 >> benchmark.log + echo python $benchmark_script -e torchscript -m $1 $benchmark_options $2 $3 >> benchmark.log if [ "$run_tests" = true ] ; then - python $optimizer_script -e torchscript -m $1 $benchmark_options $2 $3 + python $benchmark_script -e torchscript -m $1 $benchmark_options $2 $3 fi fi }