mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
FastGelu fusion for Megatron model (#7344)
* add a fastgelu pattern from Megatron model * update comment * add test
This commit is contained in:
Tianlei Wu
GitHub
parent
0da085ed48
commit
aa9ab565f5
3 changed files with 179 additions and 4 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
Loading…
Reference in a new issue