diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index 97e8dff9af..45bd8c49e9 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/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( diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py index d120dfe7dd..a2d86b84a1 100644 --- a/onnxruntime/python/tools/transformers/optimizer.py +++ b/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() diff --git a/onnxruntime/test/python/transformers/bert_model_generator.py b/onnxruntime/test/python/transformers/bert_model_generator.py index 79ceec701d..dc9a504578 100644 --- a/onnxruntime/test/python/transformers/bert_model_generator.py +++ b/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: diff --git a/onnxruntime/test/python/transformers/test_attention_fusion.py b/onnxruntime/test/python/transformers/test_attention_fusion.py index f7d81272ac..dd4d8628a0 100644 --- a/onnxruntime/test/python/transformers/test_attention_fusion.py +++ b/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() diff --git a/onnxruntime/test/python/transformers/test_data/models/attention_opt.onnx b/onnxruntime/test/python/transformers/test_data/models/attention_opt.onnx new file mode 100644 index 0000000000..ececb8701a Binary files /dev/null and b/onnxruntime/test/python/transformers/test_data/models/attention_opt.onnx differ diff --git a/onnxruntime/test/python/transformers/test_data/models/attention_with_varied_qkv_opt.onnx b/onnxruntime/test/python/transformers/test_data/models/attention_with_varied_qkv_opt.onnx new file mode 100644 index 0000000000..da048bbe5c Binary files /dev/null and b/onnxruntime/test/python/transformers/test_data/models/attention_with_varied_qkv_opt.onnx differ