mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-26 19:52:38 +00:00
e8b8d0d13b
Some initializers are added without raw=True flag. That causes those tensors cannot be saved to external data. If those tensors exceed 2GB in total, optimized model cannot be saved due to protobuf limit. This change will save attention weights and bias in raw data. Note: it is optional to use raw data for shape tensor since they are tiny. ### Motivation and Context https://github.com/microsoft/onnxruntime/issues/17212 https://github.com/microsoft/onnxruntime/issues/15349
475 lines
18 KiB
Python
475 lines
18 KiB
Python
# -------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
# --------------------------------------------------------------------------
|
|
|
|
import logging
|
|
|
|
import onnx
|
|
from onnx import numpy_helper
|
|
from onnx_model_bert_tf import BertOnnxModelTF
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BertOnnxModelKeras(BertOnnxModelTF):
|
|
def __init__(self, model, num_heads, hidden_size):
|
|
super().__init__(model, num_heads, hidden_size)
|
|
|
|
def match_mask_path(self, add_or_sub_before_softmax):
|
|
mask_nodes = self.match_parent_path(
|
|
add_or_sub_before_softmax,
|
|
["Mul", "Sub", "Reshape", "Cast"],
|
|
[1, None, 1, 0],
|
|
)
|
|
if mask_nodes is not None:
|
|
return mask_nodes
|
|
|
|
mask_nodes = self.match_parent_path(
|
|
add_or_sub_before_softmax,
|
|
["Mul", "Sub", "Cast", "Slice", "Unsqueeze"],
|
|
[1, 1, 1, 0, 0],
|
|
)
|
|
if mask_nodes is not None:
|
|
return mask_nodes
|
|
|
|
mask_nodes = self.match_parent_path(
|
|
add_or_sub_before_softmax,
|
|
["Mul", "Sub", "Cast", "Unsqueeze", "Unsqueeze"],
|
|
[1, None, 1, 0, 0],
|
|
)
|
|
return mask_nodes
|
|
|
|
def check_attention_input(self, matmul_q, matmul_k, matmul_v, parent, output_name_to_node):
|
|
reshape_nodes = []
|
|
|
|
for x in [matmul_q, matmul_k, matmul_v]:
|
|
root_input = x.input[0]
|
|
root_node = output_name_to_node[root_input]
|
|
if root_node == parent:
|
|
continue
|
|
if root_node.op_type == "Reshape" and root_node.input[0] == parent.output[0]:
|
|
reshape_nodes.append(root_node)
|
|
continue
|
|
logger.debug(f"Check attention input failed:{root_input}, {parent.output[0]}")
|
|
return False, []
|
|
|
|
return True, reshape_nodes
|
|
|
|
def fuse_attention(self):
|
|
self.input_name_to_nodes()
|
|
output_name_to_node = self.output_name_to_node()
|
|
|
|
nodes_to_remove = []
|
|
attention_count = 0
|
|
|
|
skip_layer_norm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization")
|
|
for normalize_node in skip_layer_norm_nodes:
|
|
# SkipLayerNormalization has two inputs, and one of them is the root input for attention.
|
|
parent = self.get_parent(normalize_node, 0)
|
|
if parent is None or parent.op_type not in [
|
|
"SkipLayerNormalization",
|
|
"EmbedLayerNormalization",
|
|
]:
|
|
if parent.op_type == "Add":
|
|
parent = self.get_parent(normalize_node, 1)
|
|
if parent is None or parent.op_type not in [
|
|
"SkipLayerNormalization",
|
|
"EmbedLayerNormalization",
|
|
]:
|
|
logger.debug(f"First input for skiplayernorm: {parent.op_type if parent is not None else None}")
|
|
continue
|
|
else:
|
|
logger.debug(f"First input for skiplayernorm: {parent.op_type if parent is not None else None}")
|
|
continue
|
|
else:
|
|
# TODO: shall we add back the checking of children op types.
|
|
pass
|
|
|
|
qkv_nodes = self.match_parent_path(
|
|
normalize_node,
|
|
["Add", "Reshape", "MatMul", "Reshape", "Transpose", "MatMul"],
|
|
[None, 0, 0, 0, 0, 0],
|
|
)
|
|
if qkv_nodes is None:
|
|
logger.debug("Failed to match qkv nodes")
|
|
continue
|
|
(
|
|
add,
|
|
extra_reshape_0,
|
|
matmul,
|
|
reshape_qkv,
|
|
transpose_qkv,
|
|
matmul_qkv,
|
|
) = qkv_nodes
|
|
logger.debug("Matched qkv nodes")
|
|
|
|
v_nodes = self.match_parent_path(
|
|
matmul_qkv,
|
|
["Transpose", "Reshape", "Add", "Reshape", "MatMul"],
|
|
[1, 0, 0, 0, 0],
|
|
)
|
|
if v_nodes is None:
|
|
logger.debug("Failed to match v path")
|
|
continue
|
|
(transpose_v, reshape_v, add_v, extra_reshape_1, matmul_v) = v_nodes
|
|
|
|
qk_nodes = self.match_parent_path(matmul_qkv, ["Softmax", "Sub", "MatMul"], [0, 0, 0])
|
|
if qk_nodes is not None:
|
|
(softmax_qk, sub_qk, matmul_qk) = qk_nodes
|
|
q_nodes = self.match_parent_path(
|
|
matmul_qk,
|
|
["Mul", "Transpose", "Reshape", "Add", "Reshape", "MatMul"],
|
|
[0, None, 0, 0, 0, 0],
|
|
)
|
|
if q_nodes is not None:
|
|
(
|
|
mul_q,
|
|
transpose_q,
|
|
reshape_q,
|
|
add_q,
|
|
extra_reshape_2,
|
|
matmul_q,
|
|
) = q_nodes
|
|
|
|
else:
|
|
qk_nodes = self.match_parent_path(matmul_qkv, ["Softmax", "Add", "Mul", "MatMul"], [0, 0, 0, None])
|
|
if qk_nodes is None:
|
|
qk_nodes = self.match_parent_path(matmul_qkv, ["Softmax", "Add", "Div", "MatMul"], [0, 0, 0, None])
|
|
if qk_nodes is None:
|
|
logger.debug("Failed to match qk path")
|
|
continue
|
|
(softmax_qk, add_qk, mul_qk, matmul_qk) = qk_nodes
|
|
|
|
q_nodes = self.match_parent_path(
|
|
matmul_qk,
|
|
["Transpose", "Reshape", "Add", "Reshape", "MatMul"],
|
|
[0, 0, 0, 0, 0],
|
|
)
|
|
if q_nodes is not None:
|
|
(transpose_q, reshape_q, add_q, extra_reshape_2, matmul_q) = q_nodes
|
|
|
|
if q_nodes is None:
|
|
logger.debug("Failed to match q path")
|
|
continue
|
|
|
|
k_nodes = self.match_parent_path(
|
|
matmul_qk,
|
|
["Transpose", "Reshape", "Add", "Reshape", "MatMul"],
|
|
[1, 0, 0, 0, 0],
|
|
)
|
|
if k_nodes is None:
|
|
logger.debug("Failed to match k path")
|
|
continue
|
|
(transpose_k, reshape_k, add_k, extra_reshape_3, matmul_k) = k_nodes
|
|
|
|
mask_nodes = self.match_mask_path(qk_nodes[1])
|
|
if mask_nodes is None:
|
|
logger.debug("Failed to match mask path")
|
|
continue
|
|
if not self.has_constant_input(mask_nodes[1], 1):
|
|
logger.debug("Sub node expected to have an input with constant value 1.0.")
|
|
continue
|
|
|
|
is_same_root, reshape_nodes = self.check_attention_input(
|
|
matmul_q, matmul_k, matmul_v, parent, output_name_to_node
|
|
)
|
|
if is_same_root:
|
|
mask_index = self.attention_mask.process_mask(mask_nodes[-1].input[0])
|
|
logger.debug("Create an Attention node.")
|
|
attention_node = self.attention_fusion.create_attention_node(
|
|
mask_index,
|
|
matmul_q,
|
|
matmul_k,
|
|
matmul_v,
|
|
add_q,
|
|
add_k,
|
|
add_v,
|
|
self.num_heads,
|
|
self.hidden_size,
|
|
parent.output[0],
|
|
reshape_qkv.output[0],
|
|
None,
|
|
)
|
|
if attention_node is None:
|
|
continue
|
|
|
|
self.add_node(attention_node)
|
|
attention_count += 1
|
|
|
|
nodes_to_remove.extend([reshape_qkv, transpose_qkv, matmul_qkv])
|
|
nodes_to_remove.extend(qk_nodes)
|
|
nodes_to_remove.extend(q_nodes)
|
|
nodes_to_remove.extend(k_nodes)
|
|
nodes_to_remove.extend(v_nodes)
|
|
nodes_to_remove.extend(mask_nodes)
|
|
nodes_to_remove.extend(reshape_nodes)
|
|
nodes_to_remove.append(extra_reshape_0)
|
|
self.replace_node_input(add, extra_reshape_0.output[0], matmul.output[0])
|
|
else:
|
|
logger.debug("Root node not matched.")
|
|
continue
|
|
self.remove_nodes(nodes_to_remove)
|
|
self.update_graph()
|
|
logger.info(f"Fused Attention count:{attention_count}")
|
|
|
|
def preprocess(self):
|
|
self.process_embedding()
|
|
self.fuse_mask()
|
|
self.skip_reshape()
|
|
|
|
def skip_reshape(self):
|
|
self.input_name_to_nodes()
|
|
self.output_name_to_node()
|
|
|
|
count = 0
|
|
reshape_nodes = self.get_nodes_by_op_type("Reshape")
|
|
for reshape_node in reshape_nodes:
|
|
parent = self.get_parent(reshape_node, 0)
|
|
if parent is not None and parent.op_type == "Reshape":
|
|
reshape_node.input[0] = parent.input[0]
|
|
count += 1
|
|
|
|
if count > 0:
|
|
logger.info(f"Skip consequent Reshape count: {count}")
|
|
|
|
def fuse_embedding(self, node, output_name_to_node):
|
|
assert node.op_type == "LayerNormalization"
|
|
logger.debug(f"start fusing embedding from node with output={node.output[0]}...")
|
|
word_embed_path = self.match_parent_path(node, ["Add", "Add", "Gather"], [0, 0, 0], output_name_to_node)
|
|
if word_embed_path is None:
|
|
logger.debug("failed to match word_embed_path")
|
|
return False
|
|
|
|
skip_node, add_node, gather_node = word_embed_path
|
|
|
|
word_initializer = self.get_initializer(gather_node.input[0])
|
|
if word_initializer is None:
|
|
logger.debug("failed to get word initializer")
|
|
return False
|
|
|
|
temp = numpy_helper.to_array(word_initializer)
|
|
if len(temp.shape) == 2:
|
|
logger.info(f"Found word embedding. name:{word_initializer.name}, shape:{temp.shape}")
|
|
word_embedding = word_initializer.name
|
|
else:
|
|
logger.info(f"Failed to find word embedding. name:{word_initializer.name}, shape:{temp.shape}")
|
|
return False
|
|
|
|
pos_initializer = self.get_initializer(add_node.input[1])
|
|
if pos_initializer is not None:
|
|
temp = numpy_helper.to_array(pos_initializer)
|
|
if len(temp.shape) == 3 and temp.shape[0] == 1:
|
|
tensor = numpy_helper.from_array(temp.reshape((temp.shape[1], temp.shape[2])), "position_embedding")
|
|
self.add_initializer(tensor)
|
|
logger.info(f"Found position embedding. name:{pos_initializer.name}, shape:{temp.shape[1:]}")
|
|
position_embedding = "position_embedding"
|
|
else:
|
|
logger.info(f"Failed to find position embedding. name:{pos_initializer.name}, shape:{temp.shape}")
|
|
return False
|
|
else:
|
|
pos_embed_path = self.match_parent_path(add_node, ["Gather", "Slice"], [1, 1], output_name_to_node)
|
|
if pos_embed_path is None:
|
|
logger.debug("failed to match pos_embed_path")
|
|
return False
|
|
|
|
pos_gather, pos_slice = pos_embed_path
|
|
pos_initializer = self.get_initializer(pos_gather.input[0])
|
|
if pos_initializer is None:
|
|
logger.debug("failed to get pos initializer")
|
|
return False
|
|
|
|
temp = numpy_helper.to_array(pos_initializer)
|
|
if len(temp.shape) == 2:
|
|
logger.info(f"Found word embedding. name:{pos_initializer.name}, shape:{temp.shape}")
|
|
position_embedding = pos_initializer.name
|
|
else:
|
|
logger.info(f"Failed to find position embedding. name:{pos_initializer.name}, shape:{temp.shape}")
|
|
return False
|
|
|
|
gather = self.get_parent(skip_node, 1, output_name_to_node)
|
|
if gather is None or gather.op_type != "Gather":
|
|
logger.debug("failed to get gather")
|
|
return False
|
|
|
|
segment_initializer = self.get_initializer(gather.input[0])
|
|
if segment_initializer is None:
|
|
logger.debug("failed to get segment initializer")
|
|
return False
|
|
|
|
temp = numpy_helper.to_array(segment_initializer)
|
|
if len(temp.shape) == 2:
|
|
logger.info(f"Found segment embedding. name:{segment_initializer.name}, shape:{temp.shape}")
|
|
segment_embedding = segment_initializer.name
|
|
else:
|
|
logger.info(f"Failed to find segment embedding. name:{segment_initializer.name}, shape:{temp.shape}")
|
|
return False
|
|
|
|
logger.info("Create Embedding node")
|
|
self.create_embedding_subgraph(node, word_embedding, segment_embedding, position_embedding)
|
|
return True
|
|
|
|
def process_embedding(self):
|
|
"""
|
|
Automatically detect word, segment and position embeddings.
|
|
"""
|
|
logger.info("start processing embedding layer...")
|
|
output_name_to_node = self.output_name_to_node()
|
|
for node in self.nodes():
|
|
if node.op_type == "LayerNormalization":
|
|
if self.fuse_embedding(node, output_name_to_node):
|
|
return
|
|
break
|
|
|
|
def fuse_mask(self):
|
|
nodes_to_remove = []
|
|
for node in self.nodes():
|
|
if node.op_type == "Mul" and self.has_constant_input(node, -10000):
|
|
mask_path = self.match_parent_path(node, ["Sub", "Cast", "Slice", "Unsqueeze"], [0, 1, 0, 0])
|
|
if mask_path is None:
|
|
continue
|
|
sub_node, cast_node, slice_node, unsqueeze_node = mask_path
|
|
|
|
mask_input_name = self.attention_mask.get_first_mask()
|
|
if unsqueeze_node.input[0] != mask_input_name:
|
|
print(f"Cast input {unsqueeze_node.input[0]} is not mask input {mask_input_name}")
|
|
continue
|
|
|
|
unsqueeze_added_1 = onnx.helper.make_node(
|
|
"Unsqueeze",
|
|
inputs=[mask_input_name],
|
|
outputs=["mask_fuse_unsqueeze1_output"],
|
|
name="Mask_UnSqueeze_1",
|
|
axes=[1],
|
|
)
|
|
|
|
unsqueeze_added_2 = onnx.helper.make_node(
|
|
"Unsqueeze",
|
|
inputs=["mask_fuse_unsqueeze1_output"],
|
|
outputs=["mask_fuse_unsqueeze2_output"],
|
|
name="Mask_UnSqueeze_2",
|
|
axes=[2],
|
|
)
|
|
|
|
# self.replace_node_input(cast_node, cast_node.input[0], 'mask_fuse_unsqueeze2_output')
|
|
cast_node_2 = onnx.helper.make_node(
|
|
"Cast",
|
|
inputs=["mask_fuse_unsqueeze2_output"],
|
|
outputs=["mask_fuse_cast_output"],
|
|
)
|
|
cast_node_2.attribute.extend([onnx.helper.make_attribute("to", 1)])
|
|
self.replace_node_input(sub_node, sub_node.input[1], "mask_fuse_cast_output")
|
|
|
|
nodes_to_remove.extend([slice_node, unsqueeze_node, cast_node])
|
|
self.add_node(unsqueeze_added_1)
|
|
self.add_node(unsqueeze_added_2)
|
|
self.add_node(cast_node_2)
|
|
|
|
self.remove_nodes(nodes_to_remove)
|
|
|
|
# Prune graph is done after removing nodes to remove island nodes.
|
|
if len(nodes_to_remove) > 0:
|
|
self.prune_graph()
|
|
|
|
logger.info("Fused mask" if len(nodes_to_remove) > 0 else "Failed to fuse mask")
|
|
|
|
def remove_extra_reshape(self):
|
|
skiplayernorm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization")
|
|
reshape_removed = 0
|
|
for skiplayernorm_node in skiplayernorm_nodes:
|
|
path = self.match_parent_path(
|
|
skiplayernorm_node,
|
|
[
|
|
"Add",
|
|
"Reshape",
|
|
"MatMul",
|
|
"Reshape",
|
|
"Gelu",
|
|
"Add",
|
|
"Reshape",
|
|
"MatMul",
|
|
"SkipLayerNormalization",
|
|
],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
)
|
|
if path is None:
|
|
continue
|
|
|
|
(
|
|
add_1,
|
|
reshape_1,
|
|
matmul_1,
|
|
reshape_2,
|
|
gelu,
|
|
add_2,
|
|
reshape_3,
|
|
matmul_2,
|
|
skiplayernorm,
|
|
) = path
|
|
add_2.input[0] = matmul_2.output[0]
|
|
self.remove_node(reshape_3)
|
|
matmul_1.input[0] = gelu.output[0]
|
|
self.remove_node(reshape_2)
|
|
add_1.input[0] = matmul_1.output[0]
|
|
self.remove_node(reshape_1)
|
|
reshape_removed += 3
|
|
|
|
return reshape_removed
|
|
|
|
def remove_extra_reshape_2(self):
|
|
skiplayernorm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization")
|
|
reshape_removed = 0
|
|
for skiplayernorm_node in skiplayernorm_nodes:
|
|
path = self.match_parent_path(
|
|
skiplayernorm_node,
|
|
[
|
|
"Add",
|
|
"Reshape",
|
|
"MatMul",
|
|
"Reshape",
|
|
"Gelu",
|
|
"Add",
|
|
"Reshape",
|
|
"MatMul",
|
|
"Reshape",
|
|
"SkipLayerNormalization",
|
|
],
|
|
[None, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
)
|
|
if path is None:
|
|
continue
|
|
|
|
(
|
|
add_1,
|
|
reshape_1,
|
|
matmul_1,
|
|
reshape_2,
|
|
gelu,
|
|
add_2,
|
|
reshape_3,
|
|
matmul_2,
|
|
reshape_4,
|
|
skiplayernorm,
|
|
) = path
|
|
|
|
matmul_2.input[0] = skiplayernorm.output[0]
|
|
self.remove_node(reshape_4)
|
|
|
|
add_2.input[0] = matmul_2.output[0]
|
|
self.remove_node(reshape_3)
|
|
|
|
matmul_1.input[0] = gelu.output[0]
|
|
self.remove_node(reshape_2)
|
|
|
|
add_1.input[0] = matmul_1.output[0]
|
|
self.remove_node(reshape_1)
|
|
|
|
reshape_removed += 4
|
|
|
|
return reshape_removed
|
|
|
|
def postprocess(self):
|
|
reshape_removed = self.remove_extra_reshape() + self.remove_extra_reshape_2()
|
|
logger.info(f"Remove {reshape_removed} Reshape nodes.")
|
|
|
|
self.prune_graph()
|