mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
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
This commit is contained in:
Tianlei Wu
GitHub
parent
181628ced1
commit
abe1642a0c
7 changed files with 148 additions and 47 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue