From aa9ab565f587f9a590dbed62cd7759c856d1b793 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 15 Apr 2021 00:39:33 -0700 Subject: [PATCH] FastGelu fusion for Megatron model (#7344) * add a fastgelu pattern from Megatron model * update comment * add test --- .../tools/transformers/fusion_biasgelu.py | 1 + .../tools/transformers/fusion_fastgelu.py | 116 +++++++++++++++++- .../transformers/test/test_gelu_fusions.py | 66 ++++++++++ 3 files changed, 179 insertions(+), 4 deletions(-) create mode 100644 onnxruntime/python/tools/transformers/test/test_gelu_fusions.py diff --git a/onnxruntime/python/tools/transformers/fusion_biasgelu.py b/onnxruntime/python/tools/transformers/fusion_biasgelu.py index bf0341ae60..d5d84058bb 100644 --- a/onnxruntime/python/tools/transformers/fusion_biasgelu.py +++ b/onnxruntime/python/tools/transformers/fusion_biasgelu.py @@ -30,6 +30,7 @@ class FusionBiasGelu(Fusion): return (add, matmul) = nodes + bias_weight = None # bias should be one dimension bias_index = -1 for i, input in enumerate(add.input): diff --git a/onnxruntime/python/tools/transformers/fusion_fastgelu.py b/onnxruntime/python/tools/transformers/fusion_fastgelu.py index f678aff040..a3f755fe74 100644 --- a/onnxruntime/python/tools/transformers/fusion_fastgelu.py +++ b/onnxruntime/python/tools/transformers/fusion_fastgelu.py @@ -15,10 +15,15 @@ class FusionFastGelu(Fusion): def __init__(self, model: OnnxModel): super().__init__(model, "FastGelu", "Tanh") - def fuse(self, erf_node, input_name_to_nodes: Dict, output_name_to_node: Dict): - if self.fuse_1(erf_node, input_name_to_nodes, output_name_to_node): + def fuse(self, tanh_node, input_name_to_nodes: Dict, output_name_to_node: Dict): + if self.fuse_1(tanh_node, input_name_to_nodes, output_name_to_node): + return + + if self.fuse_2(tanh_node, input_name_to_nodes, output_name_to_node): + return + + if self.fuse_3(tanh_node, input_name_to_nodes, output_name_to_node): return - self.fuse_2(erf_node, input_name_to_nodes, output_name_to_node) def fuse_1(self, tanh_node, input_name_to_nodes, output_name_to_node) -> Optional[bool]: """ @@ -105,10 +110,11 @@ class FusionFastGelu(Fusion): name=self.model.create_node_name('FastGelu')) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) + return True def fuse_2(self, tanh_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: """ - This pattern is from Tensorflow mode. + This pattern is from Tensorflow model. Fuse Gelu with tanh into one node: +---------------------------+ | | @@ -199,3 +205,105 @@ class FusionFastGelu(Fusion): fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) return True + + + def fuse_3(self, tanh_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: + """ + OpenAI's gelu implementation, also used in Megatron: + Gelu(x) = x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1.0 + 0.044715 * x * x))) + + Fuse subgraph into a FastGelu node: + +------------ Mul (B=0.79788456) -------------------+ + | | + +-------------------------------+ | + | | | + | v v + [root] --> Mul (B=0.044715) --> Mul --> Add(B=1) --> Mul --> Tanh --> Add(B=1) --> Mul--> + | ^ + | | + +-----------> Mul (B=0.5) --------------------------------------------------------+ + """ + if tanh_node.output[0] not in input_name_to_nodes: + return + + children = input_name_to_nodes[tanh_node.output[0]] + if len(children) != 1 or children[0].op_type != 'Add': + return + add_after_tanh = children[0] + + if not self.model.has_constant_input(add_after_tanh, 1.0): + return + + if add_after_tanh.output[0] not in input_name_to_nodes: + return + children = input_name_to_nodes[add_after_tanh.output[0]] + if len(children) != 1 or children[0].op_type != 'Mul': + return + mul_last = children[0] + + mul_half = self.model.match_parent(mul_last, 'Mul', None, output_name_to_node) + if mul_half is None: + return + + i = self.model.find_constant_input(mul_half, 0.5) + if i < 0: + return + + root_input = mul_half.input[0 if i == 1 else 1] + + mul_before_tanh = self.model.match_parent(tanh_node, 'Mul', 0, output_name_to_node) + if mul_before_tanh is None: + return + + add_1 = self.model.match_parent(mul_before_tanh, 'Add', None, output_name_to_node) + if add_1 is None: + return + j = self.model.find_constant_input(add_1, 1.0) + if j < 0: + return + + mul_7978 = self.model.match_parent(mul_before_tanh, 'Mul', None, output_name_to_node) + if mul_7978 is None: + return + k = self.model.find_constant_input(mul_7978, 0.79788456) + if k < 0: + return + if mul_7978.input[0 if k == 1 else 1] != root_input: + return + + mul_before_add_1 = self.model.match_parent(add_1, 'Mul', 0 if j == 1 else 1, output_name_to_node) + if mul_before_add_1 is None: + return + + if mul_before_add_1.input[0] == root_input: + another = 1 + elif mul_before_add_1.input[1] == root_input: + another = 0 + else: + return + + mul_0447 = self.model.match_parent(mul_before_add_1, 'Mul', another, output_name_to_node) + if mul_0447 is None: + return + m = self.model.find_constant_input(mul_0447, 0.044715) + if m < 0: + return + + if mul_0447.input[0 if m == 1 else 1] != root_input: + return + + subgraph_nodes = [ + mul_0447, mul_before_add_1, add_1, mul_before_tanh, tanh_node, add_after_tanh, mul_7978, mul_half, mul_last + ] + if not self.model.is_safe_to_fuse_nodes(subgraph_nodes, [mul_last.output[0]], input_name_to_nodes, + output_name_to_node): + return + + self.nodes_to_remove.extend(subgraph_nodes) + fused_node = helper.make_node('FastGelu', + inputs=[root_input], + outputs=mul_last.output, + name=self.model.create_node_name('FastGelu')) + fused_node.domain = "com.microsoft" + self.nodes_to_add.append(fused_node) + return True diff --git a/onnxruntime/python/tools/transformers/test/test_gelu_fusions.py b/onnxruntime/python/tools/transformers/test/test_gelu_fusions.py new file mode 100644 index 0000000000..fbe4a4bf24 --- /dev/null +++ b/onnxruntime/python/tools/transformers/test/test_gelu_fusions.py @@ -0,0 +1,66 @@ +import os +import sys +import unittest +import math +import torch + + +class HuggingfaceGelu(torch.nn.Module): + def forward(self, x): + return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))) + + +class HuggingfaceFastGelu(torch.nn.Module): + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x))) + + +class MegatronGelu(torch.nn.Module): + def forward(self, x): + # The original implementation using ones_like, which might cause problem for input with dynamic axes in onnx. + # return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + torch.ones_like(x).to(dtype=x.dtype)) + return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + 1.0) + + +class MegatronFastGelu(torch.nn.Module): + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))) + + +test_cases = [ + ('huggingface', 'Gelu', HuggingfaceGelu), + ('huggingface', 'FastGelu', HuggingfaceFastGelu), + ('megatron', 'Gelu', MegatronGelu), + ('megatron', 'FastGelu', MegatronFastGelu) +] + + +class TestGeluFusions(unittest.TestCase): + def verify_node_count(self, bert_model, expected_node_count, test_name): + for op_type, count in expected_node_count.items(): + if len(bert_model.get_nodes_by_op_type(op_type)) != count: + print(f"Counters is not expected in test: {test_name}") + for op, counter in expected_node_count.items(): + print("{}: {} expected={}".format(op, len(bert_model.get_nodes_by_op_type(op)), counter)) + self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count) + + def test_fusions(self): + sys.path.append(os.path.join(os.path.dirname(__file__), '..')) + from optimizer import optimize_model + + for test_case in test_cases: + source, operator, model_class = test_case + model = model_class() + dummy_input = torch.ones(3, dtype=torch.float32) + test_name = f"{operator}_{source}" + onnx_path = f"{test_name}.onnx" + torch.onnx.export(model, (dummy_input), onnx_path, input_names=['input'], output_names=['output']) + optimizer = optimize_model(onnx_path, 'bert') + # optimizer.save_model_to_file(f"{operator}_{source}_opt.onnx") + os.remove(onnx_path) + expected_node_count = {operator: 1} + self.verify_node_count(optimizer, expected_node_count, test_name) + + +if __name__ == '__main__': + unittest.main()