From d88da4406604d7468de9e879f2620dcdc72dd30f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 6 May 2021 09:43:28 -0700 Subject: [PATCH] Allow flexible order of Add inputs in Attention fusion (#7565) --- .../tools/transformers/fusion_attention.py | 33 ++++++++++--------- .../transformers/test/bert_model_generator.py | 29 +++++++++++----- .../test/test_attention_fusion.py | 16 ++++++++- 3 files changed, 53 insertions(+), 25 deletions(-) diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index 657a4e717f..e6f9ec9df2 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -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 diff --git a/onnxruntime/python/tools/transformers/test/bert_model_generator.py b/onnxruntime/python/tools/transformers/test/bert_model_generator.py index 5d1a65f281..79ceec701d 100644 --- a/onnxruntime/python/tools/transformers/test/bert_model_generator.py +++ b/onnxruntime/python/tools/transformers/test/bert_model_generator.py @@ -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") \ No newline at end of file + onnx.save(model, "bert_3d_attention.onnx") diff --git a/onnxruntime/python/tools/transformers/test/test_attention_fusion.py b/onnxruntime/python/tools/transformers/test/test_attention_fusion.py index 2543ce06ce..da70e03ac9 100644 --- a/onnxruntime/python/tools/transformers/test/test_attention_fusion.py +++ b/onnxruntime/python/tools/transformers/test/test_attention_fusion.py @@ -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 = '.'