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

Allow flexible order of Add inputs in Attention fusion (#7565)

Browse source
This commit is contained in:
Tianlei Wu 2021-05-06 09:43:28 -07:00 committed by GitHub
parent 4896744638
commit d88da44066
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 53 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
onnxruntime/python/tools/transformers/fusion_attention.py
View file

@ -150,9 +150,9 @@ class FusionAttention(Fusion):
q_weight = self.model.get_initializer(q_matmul.input[1])
k_weight = self.model.get_initializer(k_matmul.input[1])
v_weight = self.model.get_initializer(v_matmul.input[1])
q_bias = self.model.get_initializer(q_add.input[1])
k_bias = self.model.get_initializer(k_add.input[1])
v_bias = self.model.get_initializer(v_add.input[1])
q_bias = self.model.get_initializer(q_add.input[1]) or self.model.get_initializer(q_add.input[0])
k_bias = self.model.get_initializer(k_add.input[1]) or self.model.get_initializer(k_add.input[0])
v_bias = self.model.get_initializer(v_add.input[1]) or self.model.get_initializer(v_add.input[0])
if q_weight is None:
print(f"{q_matmul.input[1]} is not initializer. Please set do_constant_folding=True in torch.onnx.export")
@ -166,14 +166,14 @@ class FusionAttention(Fusion):
# Check if all matrices have the same shape
assert qw.shape == kw.shape == vw.shape
# All the matrices have the same shape. For 2d weights, the shapes would be [in_size, out_size].
# All the matrices have the same shape. For 2d weights, the shapes would be [in_size, out_size].
# For 3d weights, shape would be [in_size, a, b] where a*b = out_size
in_size = qw.shape[0]
out_size = np.prod(qw.shape[1:])
qkv_weight = np.stack((qw, kw, vw), axis=1)
qb = NumpyHelper.to_array(q_bias)
qb = NumpyHelper.to_array(q_bias)
kb = NumpyHelper.to_array(k_bias)
vb = NumpyHelper.to_array(v_bias)
@ -233,13 +233,14 @@ class FusionAttention(Fusion):
# SkipLayerNormalization has two inputs, and one of them is the root input for attention.
qkv_nodes = self.model.match_parent_path(start_node, ['Add', 'MatMul', 'Reshape', 'Transpose', 'MatMul'],
[None, 0, 0, 0, 0])
[None, None, 0, 0, 0])
einsum_node = None
if qkv_nodes is not None:
(_, matmul_qkv, reshape_qkv, transpose_qkv, matmul_qkv) = qkv_nodes
else:
# Match Albert
qkv_nodes = self.model.match_parent_path(start_node, ['Add', 'Einsum', 'Transpose', 'MatMul'], [1, 0, 0, 0])
qkv_nodes = self.model.match_parent_path(start_node, ['Add', 'Einsum', 'Transpose', 'MatMul'],
[1, None, 0, 0])
if qkv_nodes is not None:
(_, einsum_node, transpose_qkv, matmul_qkv) = qkv_nodes
else:
@ -284,16 +285,16 @@ class FusionAttention(Fusion):
if children_types.count('MatMul') != 3:
return
v_nodes = self.model.match_parent_path(matmul_qkv, ['Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, 0])
v_nodes = self.model.match_parent_path(matmul_qkv, ['Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, None])
if v_nodes is None:
logger.debug("fuse_attention: failed to match v path")
return
(_, _, add_v, matmul_v) = v_nodes
is_distill = False
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Div', 'MatMul'], [0, 0, 0, 0])
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Div', 'MatMul'], [0, 0, None, 0])
if qk_nodes is None:
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Mul', 'MatMul'], [0, 0, 0, 0])
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Add', 'Mul', 'MatMul'], [0, 0, None, 0])
if qk_nodes is None:
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'MatMul', 'Div'], [0, 0, 2, 0])
is_distill = True
@ -309,10 +310,10 @@ class FusionAttention(Fusion):
else:
(_, add_qk, _, matmul_qk) = qk_nodes
q_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [0, 0, 0, 0])
q_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [0, 0, 0, None])
if q_nodes is None:
q_nodes = self.model.match_parent_path(matmul_qk, ['Div', 'Transpose', 'Reshape', 'Add', 'MatMul'],
[0, 0, 0, 0, 0])
[0, 0, 0, 0, None])
if q_nodes is None:
logger.debug("fuse_attention: failed to match q path")
return
@ -320,10 +321,10 @@ class FusionAttention(Fusion):
add_q = q_nodes[-2]
matmul_q = q_nodes[-1]
k_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, 0])
k_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'], [1, 0, 0, None])
if k_nodes is None:
k_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Transpose', 'Reshape', 'Add', 'MatMul'],
[1, 0, 0, 0, 0])
[1, 0, 0, 0, None])
if k_nodes is None:
logger.debug("fuse_attention: failed to match k path")
return
@ -339,8 +340,8 @@ class FusionAttention(Fusion):
output_name_to_node)
else:
_, mask_nodes, _ = self.model.match_parent_paths(
add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]),
(['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], output_name_to_node)
add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [None, 0, 1, 0, 0]),
(['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [None, 0, 1, 0])], output_name_to_node)
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
return

29
onnxruntime/python/tools/transformers/test/bert_model_generator.py
View file

@ -21,7 +21,17 @@ def float_tensor(name: str, shape: List[int], random=False):
return helper.make_tensor(name, TensorProto.FLOAT, shape, weights)
def create_bert_attention(input_hidden_size=16, pruned_num_heads=2, pruned_head_size=4, use_float_mask=False):
def reverse_if(inputs, reverse=False):
if reverse:
inputs.reverse()
return inputs
def create_bert_attention(input_hidden_size=16,
pruned_num_heads=2,
pruned_head_size=4,
use_float_mask=False,
switch_add_inputs=False):
# unsqueeze in opset version 13 has two inputs (axis is moved from attribute to input).
has_unsqueeze_two_inputs = (version.parse(onnx.__version__) >= version.parse('1.8.0'))
@ -36,13 +46,13 @@ def create_bert_attention(input_hidden_size=16, pruned_num_heads=2, pruned_head_
# q nodes
helper.make_node("MatMul", ["layernorm_out", "matmul_q_weight"], ["matmul_q_out"], "matmul_q"),
helper.make_node("Add", ["matmul_q_out", "add_q_weight"], ["add_q_out"], "add_q"),
helper.make_node("Add", reverse_if(["matmul_q_out", "add_q_weight"], switch_add_inputs), ["add_q_out"], "add_q"),
helper.make_node("Reshape", ["add_q_out", "reshape_weight_1"], ["reshape_q_out"], "reshape_q"),
helper.make_node("Transpose", ["reshape_q_out"], ["transpose_q_out"], "transpose_q", perm=[0, 2, 1, 3]),
# k nodes
helper.make_node("MatMul", ["layernorm_out", "matmul_k_weight"], ["matmul_k_out"], "matmul_k"),
helper.make_node("Add", ["matmul_k_out", "add_k_weight"], ["add_k_out"], "add_k"),
helper.make_node("Add", reverse_if(["matmul_k_out", "add_k_weight"], switch_add_inputs), ["add_k_out"], "add_k"),
helper.make_node("Reshape", ["add_k_out", "reshape_weight_1"], ["reshape_k_out"], "reshape_k"),
helper.make_node("Transpose", ["reshape_k_out"], ["transpose_k_out"], "transpose_k", perm=[0, 2, 3, 1]),
@ -60,7 +70,7 @@ def create_bert_attention(input_hidden_size=16, pruned_num_heads=2, pruned_head_
# qk nodes
helper.make_node("MatMul", ["transpose_q_out", "transpose_k_out"], ["matmul_qk_out"], "matmul_qk"),
helper.make_node("Div", ["matmul_qk_out", "div_weight"], ["div_qk_out"], "div_qk"),
helper.make_node("Add", ["div_qk_out", "mul_mask_out"], ["add_qk_out"], "add_qk"),
helper.make_node("Add", reverse_if(["div_qk_out", "mul_mask_out"], switch_add_inputs), ["add_qk_out"], "add_qk"),
helper.make_node("Softmax", ["add_qk_out"], ["softmax_qk_out"], "softmax_qk", axis=3),
# v nodes
@ -74,8 +84,8 @@ def create_bert_attention(input_hidden_size=16, pruned_num_heads=2, pruned_head_
helper.make_node("Transpose", ["matmul_qkv_1_out"], ["transpose_qkv_out"], "transpose_qkv", perm=[0, 2, 1, 3]),
helper.make_node("Reshape", ["transpose_qkv_out", "reshape_weight_2"], ["reshape_qkv_out"], "reshape_qkv"),
helper.make_node("MatMul", ["reshape_qkv_out", "matmul_qkv_weight"], ["matmul_qkv_2_out"], "matmul_qkv_2"),
helper.make_node("Add", ["matmul_qkv_2_out", "add_qkv_weight"], ["add_qkv_out"], "add_qkv"),
helper.make_node("Add", ["add_qkv_out", "layernorm_out"], ["skip_output"], "add_skip"),
helper.make_node("Add", reverse_if(["matmul_qkv_2_out", "add_qkv_weight"], switch_add_inputs), ["add_qkv_out"], "add_qkv"),
helper.make_node("Add", reverse_if(["add_qkv_out", "layernorm_out"], switch_add_inputs), ["skip_output"], "add_skip"),
helper.make_node("LayerNormalization", ["skip_output", "layer_norm_weight", "layer_norm_bias"], ["output"],
"layernorm2",
axis=-1,
@ -127,6 +137,7 @@ def create_bert_attention(input_hidden_size=16, pruned_num_heads=2, pruned_head_
model = helper.make_model(graph)
return model
def create_tf2onnx_attention_3d(input_hidden_size=16, num_heads=4, head_size=4, use_float_mask=False):
# unsqueeze in opset version 13 has two inputs (axis is moved from attribute to input).
has_unsqueeze_two_inputs = (version.parse(onnx.__version__) >= version.parse('1.8.0'))
@ -143,7 +154,7 @@ def create_tf2onnx_attention_3d(input_hidden_size=16, num_heads=4, head_size=4,
# q nodes
helper.make_node("Einsum", ["layernorm_out", "einsum_q_weight"], ["einsum_q_out"], "einsum_q", equation="abc,cde->abde"),
helper.make_node("Add", ["einsum_q_out", "add_q_weight"], ["add_q_out"], "add_q"),
# k nodes
helper.make_node("Einsum", ["layernorm_out", "einsum_k_weight"], ["einsum_k_out"], "einsum_k", equation="abc,cde->abde"),
helper.make_node("Add", ["einsum_k_out", "add_k_weight"], ["add_k_out"], "add_k"),
@ -229,5 +240,7 @@ def create_tf2onnx_attention_3d(input_hidden_size=16, num_heads=4, head_size=4,
if __name__ == "__main__":
model = create_bert_attention()
onnx.save(model, "pruned_bert_attention.onnx")
model = create_bert_attention(switch_add_inputs=True)
onnx.save(model, "bert_attention_reverse_add_order.onnx")
model = create_tf2onnx_attention_3d()
onnx.save(model, "bert_3d_attention.onnx")
onnx.save(model, "bert_3d_attention.onnx")

16
onnxruntime/python/tools/transformers/test/test_attention_fusion.py
View file

@ -28,7 +28,21 @@ class TestFusion(unittest.TestCase):
'pruned_attention_opt.onnx')
expected = onnx.load(expected_model_path)
self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
def test_attention_fusion_reverse_add_order(self):
model = create_bert_attention(switch_add_inputs=True)
dir = '.'
model_path = os.path.join(dir, "bert_attention_reverse_add_order.onnx")
onnx.save(model, model_path)
optimized_model = optimize_model(model_path)
os.remove(model_path)
# reverse add input order will get same optimized model
expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'fusion',
'pruned_attention_opt.onnx')
expected = onnx.load(expected_model_path)
self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
def test_3d_attention_fusion_tf2onnx_model(self):
model = create_tf2onnx_attention_3d()
dir = '.'