diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc new file mode 100644 index 0000000000..59904f348f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/common.h" +#include "core/providers/cuda/cudnn_common.h" +#include "core/framework/tensorprotoutils.h" +#include "fast_gelu.h" +#include "fast_gelu_impl.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + FastGelu, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + FastGelu); + +REGISTER_KERNEL_TYPED(float) +REGISTER_KERNEL_TYPED(MLFloat16) + +using namespace ONNX_NAMESPACE; + +template +FastGelu::FastGelu(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { +} + +template +Status FastGelu::ComputeInternal(OpKernelContext* ctx) const { + const Tensor* input = ctx->Input(0); + + const auto input_dims = input->Shape().GetDims(); + if (input_dims.size() < 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 0 is expected to have 1 or more dimensions, got ", input_dims.size()); + } + + size_t num_inputs = OpKernel::Node().InputDefs().size(); + bool has_bias = (num_inputs == 2); + + int input_length = 1; + for (size_t i = 0; i < input_dims.size(); i++) { + input_length *= static_cast(input_dims[i]); + } + + int bias_length = 0; + const Tensor* bias = nullptr; + if (has_bias) { + bias = ctx->Input(1); + const auto bias_dims = bias->Shape().GetDims(); + if (bias_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 1 is expected to have 1 dimensions, got ", bias_dims.size()); + } + if (bias_dims[0] != input_dims[input_dims.size() - 1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 1 dimension 0 should have same length as the last dimension of input 0"); + } + bias_length = static_cast(bias_dims[0]); + } + + Tensor* output = ctx->Output(0, input->Shape()); + + typedef typename ToCudaType::MappedType CudaT; + if (!LaunchFastGeluKernel(nullptr, + input_length, + bias_length, + reinterpret_cast(input->template Data()), + has_bias ? reinterpret_cast(bias->template Data()) : nullptr, + reinterpret_cast(output->template MutableData()))) { + CUDA_CALL(cudaGetLastError()); + return Status(common::ONNXRUNTIME, common::FAIL); + } + + return Status::OK(); +} + +} //namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.h b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.h new file mode 100644 index 0000000000..1731bc5223 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.h @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/common/common.h" +#include "core/framework/op_kernel.h" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +template +class FastGelu final : public CudaKernel { + public: + FastGelu(const OpKernelInfo& op_kernel_info); + Status ComputeInternal(OpKernelContext* ctx) const override; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu new file mode 100644 index 0000000000..9a9b35556d --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu @@ -0,0 +1,102 @@ +/* + The implementation of this file is based on gelu plugin in TensorRT demo: + https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/ + +Copyright 2019 NVIDIA Corporation + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Modifications: Add (bias) before Gelu is merged into this op to get better performance. +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/shared_inc/cuda_call.h" +#include "fast_gelu_impl.h" + +using namespace onnxruntime::cuda; + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// constants for approximating the normal cdf +constexpr float A = 0.5; + +constexpr float B = 0.7978845608028654; // sqrt(2.0/M_PI) + +constexpr float C = 0.035677408136300125; // 0.044715 * sqrt(2.0/M_PI) + +template +__global__ void FastGeluKernel(const T a, const T b, const T c, int input_length, int bias_length, const T* input, const T* bias, T* output) { + const int idx = blockIdx.x * TPB + threadIdx.x; + + if (idx < input_length) { + const T x = input[idx]; + const T in = (bias == nullptr) ? x : (x + bias[idx % bias_length]); + const T cdf = a + a * _Tanh(in * (c * in * in + b)); + output[idx] = in * cdf; + } +} + +template +__global__ void FastGeluKernel2(const half2 a, const half2 b, const half2 c, int input_length, int bias_length, const half2* input, const half2* bias, half2* output) { +// half2 arithmetic functions requires cuda architecture >= 5.3 +#if __CUDA_ARCH__ >= 530 + const int idx = blockIdx.x * TPB + threadIdx.x; + + if (idx < input_length) { + const half2 x = input[idx]; + const half2 in = (bias == nullptr) ? x : (x + bias[idx % bias_length]); + const half2 cdf = a + a * _Tanh(in * (c * in * in + b)); + output[idx] = in * cdf; + } +#endif +} + +template <> +bool LaunchFastGeluKernel(cudaStream_t stream, int input_length, int bias_length, const float* input, const float* bias, float* output) { + constexpr int blockSize = 256; + const int gridSize = (input_length + blockSize - 1) / blockSize; + FastGeluKernel<<>>(A, B, C, input_length, bias_length, input, bias, output); + + return CUDA_CALL(cudaPeekAtLastError()); +} + +template <> +bool LaunchFastGeluKernel(cudaStream_t stream, int input_length, int bias_length, const half* input, const half* bias, half* output) { + constexpr int blockSize = 256; + + if (0 == (bias_length & 1) && DeviceProp::GetDeviceProps().major >= 7) { + const int n = input_length / 2; + const int gridSize = (n + blockSize - 1) / blockSize; + const half2 A2 = __floats2half2_rn(A, A); + const half2 B2 = __floats2half2_rn(B, B); + const half2 C2 = __floats2half2_rn(C, C); + const half2* input2 = reinterpret_cast(input); + const half2* bias2 = reinterpret_cast(bias); + half2* output2 = reinterpret_cast(output); + FastGeluKernel2<<>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2); + } else { + const int gridSize = (input_length + blockSize - 1) / blockSize; + FastGeluKernel<<>>(A, B, C, input_length, bias_length, input, bias, output); + } + + return CUDA_CALL(cudaPeekAtLastError()); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.h b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.h new file mode 100644 index 0000000000..9e59219fee --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.h @@ -0,0 +1,15 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +bool LaunchFastGeluKernel(cudaStream_t stream, int input_length, int bias_length, const T* input, const T* bias, T* output); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc index b4af013dab..799a86776f 100644 --- a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc @@ -9,6 +9,8 @@ using namespace onnxruntime::common; namespace onnxruntime { namespace contrib { namespace cuda { +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FastGelu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, FastGelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Gelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Gelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Gelu); @@ -48,6 +50,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, void RegisterCudaContribKernels(KernelRegistry& kernel_registry) { static const BuildKernelCreateInfoFn function_table[] = { + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 5ace5c396c..82ef303efa 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -314,6 +314,17 @@ void RegisterBertSchemas() { updateOutputShape(ctx, 1, mask_index_shape); }); + ONNX_CONTRIB_OPERATOR_SCHEMA(FastGelu) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) + .SetDoc("Gelu") + .Input(0, "X", "input tensor", "T") + .Input(1, "bias", "bias tensor", "T", OpSchema::Optional) + .Output(0, "Y", "output tensor", "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.") + .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); + ONNX_CONTRIB_OPERATOR_SCHEMA(SkipLayerNormalization) .SetDomain(kMSDomain) .SinceVersion(1) diff --git a/onnxruntime/core/providers/cuda/cu_inc/common.cuh b/onnxruntime/core/providers/cuda/cu_inc/common.cuh index 9bf34e7bc0..b440519a8a 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/common.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/common.cuh @@ -7,6 +7,7 @@ #include #include #include +#include #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/cuda_call.h" @@ -154,6 +155,14 @@ __device__ __inline__ double _Tanh(double a) { return tanh(a); } template <> __device__ __inline__ half _Tanh(half a) { return half(tanhf((float)a)); } +template <> +__device__ __inline__ half2 _Tanh(half2 a) { + float2 tmp = (__half22float2(a)); + tmp.x = tanhf(tmp.x); + tmp.y = tanhf(tmp.y); + return __float22half2_rn(tmp); +} + template __device__ __inline__ T _Pow(T a, T b); diff --git a/onnxruntime/python/tools/bert/README.md b/onnxruntime/python/tools/bert/README.md index 3ed2a998cc..fa72c62bcd 100644 --- a/onnxruntime/python/tools/bert/README.md +++ b/onnxruntime/python/tools/bert/README.md @@ -1,11 +1,11 @@ # BERT Model Optimization Tool Overview -This tool converts a BERT ONNX model exported from PyTorch, and generates a optimized model to run faster in NVidia GPU. +This tool converts a BERT ONNX model exported from PyTorch, and generates an optimized model to run faster in NVidia GPU. Currently, this script **cannot** process BERT models exported from Tensorflow since the graph has some difference. ## Export an BERT model from PyTorch -For example, after using https://github.com/huggingface/transformers to Train a BERT model in PyTorch 1.3, you can use the following function to export ONNX model. +For example, after using https://github.com/huggingface/transformers to train a BERT model in PyTorch 1.3, you can use the following function to export ONNX model. Please specify do_constant_folding=True. That's required for this tool. diff --git a/onnxruntime/python/tools/bert/bert_model_optimization.py b/onnxruntime/python/tools/bert/bert_model_optimization.py index c1e6dc55b6..6fbc6d9963 100644 --- a/onnxruntime/python/tools/bert/bert_model_optimization.py +++ b/onnxruntime/python/tools/bert/bert_model_optimization.py @@ -350,7 +350,7 @@ class BertOnnxModel(OnnxModel): # constant node names self.normalize_name = "SkipLayerNormalization" - self.gelu_name = 'Gelu' + self.gelu_name = 'FastGelu' self.attention_name = 'Attention' def get_normalize_nodes(self): @@ -510,7 +510,7 @@ class BertOnnxModel(OnnxModel): if len(matmul_child) != 1 or matmul_child[0].op_type != 'Add': continue add_node = matmul_child[0] - + children = input_name_to_nodes[add_node.output[0]] children_types = sorted([child.op_type for child in children]) @@ -525,9 +525,11 @@ class BertOnnxModel(OnnxModel): if len(subgraph_nodes) != 5: continue + nodes_to_remove.append(add_node) nodes_to_remove.extend(subgraph_nodes) + bias_input = add_node.input[1] if (add_node.input[0] == matmul_node.output[0]) else add_node.input[0] gelu_node = onnx.helper.make_node(self.gelu_name, - inputs=[add_node.output[0]], + inputs=[matmul_node.output[0], bias_input], outputs=[matmul_2.input[0]]) gelu_node.domain = "com.microsoft" nodes_to_add.append(gelu_node) @@ -882,7 +884,6 @@ def main(): else: bert_model.cast_input_to_int32() - if args.float16: bert_model.convert_model_float32_to_float16() diff --git a/onnxruntime/test/contrib_ops/fastgelu_op_test.cc b/onnxruntime/test/contrib_ops/fastgelu_op_test.cc new file mode 100644 index 0000000000..5afa3ec22f --- /dev/null +++ b/onnxruntime/test/contrib_ops/fastgelu_op_test.cc @@ -0,0 +1,172 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/common/cuda_op_test_utils.h" +#include "test/providers/provider_test_utils.h" + +using namespace onnxruntime::test; + +namespace onnxruntime { +namespace test { + +const std::vector ComputeGelu(const std::vector& input_data) { + std::vector output; + output.reserve(input_data.size()); + + for (size_t i = 0; i < input_data.size(); i++) { + float x = input_data[i]; + float y = x * (0.5f + 0.5f * tanh(x * (0.035677408136300125f * x * x + 0.7978845608028654f))); + output.push_back(y); + } + return output; +} + +const std::vector AddBias(const std::vector& input_data, const std::vector& bias_data) { + size_t bias_length = bias_data.size(); + + std::vector output; + output.reserve(input_data.size()); + + for (size_t i = 0; i < input_data.size(); i++) { + output.push_back(input_data[i] + bias_data[i % bias_length]); + } + return output; +} + +const std::vector GetExpectedResult(const std::vector& input_data, const std::vector& bias_data) { + std::vector add_bias_data = AddBias(input_data, bias_data); + return ComputeGelu(add_bias_data); +} + +static void RunFastGeluTest( + const std::vector& input_data, + const std::vector& bias_data, + const std::vector& output_data, + const std::vector& input_dims, + const std::vector& bias_dims, + const std::vector& output_dims, + bool has_bias = true, + bool use_float16 = false) { + int min_cuda_architecture = use_float16 ? 530 : 0; + if (HasCudaEnvironment(min_cuda_architecture)) { + OpTester tester("FastGelu", 1, onnxruntime::kMSDomain); + + if (use_float16) { + tester.AddInput("X", input_dims, ToFloat16(input_data)); + if (has_bias) { + tester.AddInput("bias", bias_dims, ToFloat16(bias_data)); + } + tester.AddOutput("Y", output_dims, ToFloat16(output_data)); + } else { + tester.AddInput("X", input_dims, input_data); + if (has_bias) { + tester.AddInput("bias", bias_dims, bias_data); + } + tester.AddOutput("Y", output_dims, output_data); + } + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + +// This test simulates Gelu in Bert model for float32 +static void RunFastGeluTest( + const std::vector& input_data, + const std::vector& bias_data, + int batch_size, + int sequence_length, + int hidden_size) { + std::vector output_data; + + bool has_bias = (bias_data.size() > 0); + if (has_bias) { + output_data = GetExpectedResult(input_data, bias_data); + } else { + output_data = ComputeGelu(input_data); + } + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector bias_dims = {hidden_size}; + std::vector output_dims = input_dims; + RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, has_bias); +} + +TEST(FastGeluTest, FastGeluWithBiasFloat32) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f}; + + RunFastGeluTest(input_data, bias_data, batch_size, sequence_length, hidden_size); +} + +TEST(FastGeluTest, FastGeluWithoutBiasFloat32) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector bias_data = {}; + + RunFastGeluTest(input_data, bias_data, batch_size, sequence_length, hidden_size); +} + +TEST(FastGeluTest, FastGeluWithBiasFloat16) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f}; + + std::vector output_data = { + 0.1851806640625f, 0.054046630859375f, 1.0615234375f, 3.095703125f, + 0, 0.63037109375f, 1.3984375f, 1.3984375f}; + + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector bias_dims = {hidden_size}; + std::vector output_dims = input_dims; + + RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, true, true); +} + +TEST(FastGeluTest, FastGeluWithoutBiasFloat16) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector bias_data = {}; + + std::vector output_data = { + 0.63037109375f, -0.154296875f, 0.0f, 0.8408203125f, + 0.345703125f, 0.11578369140625f, 0.1854248046875f, -0.1646728515625f }; + + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector bias_dims = {}; + std::vector output_dims = input_dims; + + RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, false, true); +} + +} // namespace test +} // namespace onnxruntime