From 7b289a79275f59438310d63f86f5c4e7bd36baea Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 3 Aug 2021 15:35:19 -0700 Subject: [PATCH] Add test to evaluate Gelu and Fastgelu precision (#8592) * test gelu and fastgelu precision --- .../tools/transformers/fusion_fastgelu.py | 15 +- .../python/transformers/test_parity_gelu.py | 282 ++++++++++++++++++ 2 files changed, 292 insertions(+), 5 deletions(-) create mode 100644 onnxruntime/test/python/transformers/test_parity_gelu.py diff --git a/onnxruntime/python/tools/transformers/fusion_fastgelu.py b/onnxruntime/python/tools/transformers/fusion_fastgelu.py index c4685f4b45..f95b1c3bd9 100644 --- a/onnxruntime/python/tools/transformers/fusion_fastgelu.py +++ b/onnxruntime/python/tools/transformers/fusion_fastgelu.py @@ -62,9 +62,10 @@ class FusionFastGelu(Fusion): if i < 0: return + root_input = mul_half.input[0 if i == 1 else 1] + + #root_node could be None when root_input is graph input root_node = self.model.get_parent(mul_half, 0 if i == 1 else 1, output_name_to_node) - if root_node is None: - return mul_before_tanh = self.model.match_parent(tanh_node, 'Mul', 0, output_name_to_node) if mul_before_tanh is None: @@ -78,7 +79,11 @@ class FusionFastGelu(Fusion): if add_before_tanh is None: return - mul_after_pow = self.model.match_parent(add_before_tanh, 'Mul', None, output_name_to_node, exclude=[root_node]) + mul_after_pow = self.model.match_parent(add_before_tanh, + 'Mul', + None, + output_name_to_node, + exclude=[root_node] if root_node else []) if mul_after_pow is None: return @@ -93,7 +98,7 @@ class FusionFastGelu(Fusion): if not self.model.has_constant_input(pow, 3.0): return - if pow.input[0] != root_node.output[0]: + if pow.input[0] != root_input: return subgraph_nodes = [ @@ -105,7 +110,7 @@ class FusionFastGelu(Fusion): self.nodes_to_remove.extend(subgraph_nodes) fused_node = helper.make_node('FastGelu', - inputs=[root_node.output[0]], + inputs=[root_input], outputs=mul_after_tanh.output, name=self.model.create_node_name('FastGelu')) fused_node.domain = "com.microsoft" diff --git a/onnxruntime/test/python/transformers/test_parity_gelu.py b/onnxruntime/test/python/transformers/test_parity_gelu.py new file mode 100644 index 0000000000..dd3b3f31db --- /dev/null +++ b/onnxruntime/test/python/transformers/test_parity_gelu.py @@ -0,0 +1,282 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +""" +Below are test results for Gelu or FastGelu FP32 kernels using CUDA: + +Formula Input(BeforeCast) MaxDiff MaxDiff(Optimized) +0(gelu_python) FP32 2.38E-07 4.77E-07 +0(gelu_python) FP16 0 6.10E-05 +1(gelu) FP32 4.77E-07 0 +1(gelu) FP16 6.10E-05 0 +2(erf_gelu) FP32 2.38E-07 9.54E-07 +2(erf_gelu) FP16 1.22E-04 1.95E-03 +3(gelu_new) FP32 2.38E-07 2.38E-07 +3(gelu_new) FP16 0 0 +4(gelu_fast) FP32 0 2.38E-07 +4(gelu_fast) FP16 0 3.05E-05 +5(openai_gelu) FP32 0 2.38E-07 +5(openai_gelu) FP16 0 3.05E-05 + +For comparison, CPU has MaxDiff=4.77E-07 for each formula. +""" + +import unittest +import torch +from torch import nn +import numpy +import math +import os + + +class Gelu(nn.Module): + def __init__(self, formula=4): + super().__init__() + self.formula = formula + + def gelu(self, x): + if self.formula == 0: + return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))) + elif self.formula == 1: + return nn.functional.gelu(x) + elif self.formula == 2: + # erf_gelu in Megatron: 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) + elif self.formula == 3: + # gelu_new in huggingface transformers + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + elif self.formula == 4: + # gelu_fast in huggingface transformers with lower precision in a constant (0.7978845608) + return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x))) + else: + # openai_gelu in Megatron + return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))) + + @staticmethod + def get_fused_op(formula): + return "Gelu" if formula in [0, 1, 2] else "FastGelu" + + def forward(self, x): + if x.dtype == torch.float16: + # This test only evaluates FP32 kernels so add data type cast for input and output. + fp16_gelu = self.gelu(x.to(torch.float32)).to(torch.float16) + return (fp16_gelu, ) + else: + fp32_gelu = self.gelu(x) + return (fp32_gelu, ) + + +def create_inputs(batch_size=1, sequence_length=1, hidden_size=768, float16=False, device=torch.device('cuda')): + float_type = torch.float16 if float16 else torch.float32 + input = torch.normal(mean=0.0, std=1.0, size=(batch_size, sequence_length, hidden_size)).to(float_type).to(device) + return input + + +def get_output_names(): + outputs = ["output"] + return outputs + + +def export_onnx(model, onnx_model_path, float16, hidden_size, device): + from pathlib import Path + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + + input_hidden_states = create_inputs(hidden_size=hidden_size, float16=float16, device=device) + with torch.no_grad(): + outputs = model(input_hidden_states) + + dynamic_axes = {'input': {0: 'batch_size', 1: 'seq_len'}, "output": {0: 'batch_size', 1: 'seq_len'}} + + torch.onnx.export(model, + args=(input_hidden_states), + f=onnx_model_path, + input_names=['input'], + output_names=["output"], + dynamic_axes=dynamic_axes, + example_outputs=outputs, + opset_version=11, + do_constant_folding=True) + print("exported:", onnx_model_path) + + +def optimize_onnx(input_onnx_path, optimized_onnx_path, expected_gelu_op_type='Gelu'): + from onnxruntime.transformers.optimizer import optimize_model + onnx_model = optimize_model(input_onnx_path, model_type='gpt2', opt_level=0) + assert len(onnx_model.get_nodes_by_op_type( + expected_gelu_op_type)) == 1, f"Expected {expected_gelu_op_type} node not found in the optimized model" + onnx_model.save_model_to_file(optimized_onnx_path) + + +def diff_outputs(torch_outputs, ort_outputs, index): + """ Returns the maximum difference between PyTorch and OnnxRuntime outputs. + """ + expected_outputs = torch_outputs[index].cpu().numpy() + diff = numpy.abs(expected_outputs - ort_outputs[index]) + return numpy.amax(diff) + + +def compare_outputs(torch_outputs, ort_outputs, atol=1e-06, verbose=True): + """ Returns True if torch and ORT outputs are close for given thresholds, and False otherwise. + """ + same = numpy.asarray([ + numpy.allclose(ort_outputs[i], torch_outputs[i].cpu().numpy(), atol=atol, rtol=0) + for i in range(len(ort_outputs)) + ]) + + max_abs_diff = [diff_outputs(torch_outputs, ort_outputs, i) for i in range(len(ort_outputs))] + + is_all_close = same.all() + if is_all_close: + for i in numpy.where(numpy.logical_not(same))[0]: + diff = numpy.fabs(ort_outputs[i] - torch_outputs[i].cpu().numpy()) + idx = numpy.unravel_index(diff.argmax(), diff.shape) + print( + f'Output {i}, diff={diff[idx]:.9f} index={idx} ort={ort_outputs[i][idx]:.9f} torch={float(torch_outputs[i][idx]):.9f}' + ) + + return is_all_close, max(max_abs_diff) + + +def create_ort_session(onnx_model_path, use_gpu=True): + from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel, __version__ as onnxruntime_version + sess_options = SessionOptions() + sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL + sess_options.intra_op_num_threads = 2 + sess_options.log_severity_level = 2 + execution_providers = ['CPUExecutionProvider'] if not use_gpu else ['CUDAExecutionProvider', 'CPUExecutionProvider'] + return InferenceSession(onnx_model_path, sess_options, providers=execution_providers) + + +def onnxruntime_inference(ort_session, input): + ort_inputs = {'input': numpy.ascontiguousarray(input.cpu().numpy())} + ort_outputs = ort_session.run(None, ort_inputs) + return ort_outputs + + +def run_parity(model, + onnx_model_path, + batch_size, + hidden_size, + sequence_length, + float16, + device, + optimized, + test_cases=100, + verbose=False): + print( + f"optimized={optimized}, onnx_model_path={onnx_model_path}, batch_size={batch_size}, hidden_size={hidden_size}, sequence_length={sequence_length}, float16={float16}, device={device}" + ) + passed_cases = 0 + max_diffs = [] + printed = False # print only one sample + ort_session = create_ort_session(onnx_model_path, device.type == 'cuda') + for i in range(test_cases): + input_hidden_states = create_inputs(batch_size, sequence_length, hidden_size, float16, device) + + with torch.no_grad(): + torch_outputs = model(input_hidden_states) + + ort_outputs = onnxruntime_inference(ort_session, input_hidden_states) + + tolerance = 1e-04 if float16 else 1e-06 + is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, atol=tolerance) + max_diffs.append(max_diff) + if is_all_close: + passed_cases += 1 + elif verbose and not printed: + printed = True + numpy.set_printoptions(precision=10, floatmode='fixed') + torch.set_printoptions(precision=10) + print("input", input_hidden_states) + print("torch_outputs", torch_outputs) + print("ort_outputs", ort_outputs) + + max_diff = max(max_diffs) + diff_count = len([i for i in max_diffs if i > 0]) + success_flag = "[FAILED]" if passed_cases < test_cases else "[OK]" + print(f"{success_flag} Passed_cases={passed_cases}/{test_cases}; Max_diff={max_diff}; Diff_count={diff_count}") + return test_cases - passed_cases + + +def run(batch_size, float16, optimized, hidden_size, device, test_cases, formula=0, sequence_length=2): + test_name = f"batch_size={batch_size}, float16={float16}, optimized={optimized}, hidden_size={hidden_size}, formula={formula}" + print(f"\nTesting ONNX parity: {test_name}") + + model = Gelu(formula=formula) + model.eval() + model.to(device) + if float16: + model.half() + + # Do not re-use onnx file from previous test since weights of model are random. + onnx_model_path = './temp/gelu_{}_{}.onnx'.format(formula, "fp16" if float16 else "fp32") + export_onnx(model, onnx_model_path, float16, hidden_size, device) + + if optimized: + optimized_onnx_path = './temp/gelu_{}_opt_{}.onnx'.format(formula, "fp16" if float16 else "fp32") + optimize_onnx(onnx_model_path, optimized_onnx_path, expected_gelu_op_type=Gelu.get_fused_op(formula)) + onnx_path = optimized_onnx_path + else: + onnx_path = onnx_model_path + + num_failure = run_parity(model, onnx_path, batch_size, hidden_size, sequence_length, float16, device, optimized, + test_cases) + + # clean up onnx file + os.remove(onnx_model_path) + if optimized: + os.remove(onnx_path) + + return num_failure, test_name + + +class TestGeluParity(unittest.TestCase): + def setUp(self): + self.optimized = True # Change it to False if you want to test parity of non optimized ONNX + self.test_cases = 100 # Number of test cases per test run + self.sequence_length = 2 + self.hidden_size = 768 + self.formula_to_test = [0, 1, 3, 4, 5] # formula 2 cannot pass precision test. + + def run_test(self, batch_size, float16, optimized, hidden_size, device, formula): + if float16 and device.type == 'cpu': # CPU does not support FP16 + return + num_failure, test_name = run(batch_size, float16, optimized, hidden_size, device, self.test_cases, formula, + self.sequence_length) + self.assertTrue(num_failure == 0, test_name) + + def run_one(self, optimized, device, hidden_size=768, formula=0): + for batch_size in [4]: + self.run_test(batch_size, + float16=False, + optimized=optimized, + hidden_size=hidden_size, + device=device, + formula=formula) + + self.run_test(batch_size, + float16=True, + optimized=optimized, + hidden_size=hidden_size, + device=device, + formula=formula) + + def test_cpu(self): + cpu = torch.device('cpu') + for i in self.formula_to_test: + self.run_one(self.optimized, cpu, hidden_size=self.hidden_size, formula=i) + + def test_cuda(self): + if not torch.cuda.is_available(): + import pytest + pytest.skip('test requires GPU and torch+cuda') + else: + gpu = torch.device('cuda') + for i in self.formula_to_test: + self.run_one(self.optimized, gpu, hidden_size=self.hidden_size, formula=i) + + +if __name__ == '__main__': + unittest.main()