Fuse attention node even in case of different Q,K hidden dimensions (#8106)

* changes to fuse attention node and create varied dimensions

* added an option to optimizer to only do offline fusion

* fixing a typo

* merge with master

* removing extra changes

* added new unit test - test_attention_fusion_for_varied_qkv_dimensions()

* Unit test succesfull for q,k,v paths with varied dimensions

* adding test model for unit test case

* optimizing attention tests

* removing debugs

* minor change

* addressing comments

* addressing comments

* changed the new option to disable_onnxruntime

* replacing asserts with debugs

* make attn fusion backward compatible for head_size, hidden_size

* preserving behavior for shape_modified_tensor

* adding new option as the last parameter

* cleaning up

* line breaks and spaces

* formatting according to python

* making the changes to fuse attention node without user input

* changes to fusion_attention.py updated

* bringing the code up to python standard

onnxruntime/python/tools/transformers/fusion_attention.py

@@ -2,6 +2,8 @@
 # Copyright (c) Microsoft Corporation.  All rights reserved.
 # Licensed under the MIT License.
 #--------------------------------------------------------------------------
+from os import name
+from sys import path
 import numpy as np
 from logging import getLogger
 from enum import Enum
@@ -145,7 +147,11 @@ class FusionAttention(Fusion):
         Returns:
             Union[NodeProto, None]: the node created or None if failed.
         """
-        assert num_heads > 0 and hidden_size > 0 and (hidden_size % num_heads) == 0
+        assert num_heads > 0
+
+        if hidden_size > 0 and (hidden_size % num_heads) != 0:
+            logger.debug(f"input hidden size {hidden_size} is not a multiple of num of heads {num_heads}")
+            return None
 
         q_weight = self.model.get_initializer(q_matmul.input[1])
         k_weight = self.model.get_initializer(k_matmul.input[1])
@@ -163,35 +169,64 @@ class FusionAttention(Fusion):
         kw = NumpyHelper.to_array(k_weight)
         vw = NumpyHelper.to_array(v_weight)
 
-        # Check if all matrices have the same shape
-        assert qw.shape == kw.shape == vw.shape
+        # assert q and k have same shape as expected
+        assert qw.shape == kw.shape
 
-        # All the matrices have the same shape. For 2d weights, the shapes would be [in_size, out_size].
+        qw_in_size = qw.shape[0]
+        kw_in_size = kw.shape[0]
+        vw_in_size = vw.shape[0]
+
+        assert qw_in_size == kw_in_size == vw_in_size
+
+        if hidden_size > 0 and hidden_size != qw_in_size:
+            logger.debug(
+                f"Input hidden size {hidden_size} is not same as weight matrix dimension of q,k,v paths {qw_in_size}, provide correct input hidden size or pass 0"
+            )
+            return None
+
+        is_qkv_diff_dims = False
+        if qw.shape != vw.shape:
+            is_qkv_diff_dims = True
+
+        # All the matrices can have the same shape or q, k matrics can have the same shape with v being different
+        # 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:])
+        qw_out_size = np.prod(qw.shape[1:])
+        kw_out_size = np.prod(qw.shape[1:])
+        vw_out_size = np.prod(vw.shape[1:])
 
-        qkv_weight = np.stack((qw, kw, vw), axis=1)
+        qkv_weight_dim = 0
+        if is_qkv_diff_dims:
+            qkv_weight = np.concatenate((qw, kw, vw), axis=1)
+            qkv_weight_dim = qw_out_size + kw_out_size + vw_out_size
+        else:
+            qkv_weight = np.stack((qw, kw, vw), axis=1)
+            qkv_weight_dim = 3 * qw_out_size
 
         qb = NumpyHelper.to_array(q_bias)
         kb = NumpyHelper.to_array(k_bias)
         vb = NumpyHelper.to_array(v_bias)
 
-        # 1d bias shape: [outsize,]. 2d bias shape: [a, b] where a*b = out_size
-        assert qb.shape == kb.shape == vb.shape
-        assert np.prod(qb.shape) == out_size
+        q_bias_shape = np.prod(qb.shape)
+        k_bias_shape = np.prod(kb.shape)
+        v_bias_shape = np.prod(vb.shape)
 
-        if out_size != hidden_size:
-            logger.debug(
-                f"Shape for weights of Q is {in_size, out_size}, which does not match hidden_size={hidden_size}")
-            return None
+        assert q_bias_shape == k_bias_shape == qw_out_size
+        assert v_bias_shape == vw_out_size
+
+        qkv_bias_dim = 0
+        if is_qkv_diff_dims:
+            qkv_bias = np.concatenate((qb, kb, vb), axis=0)
+            qkv_bias_dim = q_bias_shape + k_bias_shape + v_bias_shape
+        else:
+            qkv_bias = np.stack((qb, kb, vb), axis=0)
+            qkv_bias_dim = 3 * q_bias_shape
 
-        qkv_bias = np.stack((qb, kb, vb), axis=0)
         attention_node_name = self.model.create_node_name('Attention')
 
         weight = helper.make_tensor(name=attention_node_name + '_qkv_weight',
                                     data_type=TensorProto.FLOAT,
-                                    dims=[in_size, 3 * out_size],
+                                    dims=[qw_in_size, qkv_weight_dim],
                                     vals=qkv_weight.flatten().tolist())
 
         # Sometimes weights and bias are stored in fp16
@@ -201,7 +236,7 @@ class FusionAttention(Fusion):
 
         bias = helper.make_tensor(name=attention_node_name + '_qkv_bias',
                                   data_type=TensorProto.FLOAT,
-                                  dims=[3 * out_size],
+                                  dims=[qkv_bias_dim],
                                   vals=qkv_bias.flatten().tolist())
         if q_bias.data_type == 10:
             bias.CopyFrom(numpy_helper.from_array(NumpyHelper.to_array(bias).astype(np.float16), bias.name))
@@ -218,6 +253,10 @@ class FusionAttention(Fusion):
         attention_node.domain = "com.microsoft"
         attention_node.attribute.extend([helper.make_attribute("num_heads", num_heads)])
 
+        if is_qkv_diff_dims:
+            attention_node.attribute.extend(
+                [helper.make_attribute("qkv_hidden_sizes", [qw_out_size, kw_out_size, vw_out_size])])
+
         return attention_node
 
     def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
@@ -297,21 +336,36 @@ class FusionAttention(Fusion):
         (_, _, add_v, matmul_v) = v_nodes
 
         is_distill = False
-        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, None, 0])
+        is_distill_add = False
+        qk_paths = {
+            "path1": (['Softmax', 'Add', 'Div', 'MatMul'], [0, 0, None, 0]),
+            "path2": (['Softmax', 'Add', 'Mul', 'MatMul'], [0, 0, None, 0]),
+            "path3": (['Softmax', 'Where', 'MatMul', 'Div'], [0, 0, 2, 0]),
+            "path4": (['Softmax', 'Add', 'Where', 'MatMul'], [0, 0, 0, 2])
+        }
+
+        qk_nodes = None
+        for k, v in qk_paths.items():
+            qk_nodes = self.model.match_parent_path(matmul_qkv, v[0], v[1])
             if qk_nodes is None:
-                qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'MatMul', 'Div'], [0, 0, 2, 0])
+                continue
+            if k == "path3":
                 is_distill = True
-                if qk_nodes is None:
-                    logger.debug("fuse_attention: failed to match qk path")
-                    return
+            if k == "path4":
+                is_distill_add = True
+            break
+
+        if qk_nodes is None:
+            logger.debug("fuse_attention: failed to match qk path")
+            return
 
         add_qk = None
         matmul_qk = None
         where_qk = None
         if is_distill:
             (_, where_qk, matmul_qk, _) = qk_nodes
+        elif is_distill_add:
+            (_, _, where_qk, matmul_qk) = qk_nodes
         else:
             (_, add_qk, _, matmul_qk) = qk_nodes
 
@@ -343,6 +397,10 @@ class FusionAttention(Fusion):
                                                              [(['Expand', 'Reshape', 'Equal'], [0, 0, 0]),
                                                               (['Cast', 'Expand', 'Reshape', 'Equal'], [0, 0, 0, 0])],
                                                              output_name_to_node)
+        elif is_distill_add:
+            _, mask_nodes, _ = self.model.match_parent_paths(
+                where_qk, [(['Cast', 'Equal', 'Unsqueeze', 'Unsqueeze'], [0, 0, 0, 0]),
+                           (['Equal', 'Unsqueeze', 'Unsqueeze'], [0, 0, 0])], output_name_to_node)
         else:
             _, mask_nodes, _ = self.model.match_parent_paths(
                 add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [None, 0, 1, 0, 0]),
@@ -351,18 +409,17 @@ class FusionAttention(Fusion):
             logger.debug("fuse_attention: failed to match mask path")
             return
 
-        if matmul_v.input[0] == root_input and matmul_q.input[0] == root_input and matmul_v.input[0] == root_input:
+        if matmul_v.input[0] == root_input and matmul_q.input[0] == root_input and matmul_k.input[0] == root_input:
             mask_index = self.attention_mask.process_mask(mask_nodes[-1].input[0])
 
             attention_last_node = reshape_qkv if einsum_node is None else transpose_qkv
 
-            num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q)
-            if num_heads <= 0 or hidden_size <= 0 or (hidden_size % num_heads) != 0:
-                logger.debug("fuse_attention: failed to detect num_heads or hidden_size")
-                return
-
+            q_num_heads, q_hidden_size = self.get_num_heads_and_hidden_size(reshape_q)
+            # number of heads are same for all the paths, hence to create attention node, we pass the q_num_heads
+            # the input_hidden_size represents the input hidden size, this is used as needed but hidden sizes for Q, K are extracted appropriately
             new_node = self.create_attention_node(mask_index, matmul_q, matmul_k, matmul_v, add_q, add_k, add_v,
-                                                  num_heads, hidden_size, root_input, attention_last_node.output[0])
+                                                  q_num_heads, self.hidden_size, root_input,
+                                                  attention_last_node.output[0])
             if new_node is None:
                 return
 
@@ -375,8 +432,8 @@ class FusionAttention(Fusion):
                 shape_tensor = helper.make_tensor(name="shape_modified_tensor" + unique_index,
                                                   data_type=TensorProto.INT64,
                                                   dims=[4],
-                                                  vals=np.int64([0, 0, num_heads,
-                                                                 int(hidden_size / num_heads)]).tobytes(),
+                                                  vals=np.int64([0, 0, q_num_heads,
+                                                                 int(q_hidden_size / q_num_heads)]).tobytes(),
                                                   raw=True)
                 self.model.add_initializer(shape_tensor, self.this_graph_name)
                 self.model.add_node(

onnxruntime/python/tools/transformers/optimizer.py

@@ -40,7 +40,7 @@ MODEL_CLASSES = {
     "bert_tf": (BertOnnxModelTF, "tf2onnx", False),
     "bert_keras": (BertOnnxModelKeras, "keras2onnx", False),
     "gpt2": (Gpt2OnnxModel, "pytorch", True),
-    "gpt2_tf": (Gpt2OnnxModel, 'tf2onnx', False) # might add a class for GPT2OnnxModel for TF later. 
+    "gpt2_tf": (Gpt2OnnxModel, 'tf2onnx', False)  # might add a class for GPT2OnnxModel for TF later.
 }
 
 
@@ -214,6 +214,12 @@ def _parse_arguments():
     parser.add_argument('--only_onnxruntime', required=False, action='store_true', help="optimized by onnxruntime only")
     parser.set_defaults(only_onnxruntime=False)
 
+    parser.add_argument('--disable_onnxruntime',
+                        required=False,
+                        action='store_true',
+                        help="do not use onnxruntime to optimize")
+    parser.set_defaults(disable_onnxruntime=False)
+
     parser.add_argument('--opt_level',
                         required=False,
                         type=int,
@@ -265,7 +271,8 @@ def optimize_model(input,
                    optimization_options=None,
                    opt_level=0,
                    use_gpu=False,
-                   only_onnxruntime=False):
+                   only_onnxruntime=False,
+                   disable_onnxruntime=False):
     """ Optimize Model by OnnxRuntime and/or offline fusion logic.
 
     The following optimizes model by OnnxRuntime only, and no offline fusion logic:
@@ -282,6 +289,7 @@ def optimize_model(input,
         opt_level (int): onnxruntime graph optimization level (0, 1, 2 or 99). When the level > 0, onnxruntime will be used to optimize model first.
         use_gpu (bool): use gpu or not for onnxruntime.
         only_onnxruntime (bool): only use onnxruntime to optimize model, and no offline fusion logic is used.
+        disable_onnxruntime (bool): only use offline fusion logic to optimize model.
 
      Returns:
         object of an optimizer class.
@@ -289,12 +297,17 @@ def optimize_model(input,
     (optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type]
 
     temp_model_path = None
-    if opt_level > 1:  # Optimization specified for an execution provider.
-        temp_model_path = optimize_by_onnxruntime(input, use_gpu=use_gpu, opt_level=opt_level)
-    elif run_onnxruntime:
-        # Use Onnxruntime to do optimizations (like constant folding and cast elimation) that is not specified to exection provider.
-        # CPU provider is used here so that there is no extra node for GPU memory copy.
-        temp_model_path = optimize_by_onnxruntime(input, use_gpu=False, opt_level=1)
+
+    if disable_onnxruntime and only_onnxruntime:
+        logger.warning("Only one of the options can be true in disable_onnxruntime or only_onnxruntime")
+
+    if disable_onnxruntime is False:
+        if opt_level > 1:  # Optimization specified for an execution provider.
+            temp_model_path = optimize_by_onnxruntime(input, use_gpu=use_gpu, opt_level=opt_level)
+        elif run_onnxruntime:
+            # Use Onnxruntime to do optimizations (like constant folding and cast elimation) that is not specified to exection provider.
+            # CPU provider is used here so that there is no extra node for GPU memory copy.
+            temp_model_path = optimize_by_onnxruntime(input, use_gpu=False, opt_level=1)
 
     model = load_model(temp_model_path or input, format=None, load_external_data=True)
 
@@ -347,7 +360,8 @@ def main():
                                opt_level=args.opt_level,
                                optimization_options=optimization_options,
                                use_gpu=args.use_gpu,
-                               only_onnxruntime=args.only_onnxruntime)
+                               only_onnxruntime=args.only_onnxruntime,
+                               disable_onnxruntime=args.disable_onnxruntime)
 
     if args.float16:
         optimizer.convert_model_float32_to_float16()

onnxruntime/test/python/transformers/bert_model_generator.py

@@ -28,8 +28,9 @@ def reverse_if(inputs, reverse=False):
 
 
 def create_bert_attention(input_hidden_size=16,
-                          pruned_num_heads=2,
-                          pruned_head_size=4,
+                          num_heads=2,
+                          pruned_qk_hidden_size=16,
+                          pruned_v_hidden_size=16,
                           use_float_mask=False,
                           switch_add_inputs=False):
     # unsqueeze in opset version 13 has two inputs (axis is moved from attribute to input).
@@ -47,13 +48,13 @@ def create_bert_attention(input_hidden_size=16,
         # q nodes
         helper.make_node("MatMul", ["layernorm_out", "matmul_q_weight"], ["matmul_q_out"], "matmul_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("Reshape", ["add_q_out", "reshape_weight_qk"], ["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", 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("Reshape", ["add_k_out", "reshape_weight_qk"], ["reshape_k_out"], "reshape_k"),
         helper.make_node("Transpose", ["reshape_k_out"], ["transpose_k_out"], "transpose_k", perm=[0, 2, 3, 1]),
 
         # mask nodes
@@ -76,13 +77,13 @@ def create_bert_attention(input_hidden_size=16,
         # v nodes
         helper.make_node("MatMul", ["layernorm_out", "matmul_v_weight"], ["matmul_v_out"], "matmul_v"),
         helper.make_node("Add", ["matmul_v_out", "add_v_weight"], ["add_v_out"], "add_v"),
-        helper.make_node("Reshape", ["add_v_out", "reshape_weight_1"], ["reshape_v_out"], "reshape_v"),
+        helper.make_node("Reshape", ["add_v_out", "reshape_weight_v"], ["reshape_v_out"], "reshape_v"),
         helper.make_node("Transpose", ["reshape_v_out"], ["transpose_v_out"], "transpose_v", perm=[0, 2, 1, 3]),
 
         # qkv nodes
         helper.make_node("MatMul", ["softmax_qk_out", "transpose_v_out"], ["matmul_qkv_1_out"], "matmul_qkv_1"),
         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("Reshape", ["transpose_qkv_out", "reshape_weight_qkv"], ["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", 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"),
@@ -92,23 +93,25 @@ def create_bert_attention(input_hidden_size=16,
                          epsion=0.000009999999747378752),
     ]
 
-    pruned_hidden_size = pruned_num_heads * pruned_head_size
+    pruned_qk_head_size = int(pruned_qk_hidden_size / num_heads)
+    pruned_v_head_size = int(pruned_v_hidden_size / num_heads)
     initializers = [  # initializers
         float_tensor('layer_norm_weight', [input_hidden_size]),
         float_tensor('layer_norm_bias', [input_hidden_size]),
-        float_tensor('matmul_q_weight', [input_hidden_size, pruned_hidden_size]),
-        float_tensor('matmul_k_weight', [input_hidden_size, pruned_hidden_size]),
-        float_tensor('matmul_v_weight', [input_hidden_size, pruned_hidden_size]),
-        float_tensor('matmul_qkv_weight', [pruned_hidden_size, input_hidden_size]),
-        float_tensor('add_q_weight', [pruned_hidden_size]),
-        float_tensor('add_k_weight', [pruned_hidden_size]),
-        float_tensor('add_v_weight', [pruned_hidden_size]),
+        float_tensor('matmul_q_weight', [input_hidden_size, pruned_qk_hidden_size]),
+        float_tensor('matmul_k_weight', [input_hidden_size, pruned_qk_hidden_size]),
+        float_tensor('matmul_v_weight', [input_hidden_size, pruned_v_hidden_size]),
+        float_tensor('matmul_qkv_weight', [pruned_v_hidden_size, input_hidden_size]),
+        float_tensor('add_q_weight', [pruned_qk_hidden_size]),
+        float_tensor('add_k_weight', [pruned_qk_hidden_size]),
+        float_tensor('add_v_weight', [pruned_v_hidden_size]),
         float_tensor('add_qkv_weight', [input_hidden_size]),
-        helper.make_tensor('div_weight', TensorProto.FLOAT, [1], [math.sqrt(pruned_head_size)]),
+        helper.make_tensor('div_weight', TensorProto.FLOAT, [1], [math.sqrt(pruned_qk_head_size)]),
         helper.make_tensor('sub_weight', TensorProto.FLOAT, [1], [1.0]),
         helper.make_tensor('mul_weight', TensorProto.FLOAT, [1], [-10000]),
-        helper.make_tensor('reshape_weight_1', TensorProto.INT64, [4], [0, 0, pruned_num_heads, pruned_head_size]),
-        helper.make_tensor('reshape_weight_2', TensorProto.INT64, [3], [0, 0, pruned_hidden_size]),
+        helper.make_tensor('reshape_weight_qk', TensorProto.INT64, [4], [0, 0, num_heads, pruned_qk_head_size]),
+        helper.make_tensor('reshape_weight_v', TensorProto.INT64, [4], [0, 0, num_heads, pruned_v_head_size]),
+        helper.make_tensor('reshape_weight_qkv', TensorProto.INT64, [3], [0, 0, pruned_v_hidden_size]),
     ]
 
     if has_unsqueeze_two_inputs:

onnxruntime/test/python/transformers/test_attention_fusion.py

@@ -14,10 +14,26 @@ from bert_model_generator import create_bert_attention, create_tf2onnx_attention
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 from onnxruntime.transformers.optimizer import optimize_model
 
+
 class TestFusion(unittest.TestCase):
-    def test_attention_fusion_pruned_model(self):
+    def test_attention_fusion(self):
         model = create_bert_attention()
         dir = '.'
+        model_path = os.path.join(dir, "attention.onnx")
+        onnx.save(model, model_path)
+        optimized_model = optimize_model(model_path)
+        os.remove(model_path)
+
+        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'models', 'attention_opt.onnx')
+        expected = onnx.load(expected_model_path)
+        self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
+
+    def test_attention_fusion_pruned_model(self):
+        model = create_bert_attention(input_hidden_size=16,
+                                      num_heads=2,
+                                      pruned_qk_hidden_size=8,
+                                      pruned_v_hidden_size=8)
+        dir = '.'
         model_path = os.path.join(dir, "pruned_attention.onnx")
         onnx.save(model, model_path)
         optimized_model = optimize_model(model_path)
@@ -29,7 +45,11 @@ class TestFusion(unittest.TestCase):
         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)
+        model = create_bert_attention(input_hidden_size=16,
+                                      num_heads=2,
+                                      pruned_qk_hidden_size=8,
+                                      pruned_v_hidden_size=8,
+                                      switch_add_inputs=True)
         dir = '.'
         model_path = os.path.join(dir, "bert_attention_reverse_add_order.onnx")
         onnx.save(model, model_path)
@@ -42,6 +62,22 @@ class TestFusion(unittest.TestCase):
         expected = onnx.load(expected_model_path)
         self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
 
+    def test_attention_fusion_for_varied_qkv_dimensions(self):
+        model = create_bert_attention(input_hidden_size=16,
+                                      num_heads=2,
+                                      pruned_qk_hidden_size=24,
+                                      pruned_v_hidden_size=16)
+        dir = '.'
+        model_path = os.path.join(dir, "attention_with_varied_qkv.onnx")
+        onnx.save(model, model_path)
+        optimized_model = optimize_model(model_path)
+        os.remove(model_path)
+
+        expected_model_path = os.path.join(os.path.dirname(__file__), 'test_data', 'models',
+                                           'attention_with_varied_qkv_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 = '.'
@@ -55,5 +91,6 @@ class TestFusion(unittest.TestCase):
         expected = onnx.load(expected_model_path)
         self.assertEqual(str(optimized_model.model.graph), str(expected.graph))
 
+
 if __name__ == '__main__':
     unittest.main()