From abe1642a0cd5a12039bbd634c117d1588dd309da Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 29 Nov 2022 13:06:39 -0800 Subject: [PATCH] Update fusion for distilbert accuracy test on SQuAD (#13748) (1) Embed layer fusion to work with --use_mask_index. (2) Parse num_heads and hidden_size from a pattern of Concat shape node. (3) Fix a typo (CUDAExcecutionProvider=> CUDAExecutionProvider) in eval_squad.py (4) Update example comments in eval_squad.py to use optimized fp16 model. (5) Update tests in test_optimizer.py --- .../tools/transformers/fusion_attention.py | 35 +++++++++- .../python/tools/transformers/fusion_base.py | 18 ++++-- .../tools/transformers/fusion_embedlayer.py | 64 +++++++++++++++---- .../python/tools/transformers/fusion_shape.py | 6 +- .../transformers/models/bert/eval_squad.py | 13 ++-- .../tools/transformers/onnx_model_bert.py | 17 +++-- .../python/transformers/test_optimizer.py | 42 +++++++++--- 7 files changed, 148 insertions(+), 47 deletions(-) diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index b549980719..0487f77613 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -103,6 +103,37 @@ class FusionAttention(Fusion): self.num_heads_warning = True self.hidden_size_warning = True + def get_num_heads_and_hidden_size_from_concat(self, concat: NodeProto) -> Tuple[int, int]: + """ + Detect num_heads and hidden_size from Concat node in the following subgraph: + + SkipLayerNormalization or EmbedLayerNormalization + / \ + MatMul Shape + | | + Add Gather(indices=0) + \ | + \ Unsqueeze + \ | + \ Concat (*, -1, 12, 64) + \ / + Reshape + | + Transpose + """ + if len(concat.input) == 4: + num_heads = self.model.get_constant_value(concat.input[2]) + head_size = self.model.get_constant_value(concat.input[3]) + if ( + isinstance(num_heads, np.ndarray) + and num_heads.size == 1 + and isinstance(head_size, np.ndarray) + and head_size.size == 1 + ): + return num_heads[0], num_heads[0] * head_size[0] + + return self.num_heads, self.hidden_size + def get_num_heads_and_hidden_size(self, reshape_q: NodeProto) -> Tuple[int, int]: """Detect num_heads and hidden_size from a reshape node. @@ -112,10 +143,12 @@ class FusionAttention(Fusion): Returns: Tuple[int, int]: num_heads and hidden_size """ - # we assume that reshape fusion has done, so the shape is a tensor like [0, 0, num_heads, head_size] q_shape = self.model.get_initializer(reshape_q.input[1]) if q_shape is None: + concat = self.model.get_parent(reshape_q, 1) + if concat is not None and concat.op_type == "Concat": + return self.get_num_heads_and_hidden_size_from_concat(concat) logger.debug(f"{reshape_q.input[1]} is not initializer.") return self.num_heads, self.hidden_size # Fall back to user specified value diff --git a/onnxruntime/python/tools/transformers/fusion_base.py b/onnxruntime/python/tools/transformers/fusion_base.py index 31b3399e27..f338c0a8c4 100644 --- a/onnxruntime/python/tools/transformers/fusion_base.py +++ b/onnxruntime/python/tools/transformers/fusion_base.py @@ -2,10 +2,10 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +from collections import defaultdict from logging import getLogger from typing import List, Union -from onnx import GraphProto from onnx_model import OnnxModel logger = getLogger(__name__) @@ -29,7 +29,10 @@ class Fusion: self.node_name_to_graph_name: dict = {} self.this_graph_name: str = None # It is optional that subclass updates fused_count since we will also check nodes_to_add to get counter. - self.fused_count: int = 0 + self.fused_count: defaultdict = defaultdict(int) + + def increase_counter(self, fused_op_name): + self.fused_count[fused_op_name] += 1 def apply(self): logger.debug(f"start {self.description} fusion...") @@ -46,9 +49,14 @@ class Fusion: self.fuse(node, input_name_to_nodes, output_name_to_node) op_list = [node.op_type for node in self.nodes_to_add] - count = max(self.fused_count, op_list.count(self.fused_op_type)) - if count > 0: - logger.info(f"Fused {self.description} count: {count}") + if self.fused_count: + for key, value in self.fused_count.items(): + if value: + logger.info(f"Fused {key} count: {value}") + else: + count = op_list.count(self.fused_op_type) + if count > 0: + logger.info(f"Fused {self.description} count: {count}") self.model.remove_nodes(self.nodes_to_remove) self.model.add_nodes(self.nodes_to_add, self.node_name_to_graph_name) diff --git a/onnxruntime/python/tools/transformers/fusion_embedlayer.py b/onnxruntime/python/tools/transformers/fusion_embedlayer.py index 867a151bef..bd1d89c72f 100644 --- a/onnxruntime/python/tools/transformers/fusion_embedlayer.py +++ b/onnxruntime/python/tools/transformers/fusion_embedlayer.py @@ -219,7 +219,7 @@ class FusionEmbedLayerNoMask(Fusion): def match_position_embedding_bert(self, position_embedding_gather, input_ids, output_name_to_node): """ Match position embedding path from input_ids to Gather for BERT. - BERT Embedding Layer Pattern: + BERT Embedding Layer Pattern: (input_ids) / \ / Shape @@ -232,7 +232,7 @@ class FusionEmbedLayerNoMask(Fusion): \ | | \ Gather Slice (data[1,512], starts=0, ends=*, axes=1, steps=1) \ / | - Add Gather + Add Gather \ / Add | @@ -682,8 +682,27 @@ class FusionEmbedLayerNoMask(Fusion): class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): - def __init__(self, model: OnnxModel): + def __init__(self, model: OnnxModel, use_mask_index=False): super().__init__(model, "with mask") + self.use_mask_index = use_mask_index + + def replace_mask(self, mask_int32, attention_nodes): + # Inputs of EmbedLayerNorm: input_ids, segment_ids (optional), word_embedding, position_embedding, + # segment_embedding (optional), gamma, beta, mask (optional), position_ids (optional) + embed_node = self.embed_node + if len(embed_node.input) == 7: + embed_node.input.append(mask_int32) + logger.debug("append mask to %s", embed_node.name) + elif len(embed_node.input) > 7 and embed_node.input[7] == "": + embed_node.input[7] = mask_int32 + logger.debug("replace mask in %s", embed_node.name) + else: + logger.debug("skip mask in %s", embed_node.name) + return + + for attention_node in attention_nodes: + logger.debug("update mask_index in %s", attention_node.name) + attention_node.input[3] = embed_node.output[1] def fuse(self, node, input_name_to_nodes, output_name_to_node): # Reset attention and embed_node so that we know fusion is successful when they are not None. @@ -691,13 +710,32 @@ class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): self.embed_node = None super().fuse(node, input_name_to_nodes, output_name_to_node) - if self.attention and self.embed_node: - mask_index = self.attention.input[3] - if mask_index in output_name_to_node: - node = output_name_to_node[mask_index] - if node.op_type == "ReduceSum": - embed_node = self.embed_node - mask_input_name = node.input[0] - self.nodes_to_remove.extend([node]) - embed_node.input.append(mask_input_name) - embed_node.output[1] = mask_index + if self.embed_node is None: + return + + if not self.use_mask_index: + logger.debug("--use_mask_index is not set: EmbedLayerNormalization will not have mask") + self.increase_counter("EmbedLayerNormalization(no mask)") + return + + if self.attention is None: + logger.debug("EmbedLayerNormalization will not have mask since attention node is not found") + self.increase_counter("EmbedLayerNormalization(no mask)") + return + + mask_int32 = self.attention.input[3] + children_nodes = input_name_to_nodes[mask_int32] + 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)") + return + + node = output_name_to_node[mask_int32] + if node.op_type in ["ReduceSum", "Cast"]: + attention_nodes = [node for node in children_nodes if node.op_type == "Attention"] + if node.op_type == "ReduceSum": + mask_int32 = node.input[0] + if len(children_nodes) == len(attention_nodes): + self.nodes_to_remove.append(node) + self.replace_mask(mask_int32, attention_nodes) + self.increase_counter("EmbedLayerNormalization(with mask)") diff --git a/onnxruntime/python/tools/transformers/fusion_shape.py b/onnxruntime/python/tools/transformers/fusion_shape.py index 8abcfaa906..a6a74719b9 100644 --- a/onnxruntime/python/tools/transformers/fusion_shape.py +++ b/onnxruntime/python/tools/transformers/fusion_shape.py @@ -8,6 +8,7 @@ from typing import Dict, List, Union from fusion_base import Fusion from fusion_utils import FusionUtils +from numpy import ndarray from onnx import NodeProto, TensorProto from onnx_model import OnnxModel @@ -58,7 +59,7 @@ class FusionShape(Fusion): | | Unsqueeze(axes=0) Unsqueeze(axes=0) \ / - Concat + Concat | into (2d_input) --> Shape --> @@ -99,12 +100,11 @@ class FusionShape(Fusion): return value = self.model.get_constant_value(gather.input[1]) - from numpy import array_equal, ndarray if not (isinstance(value, ndarray) and value.size == 1 and value.item() == i): return if self.model.find_graph_output(concat_node.output[0]) is None: self.model.replace_input_of_all_nodes(concat_node.output[0], shape_output) - self.fused_count += 1 + self.increase_counter("Reshape") self.prune_graph = True diff --git a/onnxruntime/python/tools/transformers/models/bert/eval_squad.py b/onnxruntime/python/tools/transformers/models/bert/eval_squad.py index 47a380d80d..495c0f017b 100644 --- a/onnxruntime/python/tools/transformers/models/bert/eval_squad.py +++ b/onnxruntime/python/tools/transformers/models/bert/eval_squad.py @@ -7,9 +7,10 @@ # Example to evaluate raw and optimized model for CUDA in Linux: # pip3 install datasets evaluate optimum transformers onnxruntime-gpu # python3 eval_squad.py -m distilbert-base-cased-distilled-squad -# python3 -m onnxruntime.transformers.optimizer --output optimized.onnx --num_heads 12 --hidden_size 768 \ -# --input /home/$USER/.cache/huggingface/hub/distilbert-base-cased-distilled-squad/model.onnx -# python3 eval_squad.py -m distilbert-base-cased-distilled-squad --onnx optimized.onnx +# python3 -m onnxruntime.transformers.optimizer --output optimized_fp16.onnx --num_heads 12 --hidden_size 768 \ +# --input /home/$USER/.cache/huggingface/hub/distilbert-base-cased-distilled-squad/model.onnx \ +# --use_mask_index --float16 +# python3 eval_squad.py -m distilbert-base-cased-distilled-squad --onnx optimized_fp16.onnx import argparse import csv @@ -65,7 +66,7 @@ def load_onnx_model( if provider != "CPUExecutionProvider": model.device = torch.device("cuda:0") - model.model = ORTModel.load_model(onnx_path, "CUDAExecutionProvider") + model.model = ORTModel.load_model(onnx_path, provider) else: model.device = torch.device("cpu") model.model = ORTModel.load_model(onnx_path) @@ -284,8 +285,8 @@ def parse_arguments(argv=None): parser.add_argument( "--provider", required=False, - default="CUDAExcecutionProvider", - help="Select which Execution Provider to use for runs.", + default="CUDAExecutionProvider", + help="Select which Execution Provider to use for runs. Default is CUDAExecutionProvider.", ) parser.add_argument("--use_io_binding", required=False, action="store_true", help="Use IO Binding for GPU.") diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert.py b/onnxruntime/python/tools/transformers/onnx_model_bert.py index 025e0bc0ee..60bd2d7b8e 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert.py @@ -14,7 +14,7 @@ from fusion_gelu import FusionGelu from fusion_gelu_approximation import FusionGeluApproximation from fusion_gemmfastgelu import FusionGemmFastGelu from fusion_layernorm import FusionLayerNormalization, FusionLayerNormalizationTF -from fusion_options import FusionOptions +from fusion_options import AttentionMaskFormat, FusionOptions from fusion_qordered_attention import FusionQOrderedAttention from fusion_qordered_gelu import FusionQOrderedGelu from fusion_qordered_layernorm import FusionQOrderedLayerNormalization @@ -97,8 +97,8 @@ class BertOnnxModel(OnnxModel): fusion = FusionShape(self) fusion.apply() - def fuse_embed_layer(self): - fusion = FusionEmbedLayerNormalization(self) + def fuse_embed_layer(self, use_mask_index): + fusion = FusionEmbedLayerNormalization(self, use_mask_index) fusion.apply() def fuse_layer_norm(self): @@ -398,7 +398,8 @@ class BertOnnxModel(OnnxModel): self.fuse_shape() if (options is None) or options.enable_embed_layer_norm: - self.fuse_embed_layer() + use_mask_index = options.attention_mask_format == AttentionMaskFormat.MaskIndexEnd + self.fuse_embed_layer(use_mask_index) # Remove reshape nodes that having same shape of input and output based on symbolic shape inference. self.utils.remove_useless_reshape_nodes() @@ -437,20 +438,18 @@ class BertOnnxModel(OnnxModel): ops = [ "EmbedLayerNormalization", "Attention", - "QOrderedAttention", "Gelu", - "QOrderedGelu", "FastGelu", "BiasGelu", "GemmFastGelu", "LayerNormalization", - "QOrderedLayerNormalization", "SkipLayerNormalization", - "QOrderedMatMul", ] - for op in ops: + q_ops = ["QOrderedAttention", "QOrderedGelu", "QOrderedLayerNormalization", "QOrderedMatMul"] + for op in ops + q_ops: nodes = self.get_nodes_by_op_type(op) op_count[op] = len(nodes) + logger.info(f"Optimized operators:{op_count}") return op_count diff --git a/onnxruntime/test/python/transformers/test_optimizer.py b/onnxruntime/test/python/transformers/test_optimizer.py index d41f8e08a6..f3edf829a6 100644 --- a/onnxruntime/test/python/transformers/test_optimizer.py +++ b/onnxruntime/test/python/transformers/test_optimizer.py @@ -20,12 +20,14 @@ from transformers import is_tf_available if find_transformers_source(): from benchmark_helper import ConfigModifier, OptimizerInfo, Precision + from fusion_options import FusionOptions from huggingface_models import MODELS from onnx_exporter import export_onnx_model_from_pt, export_onnx_model_from_tf from onnx_model import OnnxModel from optimizer import optimize_model else: from onnxruntime.transformers.benchmark_helper import ConfigModifier, OptimizerInfo, Precision + from onnxruntime.transformers.fusion_options import FusionOptions from onnxruntime.transformers.huggingface_models import MODELS from onnxruntime.transformers.onnx_exporter import export_onnx_model_from_pt, export_onnx_model_from_tf from onnxruntime.transformers.onnx_model import OnnxModel @@ -65,7 +67,7 @@ class TestModelOptimization(unittest.TestCase): self.assertEqual(len(onnx_model.get_nodes_by_op_type(op_type)), count) - # add test function for huggingface pytorch model + # test huggingface pytorch model def _test_optimizer_on_huggingface_model( self, model_name, @@ -73,9 +75,11 @@ class TestModelOptimization(unittest.TestCase): inputs_count=1, validate_model=True, ): - # Remove cached model so that CI machine will have space - shutil.rmtree("./cache_models", ignore_errors=True) + # Remove cached model so that CI machine has enough space. Do not remove cache models in dev machine. + if not find_transformers_source(): + shutil.rmtree("./cache_models", ignore_errors=True) shutil.rmtree("./onnx_models", ignore_errors=True) + # expect fusion result list have the following keys # EmbedLayerNormalization, Attention, Gelu, FastGelu, BiasGelu, LayerNormalization, SkipLayerNormalization model_fusion_statistics = {} @@ -106,12 +110,26 @@ class TestModelOptimization(unittest.TestCase): fusion_options, ) - onnx_model = list(model_fusion_statistics.keys())[0] - fusion_result_list = list(model_fusion_statistics[onnx_model].values()) - if validate_model: self.assertEqual(is_valid_onnx_model, True) - self.assertEqual(fusion_result_list, expected_fusion_result_list) + + expected_node_count = { + "EmbedLayerNormalization": expected_fusion_result_list[0], + "Attention": expected_fusion_result_list[1], + "Gelu": expected_fusion_result_list[2], + "FastGelu": expected_fusion_result_list[3], + "BiasGelu": expected_fusion_result_list[4], + "LayerNormalization": expected_fusion_result_list[5], + "SkipLayerNormalization": expected_fusion_result_list[6], + } + + for _onnx_path, value in model_fusion_statistics.items(): + actual_node_count = value + + for op_type, count in expected_node_count.items(): + if op_type not in actual_node_count or actual_node_count[op_type] != count: + print(f"expected: {expected_node_count} got {actual_node_count}") + self.assertTrue(False) def test_gpt2_past(self): input = _get_test_model_path("gpt2_past") @@ -173,9 +191,12 @@ class TestModelOptimization(unittest.TestCase): onnx_files.append("embed_layer_norm_format3_no_cast.onnx") onnx_files.append("embed_layer_norm_format3_no_cast_opset13.onnx") + options = FusionOptions("bert") + options.use_raw_attention_mask(False) + for file in onnx_files: input_model_path = get_fusion_test_model(file) - model = optimize_model(input_model_path, "bert") + model = optimize_model(input_model_path, "bert", optimization_options=options) expected_node_count = { "EmbedLayerNormalization": 1, "Attention": 1, @@ -270,8 +291,9 @@ class TestTensorflowModelOptimization(unittest.TestCase): self.skipTest("skip TestBertOptimizationTF since tf2onnx not installed") def _test_optimizer_on_tf_model(self, model_name, expected_fusion_result_list, inputs_count, validate_model=True): - # Remove cached model so that CI machine will have space - shutil.rmtree("./cache_models", ignore_errors=True) + # Remove cached model so that CI machine has enough space. Do not remove cache models in dev machine. + if not find_transformers_source(): + shutil.rmtree("./cache_models", ignore_errors=True) shutil.rmtree("./onnx_models", ignore_errors=True) # expect fusion result list have the following keys