From 4c88fa59713a4fa8f65506b840b35dc3a859a83c Mon Sep 17 00:00:00 2001 From: zhangyaobit <1034716+zhangyaobit@users.noreply.github.com> Date: Fri, 4 Mar 2022 08:51:15 -0800 Subject: [PATCH] Add micro-benchmark for FastGelu (#10744) * Add micro-benchmark for FastGelu * Delete the bert-base case, as it is very similar to the bert-large one. * Add argument parsing and more user-friendly provider type assertion. --- .../python/tools/microbench/benchmark.py | 29 ++++++--- .../python/tools/microbench/fast_gelu.py | 60 +++++++++++++++++++ onnxruntime/python/tools/microbench/matmul.py | 17 ++++-- .../microbench/models/fast_gelu_fp16.onnx | 23 +++++++ .../microbench/models/fast_gelu_fp32.onnx | 20 +++++++ 5 files changed, 136 insertions(+), 13 deletions(-) create mode 100644 onnxruntime/python/tools/microbench/fast_gelu.py create mode 100644 onnxruntime/python/tools/microbench/models/fast_gelu_fp16.onnx create mode 100644 onnxruntime/python/tools/microbench/models/fast_gelu_fp32.onnx diff --git a/onnxruntime/python/tools/microbench/benchmark.py b/onnxruntime/python/tools/microbench/benchmark.py index 89359782a3..e7fabd0335 100644 --- a/onnxruntime/python/tools/microbench/benchmark.py +++ b/onnxruntime/python/tools/microbench/benchmark.py @@ -1,9 +1,16 @@ +from argparse import ArgumentParser import time import numpy import onnxruntime as ort import torch +def add_arguments(parser: ArgumentParser): + parser.add_argument("--provider", required=False, type=str, default="rocm", help="Execution provider to use") + parser.add_argument("--precision", required=False, type=str, default="fp16", help="Number format to use") + parser.add_argument('--profiling', type=bool, default=False, help='If enable profiling') + + def create_input_output_tensors(inputs, outputs): device = "cuda" input_tensors = {name: torch.from_numpy(array).to(device) for name, array in inputs.items()} @@ -29,20 +36,28 @@ def create_io_binding(sess, input_tensors, output_tensors): return io_binding -def create_session(onnx_file, provider, profiling): +def create_session(onnx_file, args): sess_opt = ort.SessionOptions() - sess_opt.enable_profiling = profiling - if provider == "rocm": + sess_opt.enable_profiling = args.profiling + if args.provider == "rocm": execution_provider = ["ROCMExecutionProvider"] - else: + elif args.provider == "cuda": execution_provider = ["CUDAExecutionProvider"] - + else: + raise ValueError(f"The script doesn't support provider type '{args.provider}' yet.") + sess = ort.InferenceSession(onnx_file, sess_options=sess_opt, providers=execution_provider) + + if args.provider == "rocm": + assert 'ROCMExecutionProvider' in sess.get_providers() + elif args.provider == "cuda": + assert 'CUDAExecutionProvider' in sess.get_providers() + return sess -def benchmark(onnx_file, inputs, outputs, provider, profiling=False): - sess = create_session(onnx_file, provider, profiling) +def benchmark(onnx_file, inputs, outputs, args): + sess = create_session(onnx_file, args) input_tensors, output_tensors = create_input_output_tensors(inputs, outputs) io_binding = create_io_binding(sess, input_tensors, output_tensors) diff --git a/onnxruntime/python/tools/microbench/fast_gelu.py b/onnxruntime/python/tools/microbench/fast_gelu.py new file mode 100644 index 0000000000..d5eeec0b98 --- /dev/null +++ b/onnxruntime/python/tools/microbench/fast_gelu.py @@ -0,0 +1,60 @@ +import argparse +import numpy as np +from benchmark import benchmark, add_arguments + + +def create_inputs_outputs(batch, seq_len, intermediate_dimension, data_type): + np.random.seed(0) + a = np.random.rand(batch, seq_len, intermediate_dimension).astype(data_type) + b = np.random.rand(intermediate_dimension).astype(data_type) + c = np.random.rand(batch, seq_len, intermediate_dimension).astype(data_type) + + inputs = {"A": a, "B": b} + outputs = {"return_val": c} + + return inputs, outputs + + +def add_benchmark_case(benchmark_cases, batch_size, seq_len, intermediate_dimension, data_type, model): + benchmark_cases += [ + (batch_size, seq_len, intermediate_dimension, data_type, model), + ] + + +def create_benchmark_cases(precision): + benchmark_cases = [] + if precision == "fp16": + model = "models/fast_gelu_fp16.onnx" + data_type = np.float16 + else: + model = "models/fast_gelu_fp32.onnx" + data_type = np.float32 + + # bert-large + hidden_size = 1024 + seq_len = 384 + batch_size = 1 + intermediate_dimension = hidden_size * 4 + add_benchmark_case(benchmark_cases, batch_size, seq_len, intermediate_dimension, data_type, model) + + return benchmark_cases + + +def benchmark_fast_gelu(batch, seq_len, intermediate_dimension, data_type, onnx_file, args): + inputs, outputs = create_inputs_outputs(batch, seq_len, intermediate_dimension, data_type) + time = benchmark(onnx_file, inputs, outputs, args) + return time + + +def main(): + parser = argparse.ArgumentParser() + add_arguments(parser) + args = parser.parse_args() + + for (batch, seq_len, intermediate_dimension, data_type, onnx_file) in create_benchmark_cases(args.precision): + time = benchmark_fast_gelu(batch, seq_len, intermediate_dimension, data_type, onnx_file, args) + print(f"(batch seq_len inter_dim) = ({batch} {seq_len} {intermediate_dimension}), {time:7.4f} ms") + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/microbench/matmul.py b/onnxruntime/python/tools/microbench/matmul.py index b7f11f0896..006a097301 100644 --- a/onnxruntime/python/tools/microbench/matmul.py +++ b/onnxruntime/python/tools/microbench/matmul.py @@ -1,5 +1,6 @@ +import argparse import numpy as np -from benchmark import benchmark +from benchmark import benchmark, add_arguments def create_inputs_outputs(b1, b2, m, k, n, data_type): @@ -23,7 +24,7 @@ def add_benchmark_case(benchmark_cases, batch_size, seq_len, hidden_size, interm ] -def create_benchmark_cases(precision="fp16"): +def create_benchmark_cases(precision): benchmark_cases = [] if precision == "fp16": model = "models/matmul_fp16.onnx" @@ -51,15 +52,19 @@ def create_benchmark_cases(precision="fp16"): return benchmark_cases -def benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file): +def benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file, args): inputs, outputs = create_inputs_outputs(b1, b2, m, k, n, data_type) - time = benchmark(onnx_file, inputs, outputs, "rocm") + time = benchmark(onnx_file, inputs, outputs, args) return time def main(): - for (b1, b2, m, k, n, data_type, onnx_file) in create_benchmark_cases(): - time = benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file) + parser = argparse.ArgumentParser() + add_arguments(parser) + args = parser.parse_args() + + for (b1, b2, m, k, n, data_type, onnx_file) in create_benchmark_cases(args.precision): + time = benchmark_matmul(b1, b2, m, k, n, data_type, onnx_file, args) tflops = b1 * b2 * m * k * n * 2 / time / 1000000000 print(f"(b1 b2 m k n) = ({b1} {b2} {m} {k} {n}), {time:7.4f} ms, {tflops:4.2f} tflops") diff --git a/onnxruntime/python/tools/microbench/models/fast_gelu_fp16.onnx b/onnxruntime/python/tools/microbench/models/fast_gelu_fp16.onnx new file mode 100644 index 0000000000..5487313893 --- /dev/null +++ b/onnxruntime/python/tools/microbench/models/fast_gelu_fp16.onnx @@ -0,0 +1,23 @@ +p2o:· ++ +A +B +return_val"FastGelu: com.microsoft fast_geluZ, +A' +% +! +batch + seq_len +  inter_dimZ +B + + +  inter_dimb5 + +return_val' +% +! +batch + seq_len +  inter_dimB + com.microsoft \ No newline at end of file diff --git a/onnxruntime/python/tools/microbench/models/fast_gelu_fp32.onnx b/onnxruntime/python/tools/microbench/models/fast_gelu_fp32.onnx new file mode 100644 index 0000000000..13bf295234 --- /dev/null +++ b/onnxruntime/python/tools/microbench/models/fast_gelu_fp32.onnx @@ -0,0 +1,20 @@ +p2o:· ++ +A +B +return_val"FastGelu: com.microsoft fast_geluZ, +A' +%! +batch + seq_len +  inter_dimZ +B + +  inter_dimb5 + +return_val' +%! +batch + seq_len +  inter_dimB + com.microsoft \ No newline at end of file