saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update embed layer norm fusion to work with transformers v4.9 (#8914)

Browse source
This commit is contained in:
Tianlei Wu 2021-09-02 19:48:07 -07:00 committed by GitHub
parent e348929019
commit 91f05f387a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 608 additions and 250 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/python/tools/transformers/fusion_base.py
View file

@ -25,6 +25,8 @@ class Fusion:
self.prune_graph: bool = False
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
def apply(self):
logger.debug(f"start {self.description} fusion...")
@ -41,7 +43,7 @@ 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 = op_list.count(self.fused_op_type)
count = max(self.fused_count, op_list.count(self.fused_op_type))
if count > 0:
logger.info(f"Fused {self.description} count: {count}")

684
onnxruntime/python/tools/transformers/fusion_embedlayer.py
View file

@ -3,9 +3,9 @@
# Licensed under the MIT License.
#--------------------------------------------------------------------------
from typing import Dict
from typing import Dict, List, Tuple, Union
from logging import getLogger
from onnx import helper
from onnx import helper, TensorProto, NodeProto
from onnx_model import OnnxModel
from fusion_base import Fusion
from fusion_utils import FusionUtils
@ -15,239 +15,375 @@ logger = getLogger(__name__)
class FusionEmbedLayerNoMask(Fusion):
"""
Embed Layer Normalization will fuse embeddings and mask processing into one node.
The embeddings before conversion:
(input_ids) --------> Gather ----------+ (segment_ids)
| | |
| v v
+--> Shape --> Expand -> Gather---->Add Gather
| ^ | |
| | v v
+---(optional graph) SkipLayerNormalization
Optional graph is used to generate position list (0, 1, ...) per batch. It can be a constant in some model.
(input_ids) --> Gather -----+ Slice
| |
v v
(segment_ids)--> Gather --->Add Reshape
| |
v v
SkipLayerNormalization
Fuse embedding layer into one node (EmbedLayerNormalization).
It supports the following model types: BERT, DistilBert, ALBert.
"""
def __init__(self, model: OnnxModel, description='no mask'):
super().__init__(model, "EmbedLayerNormalization", ["SkipLayerNormalization", "LayerNormalization"],
def __init__(self, model: OnnxModel, description: str = 'no mask'):
super().__init__(model, "EmbedLayerNormalization", ["LayerNormalization", "SkipLayerNormalization"],
description)
self.utils = FusionUtils(model)
self.shape_infer_helper = self.model.infer_runtime_shape({}, update=True)
# The following will be reset in each fuse call of FusionEmbedLayerNormalization
self.attention = None
self.embed_node = None
def match_segment_path(self, normalize_node, input_name_to_nodes, output_name_to_node, input_ids_cast_node):
segment_ids = None
segment_embedding_gather = None
def match_two_gather(self, add: NodeProto) -> Union[None, Tuple[NodeProto, NodeProto]]:
gather_0_path = self.model.match_parent_path(add, ['Gather'], [0])
if gather_0_path is None:
return None
if normalize_node.op_type == "SkipLayerNormalization":
segment_embedding_path = self.model.match_parent_path(normalize_node, ['Gather'], [1])
gather_1_path = self.model.match_parent_path(add, ['Gather'], [1])
if gather_1_path is None:
return None
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("Segment embedding is not found. Embed layer cannot be fused.")
return
_, segment_embedding_gather = segment_embedding_path
else:
segment_embedding_gather = segment_embedding_path[0]
elif normalize_node.op_type == "LayerNormalization":
segment_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Add', 'Gather'], [0, 0, 1])
return gather_0_path[0], gather_1_path[0]
if segment_embedding_path is None:
logger.info("Segment embedding is not found. Embed layer cannot be fused.")
return
_, _, segment_embedding_gather = segment_embedding_path
def check_attention_subgraph(self, layernorm: NodeProto, input_name_to_nodes: Dict[str, List[NodeProto]],
is_distil_bert: bool) -> bool:
"""Check that LayerNormalization has a child of Attention node or subgraph like Attention.
segment_ids = segment_embedding_gather.input[1]
Args:
layernorm (NodeProto): LayerNormalization node
input_name_to_nodes (Dict[str, List[NodeProto]]): map from input name to nodes
is_distil_bert (bool): whether it is DistilBert or not
self.nodes_to_remove.extend(segment_embedding_path)
if self.model.find_graph_input(segment_ids):
casted, segment_ids = self.utils.cast_graph_input_to_int32(segment_ids)
else:
segment_ids, segment_ids_cast_node = self.utils.cast_input_to_int32(segment_ids)
# Cast might be removed by OnnxRuntime.
_, segment_id_path, _ = self.model.match_parent_paths(
segment_ids_cast_node,
[(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], [0, 0, 1, 0, 0, 0]),
(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], output_name_to_node)
if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[0] == segment_id_path[-1].input[0]:
logger.debug("Simplify semgent id path...")
self.model.add_node(
helper.make_node('Shape', inputs=[input_ids_cast_node.input[0]], outputs=["input_shape"]),
self.this_graph_name)
self.model.add_node(
helper.make_node('ConstantOfShape',
inputs=["input_shape"],
outputs=["zeros_for_input_shape"],
value=helper.make_tensor("value", onnx.TensorProto.INT32, [1], [1])),
self.this_graph_name)
segment_ids = "zeros_for_input_shape"
return segment_ids, segment_embedding_gather
def fuse(self, node, input_name_to_nodes, output_name_to_node):
is_distill = False
if self.model.match_parent_path(node, ['Add', 'Gather'], [0, 0]) is None and self.model.match_parent_path(
node, ['Gather'], [0]) is None:
logger.debug(
"Failed to match path SkipLayerNormalization[0] <-- Add <-- Gather or SkipLayerNormalization[0] <-- Gather"
)
if node.op_type != "LayerNormalization" or self.model.match_parent_path(node, ['Add', 'Gather'],
[0, 1]) is None:
return
self.attention = self.model.find_first_child_by_type(node, 'Attention', input_name_to_nodes, recursive=False)
Returns:
bool: whether there is Attention node or subgraph like Attention
"""
self.attention = self.model.find_first_child_by_type(layernorm,
'Attention',
input_name_to_nodes,
recursive=False)
if self.attention is None:
# In case user disables attention fusion, check whether subgraph looks like Attention.
if node.output[0] not in input_name_to_nodes:
return
children = input_name_to_nodes[node.output[0]]
children_types = sorted([child.op_type for child in children])
if children_types != ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization'] and children_types != [
'MatMul', 'MatMul', 'MatMul', 'Shape', 'Shape', 'SkipLayerNormalization'
]:
logger.debug("No Attention like subgraph in children of SkipLayerNormalization")
return
if layernorm.output[0] not in input_name_to_nodes:
return False
children = input_name_to_nodes[layernorm.output[0]]
# Assume the order of embeddings are word_embedding + position_embedding + segment_embedding
normalize_node = node
add_node = None
word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Gather'], [0, 0])
if word_embedding_path is not None:
add_node, word_embedding_gather = word_embedding_path
else:
word_embedding_path = self.model.match_parent_path(normalize_node, ['Add', 'Add', 'Gather'], [0, 0, 0])
if word_embedding_path is not None:
_, add_node, word_embedding_gather = word_embedding_path
# For Albert, there is MatMul+Add after embedding layer before attention.
if len(children) == 1 and children[0].op_type == "MatMul" and children[0].output[0] in input_name_to_nodes:
grandchildren = input_name_to_nodes[children[0].output[0]]
if len(grandchildren) == 1 and grandchildren[0].op_type == "Add" and grandchildren[0].output[
0] in input_name_to_nodes:
nodes = input_name_to_nodes[grandchildren[0].output[0]]
for node in nodes:
if node.op_type == "Attention":
self.attention = node
return True
children_types = sorted([child.op_type for child in nodes])
else:
word_embedding_path = self.model.match_parent_path(normalize_node, ['Gather'], [0])
if word_embedding_path is not None:
word_embedding_gather = word_embedding_path[0]
is_distill = True
from packaging.version import Version
import onnxruntime
if Version(onnxruntime.__version__) <= Version("1.4.0"):
logger.warning(
'Please install onnxruntime with version > 1.4.0 for embedlayer fusion support for distilbert'
)
return
else:
logger.info("Word embedding path is not found. Embed layer cannot be fused.")
return
children_types = sorted([child.op_type for child in children])
# Two Shape nodes might be merged by ORT
if is_distil_bert:
# SkipLayerNormailization might exist when model has been optimized by ORT first.
if children_types != ['MatMul', 'MatMul', 'MatMul', 'Shape', 'SkipLayerNormalization'] and \
children_types != ['Add', 'MatMul', 'MatMul', 'MatMul', 'Shape', 'Shape'] and \
children_types != ['Add', 'MatMul', 'MatMul', 'MatMul', 'Shape']:
logger.debug("No Attention like subgraph in children of LayerNormalization")
return False
else:
if children_types != ['Add', 'MatMul', 'MatMul', 'MatMul'] and \
children_types != ['MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization']:
logger.debug("No Attention like subgraph in children of LayerNormalization")
return False
return True
def match_position_embedding_distilbert(self, position_embedding_gather, input_ids, output_name_to_node):
""" Match position embedding path from input_ids to Gather for DistilBert.
Pattern is like the following:
(input_ids)
|
Shape
| \
| Gather (indices=1)
| |
| Cast (optional)
| |
| Range (start=0, end=*, delta=1)
| |
| Unsqueeze
| /
Expand
|
Gather
"""
path1 = self.model.match_parent_path(position_embedding_gather, ['Expand', 'Shape'], [1, 1])
if path1 is None:
return False
expand, shape = path1
if shape.input[0] != input_ids:
return False
_, path2, _ = self.model.match_parent_paths(expand, [(['Unsqueeze', 'Range', 'Cast', 'Gather', 'Shape'], [0, 0, 1, 0, 0]), \
(['Unsqueeze', 'Range', 'Gather', 'Shape'], [0, 0, 1, 0])], output_name_to_node)
if path2 is None:
return False
range_node = path2[1]
if not (self.utils.check_node_input_value(range_node, 0, 0)
and self.utils.check_node_input_value(range_node, 2, 1)):
return False
gather_node = path2[-2]
if not (self.utils.check_node_input_value(gather_node, 1, 1)):
return False
shape_node = path2[-1]
if shape_node.input[0] != input_ids:
return False
return True
def match_position_embedding_roberta(self, position_embedding_gather, input_ids, output_name_to_node):
""" Match position embedding path from input_ids to Gather for Roberta.
Roberta Embedding Layer Pattern (* is optional since it might be removed by ORT, ? is the padding word id):
(input_ids) --> Equal(B=?) -- Not -- Cast(to=6) -- CumSum(axis=1) -- Mul -- Cast(to=7) -- Add(B=1) -- Cast(to=7)* --> Gather
| ^
V |
+------------------------------+
Roberta new pattern from transformers v4.9:
(input_ids) --> Equal(B=?) -- Not -- Cast(to=6) -- CumSum(axis=1) -- Add(B=0) -- Mul -- Cast(to=7) -- Add(B=1) --> Gather
| ^
V |
+-------------------------------------------+
start_node = position_embedding_gather
start_index = 1
# match optional Cast node.
parent = self.model.get_parent(start_node, start_index, output_name_to_node)
if parent is None:
return
if parent.op_type == "Cast":
if OnnxModel.get_node_attribute(parent, "to") != 7:
return
start_node = parent
start_index = 0
i, path, return_indices = self.model.match_parent_paths(
start_node,
[ (['Add', 'Cast', 'Mul', 'CumSum', 'Cast', 'Not', 'Equal'], [start_index, 0, 0, 0, 0, 0, 0]),
(['Add', 'Cast', 'Mul', 'Add', 'CumSum', 'Cast', 'Not', 'Equal'], [start_index, 0, 0, 0, 0, 0, 0, 0])],
output_name_to_node)
if path is not None:
# constant input of Add shall be 1.
i, value = self.model.get_constant_input(path[0])
if value != 1:
return False
_, self.padding_word_id = self.model.get_constant_input(path[-1])
return input_ids == path[-1].input[0]
"""
return False
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:
(input_ids)
/ \
/ Shape
/ |
/ Gather (indices=1)
/ |
/ Add (optional, B=0)
/ |
Gather (segment_ids) Unsqueeze (axes=0)
\ | |
\ Gather Slice (data[1,512], starts=0, ends=*, axes=1, steps=1)
\ / |
Add Gather
\ /
Add
|
LayerNormalization
"""
path = self.model.match_parent_path(position_embedding_gather, ['Slice', 'Unsqueeze'], [1, 2],
output_name_to_node)
if path is None:
return False
slice, unsqueeze = path
slice_weight = self.model.get_constant_value(slice.input[0])
if not (slice_weight is not None and len(slice_weight.shape) == 2 and slice_weight.shape[0] == 1 \
and self.utils.check_node_input_value(slice, 1, [0]) \
and self.utils.check_node_input_value(slice, 3, [1]) \
and (len(slice.input) == 4 or self.utils.check_node_input_value(slice, 4, [1]))):
return False
opset_version = self.model.get_opset_version()
if opset_version < 13:
if not FusionUtils.check_node_attribute(unsqueeze, 'axes', [0]):
return False
else:
if not self.utils.check_node_input_value(unsqueeze, 1, [0]):
return False
node = self.model.get_parent(unsqueeze, 0, output_name_to_node)
if node is None:
return False
if node.op_type == "Add":
if not self.utils.check_node_input_value(node, 1, 0):
return False
gather = self.model.get_parent(node, 0, output_name_to_node)
else:
gather = node
if gather is None or gather.op_type != "Gather":
return False
if not (self.utils.check_node_input_value(gather, 1, 1)):
return False
shape = self.model.get_parent(gather, 0, output_name_to_node)
if shape is None or shape.op_type != "Shape":
return False
return input_ids == shape.input[0]
def match_position_embedding(self, position_embedding_gather, input_ids, output_name_to_node):
if self.match_position_embedding_bert(position_embedding_gather, input_ids, output_name_to_node):
return True
# TODO: Support roberta (position starts from 2 instead of 0) in EmbedLayerNormalization kernel
# related: https://github.com/huggingface/transformers/issues/10736
#if self.match_position_embedding_roberta(position_embedding_gather, input_ids, output_name_to_node):
# return True
if self.match_position_embedding_distilbert(position_embedding_gather, input_ids, output_name_to_node):
return True
return False
def check_embedding(self, word_embedding_gather, segment_embedding_gather, position_embedding_gather):
"""Sanity check of embedding weights, and match hidden_size of weights and shape of inputs.
"""
input_ids = word_embedding_gather.input[1]
segment_ids = segment_embedding_gather.input[1] if segment_embedding_gather else None
position_ids = position_embedding_gather.input[1]
position_embedding_node_before_gather = None
position_embedding_shape = None
if self.shape_infer_helper is not None:
input_ids_shape = self.shape_infer_helper.get_edge_shape(input_ids)
position_ids_shape = self.shape_infer_helper.get_edge_shape(position_ids)
assert input_ids_shape and position_ids_shape
if not (len(input_ids_shape) == 2 and len(position_ids_shape) == 2
and input_ids_shape[1] == position_ids_shape[1]):
logger.info(
"Cannot fuse EmbedLayerNormalization: input_ids and position_ids not matched in 2nd dimension: {} vs {}"
.format(input_ids_shape, position_ids_shape))
return False
position_embedding_path = self.model.match_parent_path(normalize_node, ['Gather', 'Expand'],
[1, 1]) # for distill-bert
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_node_before_gather = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(normalize_node, ['Reshape', 'Slice'], [1, 0])
if position_embedding_path is not None:
_, position_embedding_weight_node = position_embedding_path
if segment_ids and not self.shape_infer_helper.compare_shape(input_ids, segment_ids):
logger.info(
"Cannot fuse EmbedLayerNormalization: input_ids and segment_ids does not have same shape: {} != {}".
format(input_ids_shape, self.shape_infer_helper.get_edge_shape(segment_ids)))
return False
word_embedding_table = self.model.get_constant_value(word_embedding_gather.input[0])
if word_embedding_table is None or len(word_embedding_table.shape) != 2:
logger.info("Cannot fuse EmbedLayerNormalization: word embedding table is not expected")
return False
position_embedding_table = self.model.get_constant_value(position_embedding_gather.input[0])
if position_embedding_table is None or len(position_embedding_table.shape) != 2 or (
word_embedding_table.shape[1] != position_embedding_table.shape[1]):
logger.info("Cannot fuse EmbedLayerNormalization: position embedding table is not expected")
return False
if segment_ids:
segment_embedding_table = self.model.get_constant_value(segment_embedding_gather.input[0])
if segment_embedding_table is None or len(segment_embedding_table.shape) != 2 or (
word_embedding_table.shape[1] != segment_embedding_table.shape[1]):
logger.info("Cannot fuse EmbedLayerNormalization: segment embedding table is not expected")
return False
# In normal case, word embeding table is the largest, and segment embedding table is the smallest, while postion embedding table is in between.
# TODO: use other information (like initializer names) to identify different embedding weights automatically.
if word_embedding_table.shape[0] <= position_embedding_table.shape[0]:
logger.warn(
f"word_embedding_table ({word_embedding_gather.input[0]}) size {word_embedding_table.shape[0]} <= position_embedding_table ({position_embedding_gather.input[0]}) size {position_embedding_table.shape[0]}"
)
if segment_ids:
if word_embedding_table.shape[0] <= segment_embedding_table.shape[0]:
logger.warn(
f"word_embedding_table ({word_embedding_gather.input[0]}) size {word_embedding_table.shape[0]} <= segment_embedding_table ({segment_embedding_gather.input[0]}) size {segment_embedding_table.shape[0]}"
)
if position_embedding_table.shape[0] <= segment_embedding_table.shape[0]:
logger.warn(
f"position_embedding_table ({position_embedding_gather.input[0]}) size {position_embedding_table.shape[0]} <= segment_embedding_table ({segment_embedding_gather.input[0]}) size {segment_embedding_table.shape[0]}"
)
return True
def cast_to_int32(self, input_name: str) -> Tuple[str, Union[None, NodeProto]]:
"""Cast a graph input or node input to int32.
Args:
input_name (str): name of graph input or node input
Returns:
A tuple of casted input name and the cast node.
int32_output (str): If input is int32, it is the input name, Otherwise it is output name of Cast node.
input_cast_node (Union[None, NodeProto]): Cast node. It could be None if input is int32.
"""
input_cast_node = None
graph_input = self.model.find_graph_input(input_name)
if graph_input is not None:
if graph_input.type.tensor_type.elem_type != TensorProto.INT32:
int32_output, input_cast_node = self.utils.cast_input_to_int32(input_name)
else:
position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand', 'Shape'],
[1, 1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_node_before_gather, position_embedding_shape = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(
add_node, ['Gather', 'Expand', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [1, 1, 1, 1, 0, 0])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_node_before_gather, _, _, _, position_embedding_shape = position_embedding_path
else:
# Here we will not try to get exact match. Instead, we only try identify position embedding weights.
position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Expand'], [1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_node_before_gather = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(add_node, ['Gather', 'Slice'],
[1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_node_before_gather = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(
normalize_node, ['Add', 'Gather', 'Slice'], [0, 1, 1])
if position_embedding_path is not None:
_, position_embedding_weight_node, position_embedding_node_before_gather = position_embedding_path
else:
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:
logger.info("position and word embedding is expected to be applied on same input")
return
if position_embedding_node_before_gather and position_embedding_shape:
input_parent = self.model.get_parent(position_embedding_shape, 0, output_name_to_node)
subgraph_nodes = self.model.get_parent_subgraph_nodes(position_embedding_node_before_gather,
[input_parent] if input_parent else [],
output_name_to_node)
self.nodes_to_remove.extend(subgraph_nodes)
self.nodes_to_remove.extend(word_embedding_path)
self.nodes_to_remove.extend(position_embedding_path)
self.nodes_to_remove.extend([normalize_node])
# Cast input_ids and segment_ids to int32.
input_ids_cast_node = None
if self.model.find_graph_input(input_ids):
casted, input_ids = self.utils.cast_graph_input_to_int32(input_ids)
int32_output = input_name
else:
input_ids, input_ids_cast_node = self.utils.cast_input_to_int32(input_ids)
int32_output, input_cast_node = self.utils.cast_input_to_int32(input_name)
return int32_output, input_cast_node
def create_fused_node(self, input_ids: str, layernorm: NodeProto, word_embedding_gather: NodeProto,
position_embedding_gather: NodeProto, segment_embedding_gather: Union[None, NodeProto]):
"""Create an EmbedLayerNormalization node. Note that segment embedding is optional.
Args:
input_ids (str): input_ids for word embeddings
layernorm (NodeProto): LayerNormalization or SkipLayerNormalization node.
word_embedding_gather (NodeProto): the Gather node for word embedding
position_embedding_gather (NodeProto): the Gather node for position embedding
segment_embedding_gather (Union[None, NodeProto]): the Gather node for segment embedding, or None.
Returns:
NodeProto: the EmbedLayerNormalization node created.
"""
nodes_to_add = []
input_ids, _ = self.cast_to_int32(input_ids)
node_name = self.model.create_node_name('EmbedLayerNormalization')
output_name = node_name + "_output"
if normalize_node.op_type == "LayerNormalization":
gamma = normalize_node.input[1]
beta = normalize_node.input[2]
elif normalize_node.op_type == "SkipLayerNormalization":
gamma = normalize_node.input[2]
beta = normalize_node.input[3]
if layernorm.op_type == "LayerNormalization":
gamma = layernorm.input[1]
beta = layernorm.input[2]
else: # SkipLayerNormalization
gamma = layernorm.input[2]
beta = layernorm.input[3]
embed_node_inputs = None
if is_distill == False:
segment_path = self.match_segment_path(normalize_node, input_name_to_nodes, output_name_to_node,
input_ids_cast_node)
if segment_path is None:
return
else:
segment_ids, segment_embedding_gather = segment_path
if segment_embedding_gather is not None:
segment_ids, _ = self.cast_to_int32(segment_embedding_gather.input[1])
embed_node_inputs = [
input_ids,
segment_ids,
word_embedding_gather.input[0],
position_embedding_weight_node.input[0],
segment_embedding_gather.input[0],
gamma,
beta # gamma and beta
]
else:
embed_node_inputs = [
input_ids,
'',
word_embedding_gather.input[0],
position_embedding_weight_node.input[0],
'',
gamma,
beta # gamma and beta
input_ids, segment_ids, word_embedding_gather.input[0], position_embedding_gather.input[0],
segment_embedding_gather.input[0], gamma, beta
]
else: # no segment embedding
embed_node_inputs = [
input_ids, '', word_embedding_gather.input[0], position_embedding_gather.input[0], '', gamma, beta
]
embed_node = helper.make_node('EmbedLayerNormalization',
@ -258,7 +394,7 @@ class FusionEmbedLayerNoMask(Fusion):
embed_node.domain = "com.microsoft"
# Pass attribute "epsilon" from normalize node to EmbedLayerNormalization.
for att in normalize_node.attribute:
for att in layernorm.attribute:
if att.name == 'epsilon':
embed_node.attribute.extend([att])
# Set default value to 1e-12 if no attribute is found.
@ -266,9 +402,119 @@ class FusionEmbedLayerNoMask(Fusion):
if len(embed_node.attribute) == 0:
embed_node.attribute.extend([helper.make_attribute("epsilon", 1.0E-12)])
self.model.replace_input_of_all_nodes(normalize_node.output[0], output_name)
self.nodes_to_add.append(embed_node)
self.node_name_to_graph_name[embed_node.name] = self.this_graph_name
# Make sure new EmbedLayerNormalization node is the last one in self.nodes_to_add.
nodes_to_add.append(embed_node)
for node in nodes_to_add:
self.node_name_to_graph_name[node.name] = self.this_graph_name
self.nodes_to_add.extend(nodes_to_add)
self.embed_node = embed_node
return embed_node
def finish_fusion(self, layernorm, embed_node):
self.model.replace_input_of_all_nodes(layernorm.output[0], embed_node.output[0])
# use prune graph to remove nodes that is not needed
self.prune_graph = True
def fuse_distilbert(self, layernorm, add_before_layernorm, input_name_to_nodes, output_name_to_node):
"""Fuse embedding layer for DistilBert
Args:
layernorm (NodeProto): node of LayerNormalization or SkipLayerNormalization
add_before_layernorm (NodeProto): the Add node before LayerNormalization, or the SkipLayerNormalization itself
input_name_to_nodes (Dict[str, List[NodeProto]]): map from input name to nodes
output_name_to_node (Dict[str, List[NodeProto]]): map from output name to nodes
"""
# DistilBert has no segment embedding, subgraph pattern is like
# input_ids
# | \
# | (position_embedding_subgraph)
# | |
# Gather Gather
# \ /
# Add
# |
# LayerNormalization
two_gather = self.match_two_gather(add_before_layernorm)
if two_gather is None:
return False
word_embedding_gather, position_embedding_gather = two_gather
input_ids = word_embedding_gather.input[1]
if not self.check_attention_subgraph(layernorm, input_name_to_nodes, is_distil_bert=True):
return False
if not self.match_position_embedding(position_embedding_gather, input_ids, output_name_to_node):
return False
if not self.check_embedding(word_embedding_gather, None, position_embedding_gather):
return False
embed_node = self.create_fused_node(input_ids, layernorm, word_embedding_gather, position_embedding_gather,
None)
self.finish_fusion(layernorm, embed_node)
return True
def fuse_bert(self, layernorm, add_before_layernorm, input_name_to_nodes, output_name_to_node):
"""Fuse embedding layer for Bert
Args:
layernorm (NodeProto): node of LayerNormalization or SkipLayerNormalization
add_before_layernorm (NodeProto): the Add node before LayerNormalization, or the SkipLayerNormalization itself
input_name_to_nodes (Dict[str, List[NodeProto]]): map from input name to nodes
output_name_to_node (Dict[str, List[NodeProto]]): map from output name to nodes
"""
add_2_gather = self.model.match_parent_path(add_before_layernorm, ['Add'], [0])
if add_2_gather is None:
return False
two_gather = self.match_two_gather(add_2_gather[0])
if two_gather is None:
return False
word_embedding_gather, segment_embedding_gather = two_gather
input_ids = word_embedding_gather.input[1]
if not self.check_attention_subgraph(layernorm, input_name_to_nodes, is_distil_bert=False):
return False
position_embedding_path = self.model.match_parent_path(add_before_layernorm, ['Gather'], [1])
if position_embedding_path is None:
return False
position_embedding_gather = position_embedding_path[0]
if not self.match_position_embedding(position_embedding_gather, input_ids, output_name_to_node):
if not self.match_position_embedding(segment_embedding_gather, input_ids, output_name_to_node):
return False
# position and segment are switched
temp = segment_embedding_gather
segment_embedding_gather = position_embedding_gather
position_embedding_gather = temp
if not self.check_embedding(word_embedding_gather, segment_embedding_gather, position_embedding_gather):
return False
embed_node = self.create_fused_node(input_ids, layernorm, word_embedding_gather, position_embedding_gather,
segment_embedding_gather)
self.finish_fusion(layernorm, embed_node)
return True
def fuse(self, node, input_name_to_nodes, output_name_to_node):
if node.op_type == "LayerNormalization":
first_add_path = self.model.match_parent_path(node, ['Add'], [0])
if first_add_path is None:
return
add_before_layernorm = first_add_path[0]
else: # SkipLayerNormalization
add_before_layernorm = node # Add is fused into SkipLayerNormalization
if self.fuse_distilbert(node, add_before_layernorm, input_name_to_nodes, output_name_to_node):
return
if self.fuse_bert(node, add_before_layernorm, input_name_to_nodes, output_name_to_node):
return
class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask):
@ -276,22 +522,18 @@ class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask):
super().__init__(model, "with mask")
def fuse(self, node, input_name_to_nodes, output_name_to_node):
old_count = len(self.nodes_to_add)
# Reset attention and embed_node so that we know fusion is successful when they are not None.
self.attention = None
self.embed_node = None
super().fuse(node, input_name_to_nodes, output_name_to_node)
if len(self.nodes_to_add) == old_count:
return
if self.attention is not None:
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.nodes_to_add.pop()
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
self.nodes_to_add.append(embed_node)
self.prune_graph = True

100
onnxruntime/python/tools/transformers/fusion_shape.py Normal file
View file

@ -0,0 +1,100 @@
#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
from fusion_base import Fusion
from logging import getLogger
from onnx import TensorProto, NodeProto
from onnx_model import OnnxModel
from fusion_utils import FusionUtils
from typing import Union, Dict, List
logger = getLogger(__name__)
class FusionShape(Fusion):
def __init__(self, model: OnnxModel):
super().__init__(model, "Shape", "Concat")
self.utils = FusionUtils(model)
self.shape_infer = None
self.shape_infer_done = False
def get_dimensions_from_tensor_proto(self, tensor_proto: TensorProto) -> Union[int, None]:
if tensor_proto.type.tensor_type.HasField('shape'):
return len(tensor_proto.type.tensor_type.shape.dim)
else:
return None
def get_dimensions(self, input_name: str) -> Union[int, None]:
graph_input = self.model.find_graph_input(input_name)
if graph_input:
return self.get_dimensions_from_tensor_proto(graph_input)
if not self.shape_infer_done:
self.shape_infer = self.model.infer_runtime_shape({}, update=True)
self.shape_infer_done = True
if self.shape_infer is not None:
return self.get_dimensions_from_tensor_proto(self.shape_infer.known_vi_[input_name])
return None
def fuse(self, concat_node: NodeProto, input_name_to_nodes: Dict[str, List[NodeProto]],
output_name_to_node: Dict[str, NodeProto]):
"""
Smplify subgraph like
(2d_input)
/ \
Shape shape
/ \
Gather(indices=0) Gather(indices=1)
| |
Unsqueeze(axes=0) Unsqueeze(axes=0)
\ /
Concat
|
into (2d_input) --> Shape -->
"""
opset_version = self.model.get_opset_version()
inputs = len(concat_node.input)
root = None
shape_output = None
for i in range(inputs):
path = self.model.match_parent_path(concat_node, ['Unsqueeze', 'Gather', 'Shape'], [i, 0, 0],
output_name_to_node)
if path is None:
return
unsqueeze, gather, shape = path
if i == 0:
shape_output = shape.output[0]
if root is None:
root = shape.input[0]
if self.get_dimensions(root) != inputs:
return
elif shape.input[0] != root:
return
if not FusionUtils.check_node_attribute(unsqueeze, 'axis', 0, default_value=0):
return
if opset_version < 13:
if not FusionUtils.check_node_attribute(unsqueeze, 'axes', [0]):
return
else:
if not self.utils.check_node_input_value(unsqueeze, 1, [0]):
return
value = self.model.get_constant_value(gather.input[1])
from numpy import ndarray, array_equal
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.prune_graph = True

3
onnxruntime/python/tools/transformers/fusion_skiplayernorm.py
View file

@ -19,7 +19,8 @@ class FusionSkipLayerNormalization(Fusion):
"""
def __init__(self, model: OnnxModel):
super().__init__(model, "SkipLayerNormalization", "LayerNormalization")
self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7})
# Update shape inference is needed since other fusions might add new edge which does not have shape info yet.
self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True)
def fuse(self, node, input_name_to_nodes, output_name_to_node):
add = self.model.get_parent(node, 0, output_name_to_node)

2
onnxruntime/python/tools/transformers/gpt2_helper.py
View file

@ -786,7 +786,7 @@ class Gpt2Helper:
logger.info(f"Removed the existed directory: {new_dir}")
except OSError as e:
logger.info(f"Failed to remove the directory {new_dir}: {e.strerror}")
# store each model to its own directory (for external data format).
return {
"raw": os.path.join(os.path.join(output_dir, model_name), model_name + ".onnx"),

2
onnxruntime/python/tools/transformers/onnx_exporter.py
View file

@ -9,7 +9,7 @@ import numpy
import os
import torch
from pathlib import Path
from transformers import AutoConfig, AutoTokenizer, AutoModel, LxmertConfig, TransfoXLConfig
from transformers import AutoConfig, AutoTokenizer, LxmertConfig, TransfoXLConfig
from affinity_helper import AffinitySetting
from benchmark_helper import create_onnxruntime_session, Precision
from gpt2_helper import GPT2ModelNoPastState, PRETRAINED_GPT2_MODELS, TFGPT2ModelNoPastState

10
onnxruntime/python/tools/transformers/onnx_model.py
View file

@ -698,13 +698,9 @@ class OnnxModel:
Args:
outputs (list): a list of graph outputs to retain. If it is None, all graph outputs will be kept.
"""
for node in self.model.graph.node:
# Some operators with inner graph in attributes like 'body' 'else_branch' or 'then_branch'
if node.op_type in ['Loop', 'Scan', 'If']:
# TODO: handle inner graph
logger.debug(f"Skip prune_graph since graph has operator: {node.op_type}")
return
if len(self.graphs()) > 1:
logger.debug(f"Skip prune_graph since graph has subgraph")
return
if outputs is None:
outputs = [output.name for output in self.model.graph.output]

11
onnxruntime/python/tools/transformers/onnx_model_bert.py
View file

@ -8,10 +8,11 @@ from typing import List
from onnx import GraphProto, ModelProto, TensorProto, ValueInfoProto, helper
from onnx_model import OnnxModel
from fusion_reshape import FusionReshape
from fusion_shape import FusionShape
from fusion_layernorm import FusionLayerNormalization, FusionLayerNormalizationTF
from fusion_skiplayernorm import FusionSkipLayerNormalization, FusionBiasSkipLayerNormalization
from fusion_embedlayer import FusionEmbedLayerNormalization
from fusion_attention import FusionAttention, AttentionMask, AttentionMaskFormat
from fusion_attention import FusionAttention, AttentionMask
from fusion_gelu import FusionGelu
from fusion_fastgelu import FusionFastGelu
from fusion_biasgelu import FusionBiasGelu
@ -74,6 +75,10 @@ class BertOnnxModel(OnnxModel):
fusion = FusionReshape(self)
fusion.apply()
def fuse_shape(self):
fusion = FusionShape(self)
fusion.apply()
def fuse_embed_layer(self):
fusion = FusionEmbedLayerNormalization(self)
fusion.apply()
@ -313,6 +318,8 @@ class BertOnnxModel(OnnxModel):
self.attention_mask.set_mask_format(options.attention_mask_format)
self.fuse_attention()
self.fuse_shape()
if (options is None) or options.enable_embed_layer_norm:
self.fuse_embed_layer()
@ -340,7 +347,7 @@ class BertOnnxModel(OnnxModel):
if add_dynamic_axes:
self.use_dynamic_axes()
logger.info(f"opset verion: {self.model.opset_import[0].version}")
logger.info(f"opset verion: {self.get_opset_version()}")
def get_fused_operator_statistics(self):
"""

5
onnxruntime/python/tools/transformers/onnx_model_gpt2.py
View file

@ -4,11 +4,6 @@
#--------------------------------------------------------------------------
import logging
import onnx
import sys
import argparse
import numpy as np
from collections import deque
from onnx import ModelProto, TensorProto, numpy_helper
from onnx_model_bert import BertOnnxModel
from fusion_gpt_attention_no_past import FusionGptAttentionNoPast
from fusion_gpt_attention import FusionGptAttention

8
onnxruntime/python/tools/transformers/optimizer.py
View file

@ -19,16 +19,12 @@
import logging
import coloredlogs
import onnx
import os
import sys
import argparse
import numpy as np
from typing import Dict
from collections import deque
from onnx import ModelProto, TensorProto, numpy_helper, load_model
from onnx import load_model
from onnx_model_bart import BartOnnxModel
from onnx_model_bert import BertOnnxModel, BertOptimizationOptions
from onnx_model_bert import BertOnnxModel
from onnx_model_bert_tf import BertOnnxModelTF
from onnx_model_bert_keras import BertOnnxModelKeras
from onnx_model_gpt2 import Gpt2OnnxModel

29
onnxruntime/test/python/transformers/test_optimizer.py
View file

@ -37,13 +37,17 @@ BERT_TEST_MODELS = {
}
def _get_fusion_test_model(file):
relative_path = os.path.join(os.path.dirname(__file__), '..', '..', 'testdata', 'transform', 'fusion', file)
if (os.path.exists(relative_path)):
return relative_path
return os.path.join('.', 'testdata', 'transform', 'fusion', file)
def _get_test_model_path(name):
sub_dir, file = BERT_TEST_MODELS[name]
if sub_dir == "FUSION":
relative_path = os.path.join(os.path.dirname(__file__), '..', '..', 'testdata', 'transform', 'fusion', file)
if (os.path.exists(relative_path)):
return relative_path
return os.path.join('.', 'testdata', 'transform', 'fusion', file)
return _get_fusion_test_model(file)
else:
relative_path = os.path.join(os.path.dirname(__file__), 'test_data', sub_dir, file)
if (os.path.exists(relative_path)):
@ -58,7 +62,8 @@ class TestBertOptimization(unittest.TestCase):
print(f"Counters is not expected in test: {test_name}")
for op, counter in expected_node_count.items():
print("{}: {} expected={}".format(op, len(bert_model.get_nodes_by_op_type(op)), counter))
self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count)
self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count)
# add test function for huggingface pytorch model
def _test_optimizer_on_huggingface_model(self,
@ -198,6 +203,20 @@ class TestBertOptimization(unittest.TestCase):
}
self.verify_node_count(model, expected_node_count, 'test_multiple_embed')
def test_embed_layer_norm_fusion(self):
onnx_files = []
for i in [3, 8, 9]:
onnx_files.append(f"embed_layer_norm_format{i}.onnx")
onnx_files.append(f"embed_layer_norm_format{i}_opset13.onnx")
onnx_files.append('embed_layer_norm_format3_no_cast.onnx')
onnx_files.append('embed_layer_norm_format3_no_cast_opset13.onnx')
for file in onnx_files:
input_model_path = _get_fusion_test_model(file)
model = optimize_model(input_model_path, 'bert')
expected_node_count = {'EmbedLayerNormalization': 1, 'Attention': 1, 'ReduceSum': 0}
self.verify_node_count(model, expected_node_count, file)
# def test_bert_tf2onnx_0(self):
# input = _get_test_model_path('bert_tf2onnx_0')
# model = optimize_model(input, 'bert_tf', num_heads=2, hidden_size=8)