diff --git a/onnxruntime/python/tools/transformers/fusion_layernorm.py b/onnxruntime/python/tools/transformers/fusion_layernorm.py index b3e2f608fb..a1f33140f6 100644 --- a/onnxruntime/python/tools/transformers/fusion_layernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_layernorm.py @@ -41,15 +41,13 @@ class FusionLayerNormalization(Fusion): if len(children) == 0 or len(children) > 2: return - parent = self.model.get_parent(node, 0, output_name_to_node) - if parent is None: - return - - if children[0].op_type != 'Sub' or self.model.get_parent(children[0], 0, output_name_to_node) != parent: + root_input = node.input[0] + + if children[0].op_type != 'Sub' or children[0].input[0] != root_input: return if len(children) == 2: - if children[1].op_type != 'Sub' or self.model.get_parent(children[1], 0, output_name_to_node) != parent: + if children[1].op_type != 'Sub' or children[1].input[0] != root_input: return div_node = None diff --git a/onnxruntime/test/python/transformers/parity_utilities.py b/onnxruntime/test/python/transformers/parity_utilities.py new file mode 100644 index 0000000000..cde4f8f46d --- /dev/null +++ b/onnxruntime/test/python/transformers/parity_utilities.py @@ -0,0 +1,156 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# ------------------------------------------------------------------------- + +import torch +from torch import nn +import numpy +import os +import sys + + +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 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_op=None): + # Try import optimizer from source directory so that we need not build and install package after making change. + source_dir = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'python', 'tools', 'transformers') + if (os.path.exists(source_dir) and source_dir not in sys.path): + sys.path.append(source_dir) + from optimizer import optimize_model + else: + from onnxruntime.transformers.optimizer import optimize_model + + onnx_model = optimize_model(input_onnx_path, model_type='gpt2', opt_level=0) + onnx_model.save_model_to_file(optimized_onnx_path) + + if expected_op is not None: + assert len(onnx_model.get_nodes_by_op_type(expected_op)) == 1, \ + f"Expected {expected_op} node not found in the optimized model {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): + """Compare outputs from PyTorch and OnnxRuntime + + Args: + torch_outputs (Tuple[Torch.Tensor]): PyTorch model output + ort_outputs (List[numpy.ndarray]): OnnxRuntime output + atol (float, optional): Absolute tollerance. Defaults to 1e-06. + verbose (bool, optional): Print more information. Defaults to True. + + Returns: + is_all_close(bool): whether all elements are close. + max_abs_diff(float): maximum absolute difference. + """ + 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 (not is_all_close) and verbose: + 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, + tolerance=None): + 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) + + if tolerance is None: + tolerance = 1e-03 if float16 else 1e-05 + is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, atol=tolerance, verbose=verbose) + 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 diff --git a/onnxruntime/test/python/transformers/test_parity_gelu.py b/onnxruntime/test/python/transformers/test_parity_gelu.py index dd3b3f31db..2c06fdbef7 100644 --- a/onnxruntime/test/python/transformers/test_parity_gelu.py +++ b/onnxruntime/test/python/transformers/test_parity_gelu.py @@ -26,9 +26,9 @@ 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 +from parity_utilities import * class Gelu(nn.Module): @@ -68,141 +68,14 @@ class Gelu(nn.Module): 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}") + test_name = f"device={device}, float16={float16}, optimized={optimized}, batch_size={batch_size}, sequence_length={sequence_length}, hidden_size={hidden_size}, formula={formula}" + print(f"\nTesting: {test_name}") model = Gelu(formula=formula) model.eval() @@ -216,7 +89,7 @@ def run(batch_size, float16, optimized, hidden_size, device, test_cases, formula 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)) + optimize_onnx(onnx_model_path, optimized_onnx_path, Gelu.get_fused_op(formula)) onnx_path = optimized_onnx_path else: onnx_path = onnx_model_path @@ -238,14 +111,16 @@ class TestGeluParity(unittest.TestCase): 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. + self.formula_to_test = [0, 1, 2, 3, 4, 5] + self.formula_must_pass = [0, 1, 3, 4, 5] # formula 2 cannot pass precision test. - def run_test(self, batch_size, float16, optimized, hidden_size, device, formula): + def run_test(self, batch_size, float16, optimized, hidden_size, device, formula, enable_assert=True): 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) + if enable_assert: + self.assertTrue(num_failure == 0, "Failed: " + test_name) def run_one(self, optimized, device, hidden_size=768, formula=0): for batch_size in [4]: @@ -254,14 +129,16 @@ class TestGeluParity(unittest.TestCase): optimized=optimized, hidden_size=hidden_size, device=device, - formula=formula) + formula=formula, + enable_assert=formula in self.formula_must_pass) self.run_test(batch_size, float16=True, optimized=optimized, hidden_size=hidden_size, device=device, - formula=formula) + formula=formula, + enable_assert=formula in self.formula_must_pass) def test_cpu(self): cpu = torch.device('cpu') diff --git a/onnxruntime/test/python/transformers/test_parity_huggingface_gpt_attention.py b/onnxruntime/test/python/transformers/test_parity_huggingface_gpt_attention.py index 023fff85dd..dc8beeb350 100644 --- a/onnxruntime/test/python/transformers/test_parity_huggingface_gpt_attention.py +++ b/onnxruntime/test/python/transformers/test_parity_huggingface_gpt_attention.py @@ -20,6 +20,7 @@ import onnx import numpy import os from transformers.modeling_utils import Conv1D +from parity_utilities import diff_outputs, create_ort_session, compare_outputs DEBUG_OUTPUTS = ["qk", "norm_qk", "softmax", "attn_weights"] @@ -284,51 +285,6 @@ def optimize_onnx(input_onnx_path, optimized_onnx_path, num_heads, debug): onnx_model.save_model_to_file(optimized_onnx_path) -def diff_outputs(torch_outputs, ort_outputs, index, relative=False): - """ Returns the maximum difference between PyTorch and OnnxRuntime outputs. - """ - expected_outputs = torch_outputs[index].cpu().numpy() - diff = numpy.abs(expected_outputs - ort_outputs[index]) - if relative: - return numpy.amax(diff / (numpy.abs(expected_outputs) + 1e-6)) - else: - return numpy.amax(diff) - - -def compare_outputs(torch_outputs, ort_outputs, rtol=1e-03, atol=1e-03, verbose=False): - """ Returns True if torch and ORT outputs are close for given thresholds, and False otherwise. - """ - is_all_close = True - max_abs_diff = [] - for i in range(len(ort_outputs)): - is_close = numpy.allclose(ort_outputs[i], torch_outputs[i].cpu().numpy(), rtol=rtol, atol=atol) - if not is_close: - is_all_close = False - if verbose: - print(f'output {i} ({ort_outputs[i].name}) are close: {is_close}') - max_abs_diff.append(diff_outputs(torch_outputs, ort_outputs, i)) - - if (not is_all_close) or (verbose and max(max_abs_diff) > 0): - messages = ["max_abs_diff per output:"] - output_names = ["attn_output", "present"] + DEBUG_OUTPUTS - for i, diff in enumerate(max_abs_diff): - messages.append(f"{output_names[i]}={diff},") - print(" ".join(messages)) - - 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_ENABLE_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_hidden_states, attention_mask, past): ort_inputs = { 'past': numpy.ascontiguousarray(past.cpu().numpy()), @@ -370,7 +326,7 @@ def verify_attention(model, ort_outputs = onnxruntime_inference(ort_session, input_hidden_states, attention_mask, layer_past) tolerance = 1e-03 if float16 else 1e-05 - is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, rtol=tolerance, atol=tolerance) + is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, atol=tolerance, verbose=True) max_diffs.append(max_diff) if is_all_close: passed_cases += 1 @@ -410,52 +366,22 @@ def run(batch_size, float16, optimized, hidden_size, num_attention_heads, device past_sequence_length = 0 padding_length = 0 num_failure = 0 - num_failure += verify_attention(model, - onnx_path, - batch_size, - hidden_size, - num_attention_heads, - sequence_length, - past_sequence_length, - float16, - device, - padding_length, - optimized, - test_cases) + num_failure += verify_attention(model, onnx_path, batch_size, hidden_size, num_attention_heads, sequence_length, + past_sequence_length, float16, device, padding_length, optimized, test_cases) # Test Case: with past state and padding last 2 words sequence_length = 3 past_sequence_length = 5 padding_length = 2 - num_failure += verify_attention(model, - onnx_path, - batch_size, - hidden_size, - num_attention_heads, - sequence_length, - past_sequence_length, - float16, - device, - padding_length, - optimized, - test_cases) + num_failure += verify_attention(model, onnx_path, batch_size, hidden_size, num_attention_heads, sequence_length, + past_sequence_length, float16, device, padding_length, optimized, test_cases) # Test Case: random mask one word sequence_length = 1 past_sequence_length = 128 padding_length = -1 - num_failure += verify_attention(model, - onnx_path, - batch_size, - hidden_size, - num_attention_heads, - sequence_length, - past_sequence_length, - float16, - device, - padding_length, - optimized, - test_cases) + num_failure += verify_attention(model, onnx_path, batch_size, hidden_size, num_attention_heads, sequence_length, + past_sequence_length, float16, device, padding_length, optimized, test_cases) # clean up onnx file os.remove(onnx_model_path) @@ -467,24 +393,45 @@ def run(batch_size, float16, optimized, hidden_size, num_attention_heads, device class TestGptAttentionHuggingfaceParity(unittest.TestCase): def setUp(self): - self.optimized = True # Change it to False if you want to test parity of non optimized ONNX + self.optimized = True # Change it to False if you want to test parity of non optimized ONNX self.test_cases = 10 # Number of test cases per test run def run_test(self, batch_size, float16, optimized, hidden_size, num_attention_heads, device): - if float16 and device.type=='cpu': # CPU does not support FP16 + if float16 and device.type == 'cpu': # CPU does not support FP16 return - num_failure, test_name = run(batch_size, float16, optimized, hidden_size, num_attention_heads, device, self.test_cases) + num_failure, test_name = run(batch_size, float16, optimized, hidden_size, num_attention_heads, device, + self.test_cases) self.assertTrue(num_failure == 0, test_name) def run_small(self, optimized, device): for batch_size in [64]: - self.run_test(batch_size, float16=False, optimized=optimized, hidden_size=768, num_attention_heads=12, device=device) - self.run_test(batch_size, float16=True, optimized=optimized, hidden_size=768, num_attention_heads=12, device=device) + self.run_test(batch_size, + float16=False, + optimized=optimized, + hidden_size=768, + num_attention_heads=12, + device=device) + self.run_test(batch_size, + float16=True, + optimized=optimized, + hidden_size=768, + num_attention_heads=12, + device=device) def run_large(self, optimized, device): for batch_size in [2]: - self.run_test(batch_size, float16=False, optimized=optimized, hidden_size=4096, num_attention_heads=32, device=device) - self.run_test(batch_size, float16=True, optimized=optimized, hidden_size=4096, num_attention_heads=32, device=device) + self.run_test(batch_size, + float16=False, + optimized=optimized, + hidden_size=4096, + num_attention_heads=32, + device=device) + self.run_test(batch_size, + float16=True, + optimized=optimized, + hidden_size=4096, + num_attention_heads=32, + device=device) def test_cpu(self): cpu = torch.device('cpu') @@ -507,5 +454,6 @@ class TestGptAttentionHuggingfaceParity(unittest.TestCase): gpu = torch.device('cuda') self.run_large(self.optimized, gpu) + if __name__ == '__main__': unittest.main() diff --git a/onnxruntime/test/python/transformers/test_parity_layernorm.py b/onnxruntime/test/python/transformers/test_parity_layernorm.py new file mode 100644 index 0000000000..b83da80f07 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_parity_layernorm.py @@ -0,0 +1,227 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import unittest +import torch +from torch import nn +import os +from parity_utilities import * + + +class LayerNorm(nn.Module): + def __init__(self, hidden_size, epsilon, cast_fp16=True): + super().__init__() + self.layer_norm = nn.LayerNorm(hidden_size, eps=epsilon) + # initialize weights with random value + self.layer_norm.bias.data.normal_(mean=0.0, std=0.1) + self.layer_norm.weight.data.normal_(mean=0.0, std=0.5) + self.cast_fp16 = cast_fp16 + + @staticmethod + def get_fused_op(): + return "LayerNormalization" + + def forward(self, x): + if self.cast_fp16 and x.dtype == torch.float16: + y = self.layer_norm(x.to(torch.float32)).to(torch.float16) + return (y, ) + y = self.layer_norm(x) + return (y, ) + + +def optimize_fp16_onnx_with_cast(input_onnx_path, optimized_onnx_path, epsilon): + from onnxruntime.transformers.onnx_model import OnnxModel + import onnx + m = onnx.load(input_onnx_path) + onnx_model = OnnxModel(m) + + nodes_to_remove = onnx_model.nodes() + nodes_to_add = [ + onnx.helper.make_node("Cast", ["input"], ["fp32_input"], "cast_input", to=1), + onnx.helper.make_node("Cast", ["layer_norm.weight"], ["fp32_layer_norm.weight"], "cast_weight", to=1), + onnx.helper.make_node("Cast", ["layer_norm.bias"], ["fp32_layer_norm.bias"], "cast_bias", to=1), + onnx.helper.make_node("LayerNormalization", ["fp32_input", "fp32_layer_norm.weight", "fp32_layer_norm.bias"], + ["fp32_output"], + "layer_norm", + epsilon=epsilon), # use fp32 epsilon + onnx.helper.make_node("Cast", ["fp32_output"], ["output"], "cast_output", to=10) + ] + + onnx_model.remove_nodes(nodes_to_remove) + onnx_model.add_nodes(nodes_to_add) + onnx_model.prune_graph() + onnx_model.save_model_to_file(optimized_onnx_path) + + +def optimize_fp16_onnx_no_cast(input_onnx_path, optimized_onnx_path, epsilon): + from onnxruntime.transformers.onnx_model import OnnxModel + import onnx + m = onnx.load(input_onnx_path) + onnx_model = OnnxModel(m) + + nodes_to_remove = onnx_model.nodes() + node_to_add = onnx.helper.make_node("LayerNormalization", ["input", "layer_norm.weight", "layer_norm.bias"], + ["output"], + "layer_norm", + epsilon=epsilon) + + onnx_model.remove_nodes(nodes_to_remove) + onnx_model.add_node(node_to_add) + onnx_model.prune_graph() + onnx_model.save_model_to_file(optimized_onnx_path) + + +def get_output_names(): + outputs = ["output"] + return outputs + + +def run(batch_size, + float16, + optimized, + hidden_size, + device, + test_cases, + sequence_length=2, + epsilon=0.00001, + cast_fp16=True, + cast_onnx_only=False, + verbose=False): + test_name = f"device={device}, float16={float16}, optimized={optimized}, batch_size={batch_size}, sequence_length={sequence_length}, hidden_size={hidden_size}, epsilon={epsilon}, cast_fp16={cast_fp16}, cast_onnx_only={cast_onnx_only}" + print(f"\nTesting: {test_name}") + + model = LayerNorm(hidden_size, epsilon, cast_fp16) + model.eval() + model.to(device) + + if float16 and not cast_fp16: + model.half() + + # Do not re-use onnx file from previous test since weights of model are random. + onnx_model_path = './temp/layer_norm_{}.onnx'.format("fp16" if float16 else "fp32") + export_onnx(model, onnx_model_path, float16, hidden_size, device) + + if optimized: + optimized_onnx_path = './temp/layer_norm_{}_opt.onnx'.format("fp16" if float16 else "fp32") + if (not float16) or cast_fp16: + optimize_onnx(onnx_model_path, optimized_onnx_path, + expected_op=LayerNorm.get_fused_op()) + else: + if cast_onnx_only: + optimize_fp16_onnx_with_cast(onnx_model_path, optimized_onnx_path, epsilon=epsilon) + else: + optimize_fp16_onnx_no_cast(onnx_model_path, optimized_onnx_path, epsilon=epsilon) + + 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, + verbose=verbose) + + # clean up onnx file + os.remove(onnx_model_path) + if optimized: + os.remove(onnx_path) + + return num_failure, test_name + + +class TestLayerNormParity(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.epsilon = 0.00001 + self.verbose = False + + def run_test(self, + batch_size, + float16, + optimized, + hidden_size, + device, + cast_fp16=True, + cast_onnx_only=False, + enable_assert=True): + 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, + self.sequence_length, + self.epsilon, + cast_fp16, + cast_onnx_only, + verbose=self.verbose) + if enable_assert: + self.assertTrue(num_failure == 0, "Failed: " + test_name) + + def run_one(self, optimized, device, hidden_size=768): + for batch_size in [4]: + self.run_test(batch_size, float16=False, optimized=optimized, hidden_size=hidden_size, device=device) + + self.run_test( + batch_size, + float16=True, + optimized=optimized, + hidden_size=hidden_size, + device=device, + cast_fp16=True, + cast_onnx_only=False, + enable_assert=False # This setting has small chance to exceed tollerance threshold 0.001 + ) + + self.run_test( + batch_size, + float16=True, + optimized=optimized, + hidden_size=hidden_size, + device=device, + cast_fp16=False, + cast_onnx_only=False, + enable_assert=False # This setting cannot pass tollerance threshold + ) + + self.run_test( + batch_size, + float16=True, + optimized=optimized, + hidden_size=hidden_size, + device=device, + cast_fp16=False, + cast_onnx_only=True, + enable_assert=False # This setting cannot pass tollerance threshold + ) + + def test_cpu(self): + cpu = torch.device('cpu') + self.run_one(self.optimized, cpu, hidden_size=self.hidden_size) + + 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') + self.run_one(self.optimized, gpu, hidden_size=self.hidden_size) + + +if __name__ == '__main__': + unittest.main()