From 4df356d1c9f1387897ccb9df1d170ba80e929756 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Thu, 14 Jan 2021 19:04:32 +0800 Subject: [PATCH] Train BERT Using BFloat16 on A100 (#6090) * traing bert using bf16 * Adam support bf16 * bugfix * add fusedmatmul support * fix after merge from master. * bugfix * bugfix after merge from master * fast reduction for bf16. * resolve comments * fix win build * bugfix * change header file. Co-authored-by: Vincent Wang --- .../contrib_ops/cuda/bert/fast_gelu.cc | 3 + .../contrib_ops/cuda/bert/fast_gelu_impl.cu | 40 ++- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 17 +- onnxruntime/contrib_ops/cuda/layer_norm.cc | 3 + .../contrib_ops/cuda/layer_norm_impl.cu | 4 + .../contrib_ops/cuda/math/fused_matmul.cc | 5 + .../core/optimizer/dropout_elimination.cc | 2 +- .../providers/cuda/activation/activations.cc | 7 + .../cuda/activation/activations_impl.cu | 13 +- .../core/providers/cuda/atomic/common.cuh | 21 ++ .../core/providers/cuda/cu_inc/common.cuh | 28 ++ onnxruntime/core/providers/cuda/cuda_common.h | 12 + .../providers/cuda/cuda_execution_provider.cc | 31 ++ .../providers/cuda/cuda_execution_provider.h | 10 + onnxruntime/core/providers/cuda/cuda_utils.cu | 6 + .../core/providers/cuda/cudnn_common.cc | 5 + .../core/providers/cuda/cudnn_common.h | 8 + onnxruntime/core/providers/cuda/fpgeneric.cu | 19 +- .../cuda/math/binary_elementwise_ops.cc | 25 +- .../cuda/math/binary_elementwise_ops_impl.cu | 11 + onnxruntime/core/providers/cuda/math/gemm.cc | 3 + .../core/providers/cuda/math/matmul.cc | 3 + .../core/providers/cuda/math/softmax.cc | 27 ++ .../core/providers/cuda/math/softmax_impl.cu | 3 + .../cuda/math/unary_elementwise_ops_impl.cu | 23 ++ .../cuda/math/variadic_elementwise_ops.cc | 16 +- .../math/variadic_elementwise_ops_impl.cu | 7 + onnxruntime/core/providers/cuda/nn/dropout.cc | 13 +- onnxruntime/core/providers/cuda/nn/dropout.h | 4 + .../core/providers/cuda/nn/dropout_impl.cu | 3 + .../cuda/reduction/reduction_functions.cu | 9 + .../providers/cuda/reduction/reduction_ops.cc | 334 ++++++++---------- .../cuda/shared_inc/accumulation_type.h | 4 + .../providers/cuda/shared_inc/fpgeneric.h | 100 ++++++ .../core/providers/cuda/tensor/cast_op.cc | 9 + .../core/providers/cuda/tensor/gather_nd.cc | 25 +- .../providers/cuda/tensor/gather_nd_impl.cu | 3 + .../core/graph/training_op_defs.cc | 34 +- orttraining/orttraining/models/bert/main.cc | 2 + orttraining/orttraining/models/gpt2/main.cc | 2 + .../cuda/activation/bias_gelu_grad.cc | 12 +- .../cuda/activation/bias_gelu_grad_impl.cu | 5 + .../cuda/cuda_training_kernels.cc | 69 ++++ .../training_ops/cuda/math/isfinite.cuh | 7 + .../cuda/math/mixed_precision_scale.cc | 24 +- .../cuda/math/mixed_precision_scale.cu | 8 + .../training_ops/cuda/math/softmax_grad.cc | 30 ++ .../cuda/math/softmax_grad_impl.cu | 4 + .../training_ops/cuda/nn/dropout.cc | 27 +- .../training_ops/cuda/nn/dropout_impl.cu | 7 +- .../training_ops/cuda/nn/layer_norm.cc | 3 + .../training_ops/cuda/nn/layer_norm_impl.cu | 6 +- .../training_ops/cuda/optimizer/adam.cc | 12 + .../training_ops/cuda/optimizer/adam.cu | 12 + .../cuda/optimizer/gradient_control.cc | 5 + .../cuda/optimizer/gradient_control.cu | 5 + .../training_ops/cuda/optimizer/lamb.cc | 11 + .../training_ops/cuda/optimizer/lamb.cu | 34 ++ .../cuda/reduction/reduction_all.cc | 5 + .../cuda/reduction/reduction_all.cu | 9 + .../training_ops/cuda/tensor/gather_grad.cc | 17 +- .../cuda/tensor/gather_grad_impl.cu | 3 + .../cuda/tensor/gather_nd_grad.cc | 16 +- .../cuda/tensor/gather_nd_grad_impl.cu | 3 + 64 files changed, 965 insertions(+), 263 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc index f7c06636a3..642ef3458c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc @@ -25,6 +25,9 @@ namespace cuda { REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_KERNEL_TYPED(BFloat16) +#endif using namespace ONNX_NAMESPACE; diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu index 484a8bbe24..fbf04128d2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu @@ -45,7 +45,7 @@ __global__ void FastGeluKernel(const T a, const T b, const T c, int input_length if (idx < input_length) { const T x = input[idx]; - const T in = (bias == nullptr) ? x : (x + bias[idx % bias_length]); + const T in = (bias == nullptr) ? x : (T)(x + bias[idx % bias_length]); const T cdf = a + a * _Tanh(in * (c * in * in + b)); output[idx] = in * cdf; } @@ -97,6 +97,44 @@ bool LaunchFastGeluKernel(const cudaDeviceProp& prop, cudaStream_t stream, int i return CUDA_CALL(cudaPeekAtLastError()); } +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template +__global__ void FastGeluKernel2(const nv_bfloat162 a, const nv_bfloat162 b, const nv_bfloat162 c, + int input_length, int bias_length, + const nv_bfloat162* input, const nv_bfloat162* bias, nv_bfloat162* output) { + const int idx = blockIdx.x * TPB + threadIdx.x; + + if (idx < input_length) { + const nv_bfloat162 x = input[idx]; + const nv_bfloat162 in = (bias == nullptr) ? x : (x + bias[idx % bias_length]); + const nv_bfloat162 cdf = a + a * _Tanh(in * (c * in * in + b)); + output[idx] = in * cdf; + } +} + +template <> +bool LaunchFastGeluKernel(const cudaDeviceProp& prop, cudaStream_t stream, int input_length, int bias_length, const nv_bfloat16* input, const nv_bfloat16* bias, nv_bfloat16* output) { + constexpr int blockSize = 256; + + if (0 == (bias_length & 1) && prop.major >= 7) { + const int n = input_length / 2; + const int gridSize = (n + blockSize - 1) / blockSize; + const nv_bfloat162 A2 = __floats2bfloat162_rn(A, A); + const nv_bfloat162 B2 = __floats2bfloat162_rn(B, B); + const nv_bfloat162 C2 = __floats2bfloat162_rn(C, C); + const nv_bfloat162* input2 = reinterpret_cast(input); + const nv_bfloat162* bias2 = reinterpret_cast(bias); + nv_bfloat162* 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()); +} +#endif + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index bc94cc266a..48bd6e8a33 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -80,6 +80,13 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int8_t, QAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int8_t, QAttention); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FastGelu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, TransposeMatMul); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FusedMatMul); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, BFloat16_float, LayerNormalization); +#endif + template <> KernelCreateInfo BuildKernelCreateInfo() { KernelCreateInfo info; @@ -159,7 +166,15 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo}; + BuildKernelCreateInfo, + +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, +#endif + }; for (auto& function_table_entry : function_table) { KernelCreateInfo info = function_table_entry(); diff --git a/onnxruntime/contrib_ops/cuda/layer_norm.cc b/onnxruntime/contrib_ops/cuda/layer_norm.cc index 685fd20da8..d417c62cc6 100644 --- a/onnxruntime/contrib_ops/cuda/layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/layer_norm.cc @@ -36,6 +36,9 @@ namespace cuda { REGISTER_KERNEL_TYPED(float, float) REGISTER_KERNEL_TYPED(double, double) REGISTER_KERNEL_TYPED(MLFloat16, float) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_KERNEL_TYPED(BFloat16, float) +#endif template LayerNorm::LayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { diff --git a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu index 30c2f6386c..ecbaad2b08 100644 --- a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu @@ -389,6 +389,10 @@ LAYERNORM_LINEAR_IMPL(half, float, false) LAYERNORM_LINEAR_IMPL(double, double, false) //LAYERNORM_LINEAR_IMPL(half, half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +LAYERNORM_LINEAR_IMPL(nv_bfloat16, float, true) +LAYERNORM_LINEAR_IMPL(nv_bfloat16, float, false) +#endif } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc b/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc index 98cc601a6f..a39a6ee3b5 100644 --- a/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc +++ b/onnxruntime/contrib_ops/cuda/math/fused_matmul.cc @@ -26,6 +26,11 @@ REGISTER_KERNEL_TYPED(FusedMatMul, float) REGISTER_KERNEL_TYPED(FusedMatMul, double) REGISTER_KERNEL_TYPED(FusedMatMul, MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_KERNEL_TYPED(TransposeMatMul, BFloat16) +REGISTER_KERNEL_TYPED(FusedMatMul, BFloat16) +#endif + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/dropout_elimination.cc b/onnxruntime/core/optimizer/dropout_elimination.cc index 4bca183293..84340f0301 100644 --- a/onnxruntime/core/optimizer/dropout_elimination.cc +++ b/onnxruntime/core/optimizer/dropout_elimination.cc @@ -32,7 +32,7 @@ bool EliminateDropout::SatisfyCondition(const Graph& graph, const Node& node, co // 2. ratio input is not a graph input, so it cannot be overridden // support opset 12 and above for ort training - if (graph_utils::MatchesOpSinceVersion(node, {12}) && node.InputDefs().size() > 1) { + if (graph_utils::MatchesOpSinceVersion(node, {12, 13}) && node.InputDefs().size() > 1) { if (graph_utils::IsGraphInput(graph, node.InputDefs()[1])) { return false; } diff --git a/onnxruntime/core/providers/cuda/activation/activations.cc b/onnxruntime/core/providers/cuda/activation/activations.cc index cfd2e698aa..cb985fd024 100644 --- a/onnxruntime/core/providers/cuda/activation/activations.cc +++ b/onnxruntime/core/providers/cuda/activation/activations.cc @@ -57,8 +57,15 @@ namespace cuda { REGISTER_ACTIVATION_KERNEL(name, ver, T) \ UNARY_ACTIVATION_COMPUTE(name, T) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define UNARY_ACTIVATION_OP_TYPED_BF16(name, ver) UNARY_ACTIVATION_OP_TYPED(name, ver, BFloat16) +#else +#define UNARY_ACTIVATION_OP_TYPED_BF16(name, ver) +#endif + #define UNARY_ACTIVATION_OP_HFD(name, ver) \ UNARY_ACTIVATION_OP_TYPED(name, ver, MLFloat16) \ + UNARY_ACTIVATION_OP_TYPED_BF16(name, ver) \ UNARY_ACTIVATION_OP_TYPED(name, ver, float) \ UNARY_ACTIVATION_OP_TYPED(name, ver, double) diff --git a/onnxruntime/core/providers/cuda/activation/activations_impl.cu b/onnxruntime/core/providers/cuda/activation/activations_impl.cu index a0474fe1f4..bd7a4f7bf8 100644 --- a/onnxruntime/core/providers/cuda/activation/activations_impl.cu +++ b/onnxruntime/core/providers/cuda/activation/activations_impl.cu @@ -93,9 +93,16 @@ struct OP_ThresholdedRelu : public CtxThresholdedRelu { #define SPECIALIZED_UNARY_ACTIVATION_IMPL(name, T) \ template void Impl_##name(const T* input_data, T* output_data, const Ctx##name* func_ctx, size_t count); -#define SPECIALIZED_UNARY_ACTIVATIONL_HFD(name) \ - SPECIALIZED_UNARY_ACTIVATION_IMPL(name, half) \ - SPECIALIZED_UNARY_ACTIVATION_IMPL(name, float) \ +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +#define SPECIALIZED_UNARY_ACTIVATION_IMPL_BF16(name) SPECIALIZED_UNARY_ACTIVATION_IMPL(name, nv_bfloat16) +#else +#define SPECIALIZED_UNARY_ACTIVATION_IMPL_BF16(name) +#endif + +#define SPECIALIZED_UNARY_ACTIVATIONL_HFD(name) \ + SPECIALIZED_UNARY_ACTIVATION_IMPL(name, half) \ + SPECIALIZED_UNARY_ACTIVATION_IMPL_BF16(name) \ + SPECIALIZED_UNARY_ACTIVATION_IMPL(name, float) \ SPECIALIZED_UNARY_ACTIVATION_IMPL(name, double) #define UNARY_ACTIVATION_OP_NAME(name) \ diff --git a/onnxruntime/core/providers/cuda/atomic/common.cuh b/onnxruntime/core/providers/cuda/atomic/common.cuh index 78c26715ff..098f650986 100644 --- a/onnxruntime/core/providers/cuda/atomic/common.cuh +++ b/onnxruntime/core/providers/cuda/atomic/common.cuh @@ -21,6 +21,10 @@ #include "cuda_fp16.h" #include "cuda_runtime.h" +#if CUDA_VERSION >= 11000 +#include "cuda_bf16.h" +#endif + namespace onnxruntime { namespace cuda { @@ -68,5 +72,22 @@ __device__ __forceinline__ void atomic_add(half *address, half value) { #endif } +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +__device__ __forceinline__ void atomic_add(nv_bfloat16 *address, nv_bfloat16 value) { + unsigned int * base_address = reinterpret_cast(reinterpret_cast(address) - (reinterpret_cast(address) & 2)); + unsigned int old = *base_address; + unsigned int assumed; + unsigned short x; + + do { + assumed = old; + x = reinterpret_cast(address) & 2 ? (old >> 16) : (old & 0xffff); + x = __bfloat16_as_short(__float2bfloat16(__bfloat162float(*reinterpret_cast(&x)) + __bfloat162float(value))); + old = reinterpret_cast(address) & 2 ? (old & 0xffff) | (x << 16) : (old & 0xffff0000) | x; + old = atomicCAS(base_address, assumed, old); + } while (assumed != old); +} +#endif + } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/cuda/cu_inc/common.cuh b/onnxruntime/core/providers/cuda/cu_inc/common.cuh index 93652e792d..7d5cd9f397 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/common.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/common.cuh @@ -11,6 +11,10 @@ #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/cuda_call.h" +#if CUDA_VERSION >= 11000 +#include +#endif + namespace onnxruntime { namespace cuda { @@ -222,6 +226,30 @@ __device__ __inline__ T _Gelu(T a) { return a * _Normcdf(a); } +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template <> +__device__ __inline__ nv_bfloat16 _Sqrt(nv_bfloat16 a) { return nv_bfloat16(sqrtf(static_cast(a))); } + +template <> +__device__ __inline__ nv_bfloat16 _Exp(nv_bfloat16 a) { return nv_bfloat16(expf(static_cast(a))); } + +template <> +__device__ __inline__ nv_bfloat16 _Log(nv_bfloat16 a) { return nv_bfloat16(logf(static_cast(a))); } + +template <> +__device__ __inline__ nv_bfloat16 _Tanh(nv_bfloat16 a) { return nv_bfloat16(tanhf(static_cast(a))); } + +template <> +__device__ __inline__ nv_bfloat162 _Tanh(nv_bfloat162 a) { + float2 tmp = (__bfloat1622float2(a)); + tmp.x = tanhf(tmp.x); + tmp.y = tanhf(tmp.y); + return __float22bfloat162_rn(tmp); +} + +template <> +__device__ __inline__ nv_bfloat16 _Normcdf(nv_bfloat16 a) { return nv_bfloat16(normcdff(static_cast(a))); } +#endif // We would like to use 64-bit integer to support large matrices. However, CUDA seems to support only 32-bit integer // For now, use int32_t to ensure that both Linux and Windows see this as 32 bit integer type. diff --git a/onnxruntime/core/providers/cuda/cuda_common.h b/onnxruntime/core/providers/cuda/cuda_common.h index d978c84f67..f3994bf43f 100644 --- a/onnxruntime/core/providers/cuda/cuda_common.h +++ b/onnxruntime/core/providers/cuda/cuda_common.h @@ -67,6 +67,18 @@ class ToCudaType { } }; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +template <> +class ToCudaType { + public: + typedef nv_bfloat16 MappedType; + static MappedType FromFloat(float f) { + uint16_t h = BFloat16(f).val; + return *reinterpret_cast(&h); + } +}; +#endif + inline bool CalculateFdmStrides(gsl::span p, const std::vector& dims) { int stride = 1; if (dims.empty() || p.size() < dims.size()) diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index e1a6b339a6..723eb276a2 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -1001,6 +1001,21 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, Identity); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, ScatterND); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Add); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Sub); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Mul); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Div); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Cast); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Softmax); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, MatMul); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Relu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Sigmoid); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum); +#endif + template <> KernelCreateInfo BuildKernelCreateInfo() { KernelCreateInfo info; @@ -1702,6 +1717,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -1715,6 +1731,21 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, +#endif }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index 888c7dcf57..d70ca20be2 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h @@ -122,6 +122,13 @@ class CUDAExecutionProvider : public IExecutionProvider { constant_ones_half_ = cuda::CreateConstantOnes(); } return reinterpret_cast(constant_ones_half_->GetBuffer(count)); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + } else if (std::is_same::value) { + if (!constant_ones_bfloat16_) { + constant_ones_bfloat16_ = cuda::CreateConstantOnes(); + } + return reinterpret_cast(constant_ones_bfloat16_->GetBuffer(count)); +#endif } else { return nullptr; } @@ -143,6 +150,9 @@ class CUDAExecutionProvider : public IExecutionProvider { std::unique_ptr> constant_ones_float_; std::unique_ptr> constant_ones_double_; std::unique_ptr> constant_ones_half_; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + std::unique_ptr> constant_ones_bfloat16_; +#endif AllocatorPtr allocator_; }; diff --git a/onnxruntime/core/providers/cuda/cuda_utils.cu b/onnxruntime/core/providers/cuda/cuda_utils.cu index 19ca3e8d65..a0ee56c6a3 100644 --- a/onnxruntime/core/providers/cuda/cuda_utils.cu +++ b/onnxruntime/core/providers/cuda/cuda_utils.cu @@ -71,6 +71,9 @@ std::unique_ptr> CreateConstantOnes() { template std::unique_ptr> CreateConstantOnes(); template std::unique_ptr> CreateConstantOnes(); template std::unique_ptr> CreateConstantOnes(); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template std::unique_ptr> CreateConstantOnes(); +#endif #define SPECIALIZED_FILL(T) \ template void Fill(T * output, T value, int64_t count); @@ -82,6 +85,9 @@ SPECIALIZED_FILL(int64_t) SPECIALIZED_FILL(float) SPECIALIZED_FILL(double) SPECIALIZED_FILL(__half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_FILL(nv_bfloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cudnn_common.cc b/onnxruntime/core/providers/cuda/cudnn_common.cc index 200fa352a5..24f8220eb9 100644 --- a/onnxruntime/core/providers/cuda/cudnn_common.cc +++ b/onnxruntime/core/providers/cuda/cudnn_common.cc @@ -161,6 +161,11 @@ const float Consts::Zero = 0; const float Consts::One = 1; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +const float Consts::Zero = 0; +const float Consts::One = 1; +#endif + template <> const int8_t Consts::Zero = 0; diff --git a/onnxruntime/core/providers/cuda/cudnn_common.h b/onnxruntime/core/providers/cuda/cudnn_common.h index 1addd77c75..39c1b74405 100644 --- a/onnxruntime/core/providers/cuda/cudnn_common.h +++ b/onnxruntime/core/providers/cuda/cudnn_common.h @@ -126,6 +126,14 @@ struct Consts { static const float One; }; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +template<> +struct Consts { + static const float Zero; + static const float One; +}; +#endif + inline double ClampCudnnBatchNormEpsilon(double epsilon) { if (epsilon < CUDNN_BN_MIN_EPSILON) { if (CUDNN_BN_MIN_EPSILON - epsilon > FLT_EPSILON) diff --git a/onnxruntime/core/providers/cuda/fpgeneric.cu b/onnxruntime/core/providers/cuda/fpgeneric.cu index 0d096148b5..417d49ae2c 100644 --- a/onnxruntime/core/providers/cuda/fpgeneric.cu +++ b/onnxruntime/core/providers/cuda/fpgeneric.cu @@ -55,6 +55,14 @@ __global__ void CopyVectorHalf(const half* x, int incx, half* y, int incy, int n y[id * incy] = x[id * incx]; } +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +__global__ void CopyVectorBFloat16(const nv_bfloat16* x, int incx, nv_bfloat16* y, int incy, int n) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + if (id >= n) return; + y[id * incy] = x[id * incx]; +} +#endif + } // namespace cublasStatus_t cublasTransposeHelper(cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, const half*, const half* A, int, const half*, const half*, int, half* C, int) { @@ -74,4 +82,13 @@ cublasStatus_t cublasCopyHelper(cublasHandle_t, int n, const half* x, int incx, dim3 dimBlock(COPY_BLOCK_DIM, 1, 1); CopyVectorHalf<<>>(x, incx, y, incy, n); return CUBLAS_STATUS_SUCCESS; -} \ No newline at end of file +} + +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +cublasStatus_t cublasCopyHelper(cublasHandle_t, int n, const nv_bfloat16* x, int incx, nv_bfloat16* y, int incy) { + dim3 dimGrid((unsigned int)(n + COPY_BLOCK_DIM - 1) / COPY_BLOCK_DIM, 1, 1); + dim3 dimBlock(COPY_BLOCK_DIM, 1, 1); + CopyVectorBFloat16<<>>(x, incx, y, incy, n); + return CUBLAS_STATUS_SUCCESS; +} +#endif diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc index f8336dd4db..c733fb85df 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc @@ -188,9 +188,23 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin // D: double // O: bool -#define BINARY_OP_VERSIONED_HFD(name, startver, endver) \ - BINARY_OP_VERSIONED_TYPED(name, startver, endver, MLFloat16) \ - BINARY_OP_VERSIONED_TYPED(name, startver, endver, float) \ +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define BINARY_OP_TYPED_BF16(name, ver) BINARY_OP_TYPED(name, ver, BFloat16) +#define BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED_BF16(name, ver) BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, BFloat16) +#define BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED_BF16(name, ver) BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, BFloat16) +#define BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED_BF16(name, startver, endver) BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, BFloat16) +#define BINARY_OP_TYPED_VERSIONED_V_BF16(name, class_name, startver, endver) BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, BFloat16) +#else +#define BINARY_OP_TYPED_BF16(name, ver) +#define BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED_BF16(name, ver) +#define BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED_BF16(name, ver) +#define BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED_BF16(name, startver, endver) +#define BINARY_OP_TYPED_VERSIONED_V_BF16(name, class_name, startver, endver) +#endif + +#define BINARY_OP_VERSIONED_HFD(name, startver, endver) \ + BINARY_OP_VERSIONED_TYPED(name, startver, endver, MLFloat16) \ + BINARY_OP_VERSIONED_TYPED(name, startver, endver, float) \ BINARY_OP_VERSIONED_TYPED(name, startver, endver, double) #define BINARY_OP_VERSIONED_UZILHFD(name, startver, endver) \ @@ -202,6 +216,7 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin #define BINARY_OP_HFD(name, ver) \ BINARY_OP_TYPED(name, ver, MLFloat16) \ + BINARY_OP_TYPED_BF16(name, ver) \ BINARY_OP_TYPED(name, ver, float) \ BINARY_OP_TYPED(name, ver, double) @@ -224,6 +239,7 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin #define BINARY_OP_REGISTER_HFD(name, ver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, MLFloat16) \ + BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED_BF16(name, ver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, float) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, double) @@ -240,16 +256,19 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, int32_t) \ BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, int64_t) \ BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, MLFloat16) \ + BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED_BF16(name, ver) \ BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, float) \ BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, double) #define BINARY_OP_REGISTER_VERSIONED_HFD(name, startver, endver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, MLFloat16) \ + BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED_BF16(name, startver, endver) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, float) \ BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, double) #define BINARY_OP_REGISTER_VERSIONED_CLASS_HFD(name, class_name, startver, endver) \ BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, MLFloat16) \ + BINARY_OP_TYPED_VERSIONED_V_BF16(name, class_name, startver, endver) \ BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, float) \ BINARY_OP_TYPED_VERSIONED_V(name, class_name, startver, endver, double) diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu index 6fc6d7dd1f..f0cb62faaa 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu @@ -73,12 +73,21 @@ namespace cuda { const TArray* rhs_padded_strides, const T2* rhs_data, \ const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_t count); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, nv_bfloat16) +#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2_BF16(name) SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, nv_bfloat16, nv_bfloat16) +#else +#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) +#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2_BF16(name) +#endif + #define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint32_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint64_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int64_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) @@ -96,6 +105,7 @@ namespace cuda { #define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_BF16(x) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) @@ -156,6 +166,7 @@ BINARY_OPS2() SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, int32_t, int32_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, int64_t, int64_t) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, half, half) \ + SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2_BF16(name) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, float, float) \ SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(name, bool, double, double) diff --git a/onnxruntime/core/providers/cuda/math/gemm.cc b/onnxruntime/core/providers/cuda/math/gemm.cc index 04f47fb3d8..03819891e9 100644 --- a/onnxruntime/core/providers/cuda/math/gemm.cc +++ b/onnxruntime/core/providers/cuda/math/gemm.cc @@ -53,6 +53,9 @@ namespace cuda { REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(double) REGISTER_KERNEL_TYPED(MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_KERNEL_TYPED(BFloat16) +#endif template Status Gemm::ComputeInternal(OpKernelContext* ctx) const { diff --git a/onnxruntime/core/providers/cuda/math/matmul.cc b/onnxruntime/core/providers/cuda/math/matmul.cc index cac06e6d8d..e0ec185bdb 100644 --- a/onnxruntime/core/providers/cuda/math/matmul.cc +++ b/onnxruntime/core/providers/cuda/math/matmul.cc @@ -41,6 +41,9 @@ namespace cuda { REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(double) REGISTER_KERNEL_TYPED(MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_KERNEL_TYPED(BFloat16) +#endif // StridedBatchedGemm can be used for the following GEMM computation // C[pnm] = A[pnk]*B[km] or C[pnm] = A[pnk]*B[pkm] diff --git a/onnxruntime/core/providers/cuda/math/softmax.cc b/onnxruntime/core/providers/cuda/math/softmax.cc index f80eafe2fa..09753b66a2 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.cc +++ b/onnxruntime/core/providers/cuda/math/softmax.cc @@ -57,6 +57,30 @@ SPECIALIZED_SOFTMAX_HELPER_IMPL(float) SPECIALIZED_SOFTMAX_HELPER_IMPL(double) SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +// cudnnSoftmaxForward/Backward doesn't support BFloat16. +#define SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(is_log_softmax) \ + template <> \ + Status SoftMaxComputeHelper( \ + const BFloat16* X, \ + const TensorShape& input_shape, \ + BFloat16* Y, \ + cudnnHandle_t, \ + int64_t axis) { \ + typedef typename ToCudaType::MappedType CudaT; \ + int64_t N = input_shape.SizeToDimension(axis); \ + int64_t D = input_shape.SizeFromDimension(axis); \ + auto Y_data = reinterpret_cast(Y); \ + auto X_data = reinterpret_cast(X); \ + dispatch_softmax_forward, is_log_softmax>( \ + Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); \ + return Status::OK(); \ + } + +SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(true) +SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(false) +#endif + #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ Softmax, \ @@ -217,6 +241,9 @@ Status Softmax::ComputeInternal(OpKernelContext* ctx) const { SPECIALIZED_COMPUTE(float) SPECIALIZED_COMPUTE(double) SPECIALIZED_COMPUTE(MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +SPECIALIZED_COMPUTE(BFloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/math/softmax_impl.cu b/onnxruntime/core/providers/cuda/math/softmax_impl.cu index 6aec65a38c..f4658e93fe 100644 --- a/onnxruntime/core/providers/cuda/math/softmax_impl.cu +++ b/onnxruntime/core/providers/cuda/math/softmax_impl.cu @@ -214,6 +214,9 @@ template void dispatch_softmax_forward(output_t SPECIALIZED_SOFTMAX_IMPL(float, float, float) SPECIALIZED_SOFTMAX_IMPL(half, half, float) SPECIALIZED_SOFTMAX_IMPL(double, double, double) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_SOFTMAX_IMPL(nv_bfloat16, nv_bfloat16, float) +#endif } } \ No newline at end of file diff --git a/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu index 41a93cf4cc..fe60e66856 100644 --- a/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu @@ -91,11 +91,24 @@ struct ViaTypeMap { typedef float ViaT; }; +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template <> +struct ViaTypeMap { + typedef float ViaT; +}; +#endif + template struct OP_Cast { __device__ __inline__ OutT operator()(const InT& a) const { const bool any_float16 = std::is_same::value || std::is_same::value; +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + const bool any_bf16 = std::is_same::value || std::is_same::value; + typedef typename std::conditional::type T1; + typedef typename std::conditional::type T; +#else typedef typename std::conditional::type T; +#endif typedef typename ViaTypeMap::ViaT ViaT; return (OutT)((ViaT)a); } @@ -115,8 +128,15 @@ void Impl_Cast( #define SPECIALIZED_CAST_IMPL2(InT, OutT) \ template void Impl_Cast(const InT* input_data, OutT* output_data, size_t count); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +#define SPECIALIZED_CAST_IMPL2_BF16(T) SPECIALIZED_CAST_IMPL2(T, nv_bfloat16) +#else +#define SPECIALIZED_CAST_IMPL2_BF16(T) +#endif + #define SPECIALIZED_CAST_FROM(T) \ SPECIALIZED_CAST_IMPL2(T, half) \ + SPECIALIZED_CAST_IMPL2_BF16(T) \ SPECIALIZED_CAST_IMPL2(T, float) \ SPECIALIZED_CAST_IMPL2(T, double) \ SPECIALIZED_CAST_IMPL2(T, int8_t) \ @@ -141,6 +161,9 @@ SPECIALIZED_CAST_FROM(uint16_t) SPECIALIZED_CAST_FROM(uint32_t) SPECIALIZED_CAST_FROM(uint64_t) SPECIALIZED_CAST_FROM(bool) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_CAST_FROM(nv_bfloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc index c6e1fb1ded..beb3e829f3 100644 --- a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc +++ b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc @@ -194,22 +194,28 @@ Status VariadicElementwiseOp namespace { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double, BFloat16 +#else +#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double +#endif + using SumOp = VariadicElementwiseOp< variadic_elementwise_ops::Sum, - MLFloat16, float, double>; + ALL_IEEE_FLOAT_DATA_TYPES>; using MinOp = VariadicElementwiseOp< variadic_elementwise_ops::Min, - uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double>; + uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>; using MaxOp = VariadicElementwiseOp< variadic_elementwise_ops::Max, - uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double>; + uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>; const auto k_uzilhfd_datatypes = - BuildKernelDefConstraints(); + BuildKernelDefConstraints(); const auto k_hfd_datatypes = - BuildKernelDefConstraints(); + BuildKernelDefConstraints(); } // namespace diff --git a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu index a960226524..ad975f85c3 100644 --- a/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu @@ -121,8 +121,15 @@ void Impl_NoBroadcastInputBatch( // D: double // O: bool +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +#define SPECIALIZE_IMPL_BF16(VariadicElementwiseOpTag) SPECIALIZE_IMPL(nv_bfloat16, VariadicElementwiseOpTag) +#else +#define SPECIALIZE_IMPL_BF16(VariadicElementwiseOpTag) +#endif + #define SPECIALIZE_IMPL_HFD(VariadicElementwiseOpTag) \ SPECIALIZE_IMPL(half, VariadicElementwiseOpTag) \ + SPECIALIZE_IMPL_BF16(VariadicElementwiseOpTag) \ SPECIALIZE_IMPL(float, VariadicElementwiseOpTag) \ SPECIALIZE_IMPL(double, VariadicElementwiseOpTag) diff --git a/onnxruntime/core/providers/cuda/nn/dropout.cc b/onnxruntime/core/providers/cuda/nn/dropout.cc index fad4c4e787..13e154e5db 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout.cc +++ b/onnxruntime/core/providers/cuda/nn/dropout.cc @@ -6,6 +6,15 @@ namespace onnxruntime { namespace cuda { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()} +#else +#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() +#endif + ONNX_OPERATOR_VERSIONED_KERNEL_EX( Dropout, kOnnxDomain, @@ -25,8 +34,8 @@ ONNX_OPERATOR_KERNEL_EX( 13, kCudaExecutionProvider, KernelDefBuilder() - .TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes()) - .TypeConstraint("T1", DataTypeImpl::AllIEEEFloatTensorTypes()) + .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) + .TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) .TypeConstraint("T2", DataTypeImpl::GetTensorType()) .InputMemoryType(1) .InputMemoryType(2), diff --git a/onnxruntime/core/providers/cuda/nn/dropout.h b/onnxruntime/core/providers/cuda/nn/dropout.h index 1c25f7ac7a..575a037d13 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout.h +++ b/onnxruntime/core/providers/cuda/nn/dropout.h @@ -105,7 +105,11 @@ Status Dropout::ComputeInternal(OpKernelContext* context) con PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default(); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + utils::MLTypeCallDispatcher t_disp(X->GetElementType()); +#else utils::MLTypeCallDispatcher t_disp(X->GetElementType()); +#endif t_disp.Invoke(GetDeviceProp(), N, ratio_data, generator, *X, *Y, mask_data); return Status::OK(); diff --git a/onnxruntime/core/providers/cuda/nn/dropout_impl.cu b/onnxruntime/core/providers/cuda/nn/dropout_impl.cu index 0abe3f774a..ded4a87c40 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout_impl.cu +++ b/onnxruntime/core/providers/cuda/nn/dropout_impl.cu @@ -99,6 +99,9 @@ void DropoutKernelImpl( SPECIALIZED_DROPOUT_IMPL(float) SPECIALIZED_DROPOUT_IMPL(double) SPECIALIZED_DROPOUT_IMPL(half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_DROPOUT_IMPL(nv_bfloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index f78d570d1c..ad8533b00c 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -353,6 +353,9 @@ INSTANTIATE_REDUCE_SUM(double, double); INSTANTIATE_REDUCE_SQUARE_SUM(half, float); INSTANTIATE_REDUCE_SQUARE_SUM(float, float); INSTANTIATE_REDUCE_SQUARE_SUM(double, double); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_REDUCE_SQUARE_SUM(nv_bfloat16, float); +#endif #undef INSTANTIATE_REDUCE_SQUARE_SUM #define INSTANTIATE_REDUCE_L2_NORM(TIn, TOut) \ @@ -465,6 +468,9 @@ Status reduce_matrix_rows(const TIn* input, TOut* output, int m, int n, bool res INSTANTIATE_REDUCE_MATRIX_ROWS(half); INSTANTIATE_REDUCE_MATRIX_ROWS(float); INSTANTIATE_REDUCE_MATRIX_ROWS(double); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_REDUCE_MATRIX_ROWS(nv_bfloat16); +#endif #undef INSTANTIATE_REDUCE_MATRIX_ROWS template @@ -478,6 +484,9 @@ Status reduce_matrix_columns(const TIn* input, TOut* output, int m, int n, void* INSTANTIATE_REDUCE_MATRIX_COLUMNS(half); INSTANTIATE_REDUCE_MATRIX_COLUMNS(float); INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_REDUCE_MATRIX_COLUMNS(nv_bfloat16); +#endif #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index 8f5b1bb45a..1b4e524e2a 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -667,17 +667,104 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe calculate_log_, calculate_sqt_, log_sum_exp_, fast_reduction); } -template <> -template <> -Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { - typedef typename ToCudaType::MappedType CudaT; +#define SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(T) \ + template <> \ + template <> \ + Status ReduceKernel::ComputeImpl( \ + OpKernelContext * ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { \ + typedef typename ToCudaType::MappedType CudaT; \ + const Tensor* X = ctx->Input(0); \ + std::vector axes; \ + size_t num_inputs = ctx->InputCount(); \ + if (num_inputs == 2) { \ + const Tensor* axes_tensor = ctx->Input(1); \ + ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); \ + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); \ + auto nDims = static_cast(axes_tensor->Shape()[0]); \ + const auto* data = axes_tensor->template Data(); \ + axes.assign(data, data + nDims); \ + } else { \ + axes.assign(axes_.begin(), axes_.end()); \ + } \ + \ + if (axes.empty() && noop_with_empty_axes_) { \ + auto* Y = ctx->Output(0, X->Shape()); \ + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), X->SizeInBytes(), \ + cudaMemcpyDeviceToDevice)); \ + return Status::OK(); \ + } \ + \ + PrepareReduceMetadata prepare_reduce_metadata; \ + ORT_RETURN_IF_ERROR(PrepareForReduce(X, keepdims_, axes, prepare_reduce_metadata)); \ + \ + Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); \ + \ + int64_t input_count = prepare_reduce_metadata.input_count; \ + int64_t output_count = prepare_reduce_metadata.output_count; \ + std::vector& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn; \ + std::vector& output_dims_cudnn = prepare_reduce_metadata.output_dims_cudnn; \ + \ + if (input_count == 0) { \ + assert(Y->Shape().Size() == 0); \ + return Status::OK(); \ + } \ + \ + if (input_count == output_count) { \ + if (Y->template MutableData() != X->template Data()) { \ + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), \ + input_count * sizeof(T), cudaMemcpyDeviceToDevice)); \ + } \ + return Status::OK(); \ + } \ + \ + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes())); \ + \ + size_t indices_bytes = 0; \ + size_t workspace_bytes = 0; \ + CudnnTensor input_tensor; \ + CudnnTensor output_tensor; \ + CudnnReduceDescriptor reduce_desc; \ + \ + cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT; \ + IAllocatorUniquePtr temp_X = GetScratchBuffer(input_count); \ + Impl_Cast(reinterpret_cast(X->template Data()), temp_X.get(), X->Shape().Size()); \ + \ + ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, CUDNN_REDUCE_TENSOR_FLATTENED_INDICES)); \ + ORT_RETURN_IF_ERROR(input_tensor.Set(input_dims_cudnn, cudnn_type_X)); \ + ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_cudnn, cudnn_type_X)); \ + CUDNN_RETURN_IF_ERROR( \ + cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes)); \ + CUDNN_RETURN_IF_ERROR( \ + cudnnGetReductionWorkspaceSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes)); \ + IAllocatorUniquePtr indices_cuda = GetScratchBuffer(indices_bytes); \ + IAllocatorUniquePtr workspace_cuda = GetScratchBuffer(workspace_bytes); \ + \ + const auto one = Consts::One; \ + const auto zero = Consts::Zero; \ + auto temp_Y = GetScratchBuffer(output_count); \ + CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, \ + workspace_cuda.get(), workspace_bytes, &one, input_tensor, temp_X.get(), \ + &zero, output_tensor, temp_Y.get())); \ + \ + Impl_Cast(temp_Y.get(), reinterpret_cast(Y->template MutableData()), output_count); \ + \ + return Status::OK(); \ + } +SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int32_t) +SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int8_t) +SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(uint8_t) + +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +template <> +template <> +Status ReduceKernel::ComputeImpl( + OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { + typedef typename ToCudaType::MappedType CudaT; const Tensor* X = ctx->Input(0); std::vector axes; - size_t num_inputs = ctx->InputCount(); if (num_inputs == 2) { - //override the attribute value with the input value for reduction_axes const Tensor* axes_tensor = ctx->Input(1); ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); @@ -688,19 +775,15 @@ Status ReduceKernel::ComputeImpl( axes.assign(axes_.begin(), axes_.end()); } - // empty axes and no-op if (axes.empty() && noop_with_empty_axes_) { auto* Y = ctx->Output(0, X->Shape()); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), X->SizeInBytes(), cudaMemcpyDeviceToDevice)); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), + X->SizeInBytes(), cudaMemcpyDeviceToDevice)); return Status::OK(); } PrepareReduceMetadata prepare_reduce_metadata; - - ORT_RETURN_IF_ERROR(PrepareForReduce(X, - keepdims_, - axes, - prepare_reduce_metadata)); + ORT_RETURN_IF_ERROR(PrepareForReduce(X, keepdims_, axes, prepare_reduce_metadata)); Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); @@ -709,22 +792,40 @@ Status ReduceKernel::ComputeImpl( std::vector& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn; std::vector& output_dims_cudnn = prepare_reduce_metadata.output_dims_cudnn; - // special case when there is a dim value of 0 in the shape. if (input_count == 0) { assert(Y->Shape().Size() == 0); return Status::OK(); } - // cudnnReduceTensor for ReduceSum has issue if input and output has same size, we just need to copy the data for this case if (input_count == output_count) { - if (Y->template MutableData() != X->template Data()) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), input_count * sizeof(int32_t), cudaMemcpyDeviceToDevice)); + if (Y->template MutableData() != X->template Data()) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), + input_count * sizeof(BFloat16), cudaMemcpyDeviceToDevice)); } return Status::OK(); } - // This reduction keep adding values to this buffer. If a non-zero value, say 1000, is here, the sum will start with 1000. - // Therefore zeroing out the memory is required + if (fast_reduction_ && !ctx->GetUseDeterministicCompute()) { + int m{}, n{}; + const auto applicable_matrix_reduction = + get_applicable_matrix_reduction(cudnn_reduce_op, X->Shape().GetDims(), axes, m, n); + switch (applicable_matrix_reduction) { + case ApplicableMatrixReduction::Rows: { + return reduce_matrix_rows(reinterpret_cast(X->template Data()), + reinterpret_cast(Y->template MutableData()), m, n); + } + case ApplicableMatrixReduction::Columns: { + const auto buffer_size_bytes = compute_reduce_matrix_columns_buffer_size(m, n); + auto buffer = cuda_ep_->GetScratchBuffer(buffer_size_bytes); + return reduce_matrix_columns(reinterpret_cast(X->template Data()), + reinterpret_cast(Y->template MutableData()), m, n, buffer.get(), + buffer_size_bytes); + } + default: + break; + } + } + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes())); size_t indices_bytes = 0; @@ -735,194 +836,31 @@ Status ReduceKernel::ComputeImpl( cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT; IAllocatorUniquePtr temp_X = GetScratchBuffer(input_count); - Impl_Cast(reinterpret_cast(X->template Data()), temp_X.get(), X->Shape().Size()); + Impl_Cast(reinterpret_cast(X->template Data()), temp_X.get(), + X->Shape().Size()); ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, CUDNN_REDUCE_TENSOR_FLATTENED_INDICES)); ORT_RETURN_IF_ERROR(input_tensor.Set(input_dims_cudnn, cudnn_type_X)); ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_cudnn, cudnn_type_X)); - CUDNN_RETURN_IF_ERROR(cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes)); - CUDNN_RETURN_IF_ERROR(cudnnGetReductionWorkspaceSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes)); + CUDNN_RETURN_IF_ERROR( + cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes)); + CUDNN_RETURN_IF_ERROR( + cudnnGetReductionWorkspaceSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes)); IAllocatorUniquePtr indices_cuda = GetScratchBuffer(indices_bytes); IAllocatorUniquePtr workspace_cuda = GetScratchBuffer(workspace_bytes); const auto one = Consts::One; const auto zero = Consts::Zero; auto temp_Y = GetScratchBuffer(output_count); - CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(CudnnHandle(), - reduce_desc, - indices_cuda.get(), - indices_bytes, - workspace_cuda.get(), - workspace_bytes, - &one, - input_tensor, - temp_X.get(), - &zero, - output_tensor, - temp_Y.get())); + CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, + workspace_cuda.get(), workspace_bytes, &one, input_tensor, temp_X.get(), + &zero, output_tensor, temp_Y.get())); - Impl_Cast(temp_Y.get(), Y->template MutableData(), output_count); - - return Status::OK(); -} - -template <> -template <> -Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { - typedef typename ToCudaType::MappedType CudaT; - - const Tensor* X = ctx->Input(0); - PrepareReduceMetadata prepare_reduce_metadata; - - ORT_RETURN_IF_ERROR(PrepareForReduce(X, - keepdims_, - axes_, - prepare_reduce_metadata)); - - Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); - - int64_t input_count = prepare_reduce_metadata.input_count; - int64_t output_count = prepare_reduce_metadata.output_count; - std::vector& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn; - std::vector& output_dims_cudnn = prepare_reduce_metadata.output_dims_cudnn; - - // special case when there is a dim value of 0 in the shape. - if (input_count == 0) { - assert(Y->Shape().Size() == 0); - return Status::OK(); - } - - // cudnnReduceTensor has issue if input and output has same size, we just need to copy the data for this case - auto* const dst = Y->template MutableData(); - const auto* const src = X->template Data(); - if (input_count == output_count) { - if (src != dst) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst, src, input_count * sizeof(int8_t), cudaMemcpyDeviceToDevice)); - } - return Status::OK(); - } - - // This reduction keep adding values to this buffer. If a non-zero value, say 1000, is here, the sum will start with 1000. - // Therefore zeroing out the memory is required - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes())); - - size_t indices_bytes = 0; - size_t workspace_bytes = 0; - CudnnTensor input_tensor; - CudnnTensor output_tensor; - CudnnReduceDescriptor reduce_desc; - - cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT; - IAllocatorUniquePtr temp_X = GetScratchBuffer(input_count); - Impl_Cast(reinterpret_cast(src), temp_X.get(), X->Shape().Size()); - - ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, CUDNN_REDUCE_TENSOR_FLATTENED_INDICES)); - ORT_RETURN_IF_ERROR(input_tensor.Set(input_dims_cudnn, cudnn_type_X)); - ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_cudnn, cudnn_type_X)); - CUDNN_RETURN_IF_ERROR(cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes)); - CUDNN_RETURN_IF_ERROR(cudnnGetReductionWorkspaceSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes)); - IAllocatorUniquePtr indices_cuda = GetScratchBuffer(indices_bytes); - IAllocatorUniquePtr workspace_cuda = GetScratchBuffer(workspace_bytes); - - const auto one = Consts::One; - const auto zero = Consts::Zero; - auto temp_Y = GetScratchBuffer(output_count); - CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(CudnnHandle(), - reduce_desc, - indices_cuda.get(), - indices_bytes, - workspace_cuda.get(), - workspace_bytes, - &one, - input_tensor, - temp_X.get(), - &zero, - output_tensor, - temp_Y.get())); - - Impl_Cast(temp_Y.get(), dst, output_count); - - return Status::OK(); -} - -template <> -template <> -Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { - typedef typename ToCudaType::MappedType CudaT; - - const Tensor* X = ctx->Input(0); - PrepareReduceMetadata prepare_reduce_metadata; - - ORT_RETURN_IF_ERROR(PrepareForReduce(X, - keepdims_, - axes_, - prepare_reduce_metadata)); - - Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); - - int64_t input_count = prepare_reduce_metadata.input_count; - int64_t output_count = prepare_reduce_metadata.output_count; - std::vector& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn; - std::vector& output_dims_cudnn = prepare_reduce_metadata.output_dims_cudnn; - - // special case when there is a dim value of 0 in the shape. - if (input_count == 0) { - assert(Y->Shape().Size() == 0); - return Status::OK(); - } - - // cudnnReduceTensor has issue if input and output has same size, we just need to copy the data for this case - auto* const dst = Y->template MutableData(); - const auto* const src = X->template Data(); - if (input_count == output_count) { - if (src != dst) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst, src, input_count * sizeof(uint8_t), cudaMemcpyDeviceToDevice)); - } - return Status::OK(); - } - - // This reduction keep adding values to this buffer. If a non-zero value, say 1000, is here, the sum will start with 1000. - // Therefore zeroing out the memory is required - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes())); - - size_t indices_bytes = 0; - size_t workspace_bytes = 0; - CudnnTensor input_tensor; - CudnnTensor output_tensor; - CudnnReduceDescriptor reduce_desc; - - cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT; - IAllocatorUniquePtr temp_X = GetScratchBuffer(input_count); - Impl_Cast(reinterpret_cast(src), temp_X.get(), X->Shape().Size()); - - ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, CUDNN_REDUCE_TENSOR_FLATTENED_INDICES)); - ORT_RETURN_IF_ERROR(input_tensor.Set(input_dims_cudnn, cudnn_type_X)); - ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_cudnn, cudnn_type_X)); - CUDNN_RETURN_IF_ERROR(cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes)); - CUDNN_RETURN_IF_ERROR(cudnnGetReductionWorkspaceSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes)); - IAllocatorUniquePtr indices_cuda = GetScratchBuffer(indices_bytes); - IAllocatorUniquePtr workspace_cuda = GetScratchBuffer(workspace_bytes); - - const auto one = Consts::One; - const auto zero = Consts::Zero; - auto temp_Y = GetScratchBuffer(output_count); - CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(CudnnHandle(), - reduce_desc, - indices_cuda.get(), - indices_bytes, - workspace_cuda.get(), - workspace_bytes, - &one, - input_tensor, - temp_X.get(), - &zero, - output_tensor, - temp_Y.get())); - - Impl_Cast(temp_Y.get(), dst, output_count); + Impl_Cast(temp_Y.get(), reinterpret_cast(Y->template MutableData()), output_count); return Status::OK(); } +#endif namespace ReductionOps { @@ -972,8 +910,15 @@ template Tensor ReduceCompute( } // namespace ReductionOps +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define REGISTER_KERNEL_TYPED_BF16(name) REGISTER_KERNEL_TYPED(name, BFloat16) +#else +#define REGISTER_KERNEL_TYPED_BF16(name) +#endif + #define REGISTER_KERNEL_HFD(name) \ REGISTER_KERNEL_TYPED(name, MLFloat16) \ + REGISTER_KERNEL_TYPED_BF16(name) \ REGISTER_KERNEL_TYPED(name, float) \ REGISTER_KERNEL_TYPED(name, double) @@ -1009,6 +954,9 @@ REGISTER_KERNEL_TYPED_13(ReduceSum, MLFloat16) REGISTER_KERNEL_TYPED_13(ReduceSum, float) REGISTER_KERNEL_TYPED_13(ReduceSum, double) REGISTER_KERNEL_TYPED_13(ReduceSum, int32_t) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_KERNEL_TYPED_13(ReduceSum, BFloat16) +#endif REGISTER_KERNEL_HFD(ReduceLogSum) REGISTER_KERNEL_HFD(ReduceSumSquare) diff --git a/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h b/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h index bf2582426a..d0eafc9d23 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h +++ b/onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h @@ -17,6 +17,10 @@ template <> struct AccumulationType { using type = float; }; template <> struct AccumulationType { using type = double; }; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +template <> +struct AccumulationType { using type = float; }; +#endif template using AccumulationType_t = typename AccumulationType::type; diff --git a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h index 4bef68b343..90fdd2aea5 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h @@ -110,6 +110,37 @@ inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, #endif } +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, int n, int k, + const nv_bfloat16* alpha, + const nv_bfloat16* A, int lda, + const nv_bfloat16* B, int ldb, + const nv_bfloat16* beta, + nv_bfloat16* C, int ldc, + const cudaDeviceProp& prop) { + onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_DEFAULT_MATH); + + float h_a = onnxruntime::BFloat16(*reinterpret_cast(alpha)).ToFloat(); + float h_b = onnxruntime::BFloat16(*reinterpret_cast(beta)).ToFloat(); + + // accumulating in FP32 + return cublasGemmEx(handle, + transa, + transb, + m, n, k, + &h_a, + A, CUDA_R_16BF, lda, + B, CUDA_R_16BF, ldb, + &h_b, + C, CUDA_R_16BF, ldc, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT); +} +#endif + // batched gemm inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle, cublasOperation_t transa, @@ -209,6 +240,38 @@ inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle, #endif } +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, int n, int k, + const nv_bfloat16* alpha, + const nv_bfloat16* Aarray[], int lda, + const nv_bfloat16* Barray[], int ldb, + const nv_bfloat16* beta, + nv_bfloat16* Carray[], int ldc, + int batch_count, + const cudaDeviceProp& prop) { + onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TENSOR_OP_MATH); + float h_a = onnxruntime::BFloat16(*reinterpret_cast(alpha)).ToFloat(); + float h_b = onnxruntime::BFloat16(*reinterpret_cast(beta)).ToFloat(); + + // accumulating in FP32 + return cublasGemmBatchedEx(handle, + transa, + transb, + m, n, k, + &h_a, + (const void**)Aarray, CUDA_R_16BF, lda, + (const void**)Barray, CUDA_R_16BF, ldb, + &h_b, + (void**)Carray, CUDA_R_16BF, ldc, + batch_count, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); +} +#endif + // strided batched gemm inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, cublasOperation_t transa, @@ -319,6 +382,40 @@ inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, #endif } +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, int n, int k, + const nv_bfloat16* alpha, + const nv_bfloat16* A, int lda, + long long int strideA, + const nv_bfloat16* B, int ldb, + long long int strideB, + const nv_bfloat16* beta, + nv_bfloat16* C, int ldc, + long long int strideC, + int batch_count, + const cudaDeviceProp& prop) { + onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TENSOR_OP_MATH); + float h_a = onnxruntime::BFloat16(*reinterpret_cast(alpha)).ToFloat(); + float h_b = onnxruntime::BFloat16(*reinterpret_cast(beta)).ToFloat(); + // accumulating in FP32 + return cublasGemmStridedBatchedEx(handle, + transa, + transb, + m, n, k, + &h_a, + A, CUDA_R_16BF, lda, strideA, + B, CUDA_R_16BF, ldb, strideB, + &h_b, + C, CUDA_R_16BF, ldc, strideC, + batch_count, + CUDA_R_32F, + CUBLAS_GEMM_DEFAULT); +} +#endif + // transpose using geam inline cublasStatus_t cublasTransposeHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, const float* alpha, const float* A, int lda, const float* beta, const float* B, int ldb, float* C, int ldc) { return cublasSgeam(handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc); @@ -336,6 +433,9 @@ inline cublasStatus_t cublasCopyHelper(cublasHandle_t handle, int n, const doubl return cublasDcopy(handle, n, x, incx, y, incy); } cublasStatus_t cublasCopyHelper(cublasHandle_t handle, int n, const half* x, int incx, half* y, int incy); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +cublasStatus_t cublasCopyHelper(cublasHandle_t handle, int n, const nv_bfloat16* x, int incx, nv_bfloat16* y, int incy); +#endif diff --git a/onnxruntime/core/providers/cuda/tensor/cast_op.cc b/onnxruntime/core/providers/cuda/tensor/cast_op.cc index 2f0f71ce88..156597e3a0 100644 --- a/onnxruntime/core/providers/cuda/tensor/cast_op.cc +++ b/onnxruntime/core/providers/cuda/tensor/cast_op.cc @@ -11,6 +11,9 @@ namespace cuda { const std::vector castOpTypeConstraints{ DataTypeImpl::GetTensorType(), +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + DataTypeImpl::GetTensorType(), +#endif DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), @@ -76,6 +79,9 @@ Status Cast::ComputeInternal(OpKernelContext* context) const { switch (to_) { CASE(TensorProto_DataType_FLOAT16, MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + CASE(TensorProto_DataType_BFLOAT16, BFloat16) +#endif CASE(TensorProto_DataType_FLOAT, float) CASE(TensorProto_DataType_DOUBLE, double) CASE(TensorProto_DataType_INT8, int8_t) @@ -113,6 +119,9 @@ SPECIALIZE_IMPL(uint16_t) SPECIALIZE_IMPL(uint32_t) SPECIALIZE_IMPL(uint64_t) SPECIALIZE_IMPL(bool) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +SPECIALIZE_IMPL(BFloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc index 253c5d00a0..209fd57eca 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc +++ b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc @@ -109,6 +109,21 @@ Status GatherNDBase::PrepareCompute( .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ GatherND); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define GATHER_ND_T_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()} +#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, BFloat16 +#else +#define GATHER_ND_T_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()} +#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t +#endif + #define REGISTER_KERNEL_TYPED_GATHER_ND(TIndex, ver) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ GatherND, \ @@ -117,13 +132,7 @@ Status GatherNDBase::PrepareCompute( TIndex, \ kCudaExecutionProvider, \ KernelDefBuilder() \ - .TypeConstraint("T", \ - std::vector{ \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - }) \ + .TypeConstraint("T", GATHER_ND_T_TENSOR_TYPES) \ .TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \ GatherND); @@ -187,7 +196,7 @@ Status GatherND::ComputeInternal(OpKernelContext* context) const { const void* const kernel_input_data = input_tensor->DataRaw(); void* const kernel_output_data = output_tensor->MutableDataRaw(); - utils::MLTypeCallDispatcher + utils::MLTypeCallDispatcher t_disp(input_tensor->GetElementType()); t_disp.Invoke(num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get()); diff --git a/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu b/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu index dfb3a0aff9..64c89846d5 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu @@ -110,6 +110,9 @@ SPECIALIZED_IMPL(int64_t) SPECIALIZED_IMPL(half) SPECIALIZED_IMPL(double) #endif +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_IMPL(nv_bfloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 360eaafd33..ceff4c2cd6 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -747,7 +747,7 @@ void RegisterTrainingOpSchemas() { "Constrain input and output gradient types to float tensors.") .TypeConstraint( "T2", - OpSchema::all_tensor_types(), + OpSchema::all_tensor_types_with_bfloat(), "reset_signal can be of any tensor type.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateShapeAndTypeFromFirstInput(ctx); @@ -1135,11 +1135,11 @@ Example 4: .Output(1, "mask", "The output mask of dropout.", "T2", OpSchema::Optional) .TypeConstraint( "T", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeConstraint( "T1", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain input 'ratio' types to float tensors.") .TypeConstraint( "T2", @@ -1563,24 +1563,6 @@ Example 4: {{"X_1"}, "Cos", {"X"}}, {{"dX"}, "Mul", {"X_1", "dY"}}})); - ONNX_CONTRIB_OPERATOR_SCHEMA(ReshapeGrad) - .SetDomain(kOnnxDomain) - .SinceVersion(9) - .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL) - .SetDoc("Gradient function for Reshape") - .AllowUncheckedAttributes() - .Input(0, "X", "Input tensor", "T") - .Input(1, "dY", "Reshape output's grad", "T") - .Output(0, "dX", "REshape input's grad", "T") - .TypeConstraint( - "T", - {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, - "Constrain input and output types to all numeric tensors.") - .FunctionBody(ONNX_NAMESPACE::FunctionBodyHelper::BuildNodes( - {// nodes: {outputs, op, inputs, attributes} - {{"x_shape"}, "Shape", {"X"}}, - {{"dX"}, "Reshape", {"dY", "x_shape"}}})); - ONNX_CONTRIB_OPERATOR_SCHEMA(SummaryScalar) .SetDomain(kMSDomain) .SinceVersion(1) @@ -1893,7 +1875,7 @@ Return true if all elements are true and false otherwise. static_cast(0)) .TypeConstraint( "SrcT", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain input types to float tensors.") .TypeConstraint( "ScaleT", @@ -1901,7 +1883,7 @@ Return true if all elements are true and false otherwise. "Constrain scale types to float tensors.") .TypeConstraint( "DstT", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain output types to float tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { bool fuse_outputs = static_cast(getAttribute(ctx, "fuse_outputs", int64_t(0))); @@ -1978,11 +1960,11 @@ Return true if all elements are true and false otherwise. .Output(0, "Y", "output", "TOut") .TypeConstraint( "TIn", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain input types to float tensors.") .TypeConstraint( "TOut", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain scale types to float tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { updateOutputShape(ctx, 0, {}); @@ -2177,7 +2159,7 @@ Return true if all elements are true and false otherwise. .Output(0, "dX", "Gradient of the input.", "T") .TypeConstraint( "T", - {"tensor(float16)", "tensor(float)", "tensor(double)"}, + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index 110bfc9cb3..26a3ff8069 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/orttraining/orttraining/models/bert/main.cc @@ -110,6 +110,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet ("max_eval_steps", "Maximum number of eval steps.", cxxopts::value()->default_value("100")) ("seed", "Random seed.", cxxopts::value()->default_value("-1")) ("use_mixed_precision", "Whether to use a mix of fp32 and fp16 arithmetic on GPU.", cxxopts::value()->default_value("false")) + ("use_bfloat16", "Whether to use BFloat16 arithmetic on GPU.", cxxopts::value()->default_value("false")) ("enable_adasum", "Whether to use Adasum for allreduction.", cxxopts::value()->default_value("false")) ("allreduce_in_fp16", "Whether to do AllReduce in fp16. If false, AllReduce will be done in fp32", cxxopts::value()->default_value("true")) ("loss_scale", "Loss scaling, positive power of 2 values can improve fp16 convergence. " @@ -289,6 +290,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet } params.use_mixed_precision = flags["use_mixed_precision"].as(); + params.use_bfloat16 = flags["use_bfloat16"].as(); params.allreduce_in_mixed_precision_type = flags["allreduce_in_fp16"].as() && params.use_mixed_precision; if (params.use_mixed_precision) { printf("Mixed precision training is enabled.\n"); diff --git a/orttraining/orttraining/models/gpt2/main.cc b/orttraining/orttraining/models/gpt2/main.cc index 0dc7ef1fa2..09ee3bd7f7 100644 --- a/orttraining/orttraining/models/gpt2/main.cc +++ b/orttraining/orttraining/models/gpt2/main.cc @@ -64,6 +64,7 @@ Status ParseArguments(int argc, char* argv[], GPT2Parameters& params, OrtParamet cxxopts::value()->default_value("1")) ("seed", "Random seed.", cxxopts::value()->default_value("-1")) ("use_mixed_precision", "Whether to use a mix of fp32 and fp16 arithmetic on GPU.", cxxopts::value()->default_value("false")) + ("use_bfloat16", "Whether to use BFloat16 arithmetic on GPU.", cxxopts::value()->default_value("false")) ("enable_adasum", "Whether to use Adasum for allreduction.", cxxopts::value()->default_value("false")) ("allreduce_in_fp16", "Whether to do AllReduce in fp16. If false, AllReduce will be done in fp32", cxxopts::value()->default_value("true")) ("loss_scale", "Loss scaling, positive power of 2 values can improve fp16 convergence. " @@ -150,6 +151,7 @@ Status ParseArguments(int argc, char* argv[], GPT2Parameters& params, OrtParamet } params.use_mixed_precision = flags["use_mixed_precision"].as(); + params.use_bfloat16 = flags["use_bfloat16"].as(); params.allreduce_in_mixed_precision_type = flags["allreduce_in_fp16"].as() && params.use_mixed_precision; if (params.use_mixed_precision) { printf("Mixed precision training is enabled.\n"); diff --git a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc index a18de38cf8..e219f97951 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc @@ -10,13 +10,19 @@ namespace onnxruntime { namespace cuda { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double, BFloat16 +#else +#define ALL_IEEE_FLOAT_DATA_TYPES MLFloat16, float, double +#endif + ONNX_OPERATOR_KERNEL_EX( BiasGeluGrad_dX, kMSDomain, 1, kCudaExecutionProvider, KernelDefBuilder() - .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("T", BuildKernelDefConstraints()) .MayInplace(0, 0), BiasGeluGrad_dX); @@ -26,7 +32,7 @@ ONNX_OPERATOR_KERNEL_EX( 1, kCudaExecutionProvider, KernelDefBuilder() - .TypeConstraint("T", BuildKernelDefConstraints()) + .TypeConstraint("T", BuildKernelDefConstraints()) .MayInplace(0, 0), BiasGeluGrad_dX); @@ -70,7 +76,7 @@ Status BiasGeluGrad_dX::ComputeInternal(OpKernelContext* co utils::MLTypeCallDispatcher< KernelLaunchDispatcher, - MLFloat16, float, double> + ALL_IEEE_FLOAT_DATA_TYPES> dispatcher{X->GetElementType()}; dispatcher.Invoke(input_size, bias_size, *dY, *X, *B, *dX); diff --git a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu index da5d621f1d..f42a2360aa 100644 --- a/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu @@ -87,6 +87,11 @@ SPECIALIZED_BIAS_GELU_GRAD_IMPL(half, gelu_computation_mode::Approximation); SPECIALIZED_BIAS_GELU_GRAD_IMPL(float, gelu_computation_mode::Approximation); SPECIALIZED_BIAS_GELU_GRAD_IMPL(double, gelu_computation_mode::Approximation); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_BIAS_GELU_GRAD_IMPL(nv_bfloat16, gelu_computation_mode::Default); +SPECIALIZED_BIAS_GELU_GRAD_IMPL(nv_bfloat16, gelu_computation_mode::Approximation); +#endif + #undef SPECIALIZED_BIAS_GELU_GRAD_IMPL } // namespace cuda diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 16e9817334..eee66a0285 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/providers/cuda/cuda_fwd.h" +#include "core/providers/cuda/cuda_pch.h" #include "core/framework/kernel_registry.h" using namespace onnxruntime::common; @@ -119,6 +120,40 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Scale); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Scale); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +// Adam +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_float_float_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_float_float_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_BFloat16_float_float_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_BFloat16_BFloat16_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_BFloat16_float_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_BFloat16_BFloat16_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_BFloat16_float_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_BFloat16_BFloat16_BFloat16_BFloat16, AdamOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_BFloat16_BFloat16_float_BFloat16, AdamOptimizer); +// Lamb +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float_float_float_BFloat16, LambOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_BFloat16_float_BFloat16_BFloat16, LambOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_BFloat16_float_float_BFloat16, LambOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double_double_double_double_BFloat16, LambOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_BFloat16_BFloat16_BFloat16, LambOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_BFloat16_float_BFloat16, LambOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_float_BFloat16_BFloat16, LambOptimizer); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_float_float_BFloat16, LambOptimizer); + +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, InPlaceAccumulator); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, InPlaceAccumulator); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, InPlaceAccumulator); + +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, SoftmaxGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, MixedPrecisionScale); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, LayerNormalizationGrad); + +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, ReduceAllL2); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, ReduceAllL2); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, ReduceAllL2); +#endif + #if defined(ORT_USE_NCCL) || defined(USE_MPI) // P2P communication operators. class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Send); @@ -249,6 +284,40 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + // Adam + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + // Lamb + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, +#endif + // P2P communication operators. #if defined(ORT_USE_NCCL) || defined(USE_MPI) BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/math/isfinite.cuh b/orttraining/orttraining/training_ops/cuda/math/isfinite.cuh index 7de0fb260d..c2cb533451 100644 --- a/orttraining/orttraining/training_ops/cuda/math/isfinite.cuh +++ b/orttraining/orttraining/training_ops/cuda/math/isfinite.cuh @@ -22,5 +22,12 @@ __device__ __forceinline__ bool _IsFiniteScalar(const half value) { #endif } +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template<> +__device__ __forceinline__ bool _IsFiniteScalar(const nv_bfloat16 value) { + return isfinite(float(value)); +} +#endif + } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc index d78e5e75c4..88f800b96b 100644 --- a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc +++ b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc @@ -8,6 +8,15 @@ using namespace onnxruntime::common; namespace onnxruntime { namespace cuda { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()} +#else +#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() +#endif + #define REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(SrcT) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ MixedPrecisionScale, \ @@ -18,7 +27,7 @@ namespace cuda { KernelDefBuilder() \ .TypeConstraint("SrcT", DataTypeImpl::GetTensorType()) \ .TypeConstraint("ScaleT", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("DstT", DataTypeImpl::AllIEEEFloatTensorTypes()), \ + .TypeConstraint("DstT", ALL_IEEE_FLOAT_TENSOR_TYPES), \ MixedPrecisionScale); Status BytesPerElement(ONNX_NAMESPACE::TensorProto_DataType to, size_t& bytes_per_elem) { @@ -32,6 +41,11 @@ Status BytesPerElement(ONNX_NAMESPACE::TensorProto_DataType to, size_t& bytes_pe case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: bytes_per_elem = sizeof(MLFloat16); break; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: + bytes_per_elem = sizeof(BFloat16); + break; +#endif default: return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unexpected 'to' argument value: ", to); } @@ -102,6 +116,9 @@ Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) cons switch (to_) { CASE(TensorProto_DataType_FLOAT16, MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + CASE(TensorProto_DataType_BFLOAT16, BFloat16) +#endif CASE(TensorProto_DataType_FLOAT, float) default: return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unexpected 'to' argument value: ", to_); @@ -117,5 +134,10 @@ REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(float) template Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) const; template Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) const; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(BFloat16) +template Status MixedPrecisionScale::ComputeInternal(OpKernelContext* context) const; +#endif + } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu index 922d32cf50..b557e28a91 100644 --- a/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu +++ b/orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu @@ -42,5 +42,13 @@ SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(half, float) SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(float, half) SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(float, float) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(nv_bfloat16, nv_bfloat16) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(nv_bfloat16, float) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(float, nv_bfloat16) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(nv_bfloat16, half) +SPECIALIZE_MIXEDPRECISIONSCALE_IMPL(half, nv_bfloat16) +#endif + } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc b/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc index 6ed76e3ee0..270397c4d9 100644 --- a/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc @@ -60,6 +60,33 @@ Status SoftMaxGradComputeHelper( return Status::OK(); } +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +// cudnnSoftmaxForward/Backward doesn't support BFloat16. +#define SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(is_log_softmax) \ + template <> \ + Status SoftMaxGradComputeHelper( \ + const BFloat16* dY, \ + const TensorShape& input_shape, \ + const BFloat16* Y, \ + BFloat16* dX, \ + cudnnHandle_t, \ + int64_t axis) { \ + typedef typename ToCudaType::MappedType CudaT; \ + const int64_t normalized_axis = HandleNegativeAxis(axis, input_shape.NumDimensions()); \ + int64_t N = input_shape.SizeToDimension(normalized_axis); \ + int64_t D = input_shape.SizeFromDimension(normalized_axis); \ + auto dY_data = reinterpret_cast(dY); \ + auto Y_data = reinterpret_cast(Y); \ + auto dX_data = reinterpret_cast(dX); \ + dispatch_softmax_backward, is_log_softmax>( \ + dX_data, dY_data, Y_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); \ + return Status::OK(); \ + } + +SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(true) +SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(false) +#endif + #define REGISTER_GRADIENT_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ SoftmaxGrad, \ @@ -103,6 +130,9 @@ Status SoftmaxGrad::ComputeInternal(OpKernelContext* ctx) const { SPECIALIZED_GRADIENT(float) SPECIALIZED_GRADIENT(double) SPECIALIZED_GRADIENT(MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +SPECIALIZED_GRADIENT(BFloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu index 50eae2c11f..3b1bf2e508 100644 --- a/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu @@ -201,5 +201,9 @@ SPECIALIZED_SOFTMAX_GRAD_IMPL(float, float, float) SPECIALIZED_SOFTMAX_GRAD_IMPL(half, half, float) SPECIALIZED_SOFTMAX_GRAD_IMPL(double, double, double) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_SOFTMAX_GRAD_IMPL(nv_bfloat16, nv_bfloat16, float) +#endif + } } \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout.cc b/orttraining/orttraining/training_ops/cuda/nn/dropout.cc index 425e68827b..9d9d93acd4 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout.cc +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout.cc @@ -10,6 +10,17 @@ namespace onnxruntime { namespace cuda { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()} +#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double, BFloat16 +#else +#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() +#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double +#endif + // Temporary for backward compatibility, will eventually get rid of TrainableDropout when PyTorch exporter will move to // opset-12. ONNX_OPERATOR_KERNEL_EX( @@ -30,8 +41,8 @@ ONNX_OPERATOR_KERNEL_EX( 1, \ kCudaExecutionProvider, \ KernelDefBuilder() \ - .TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes()) \ - .TypeConstraint("T1", DataTypeImpl::AllIEEEFloatTensorTypes()) \ + .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) \ + .TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) \ .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ .InputMemoryType(2), \ DropoutGrad); @@ -70,13 +81,13 @@ Status DropoutGrad::ComputeInternal(OpKernelContext* context) const { float ratio_data = default_ratio_; auto ratio = context->Input(2); if (ratio) { - utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); + utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); t_disp.Invoke(ratio, ratio_data); } auto dX = context->Output(0, shape); - utils::MLTypeCallDispatcher t_disp(dY->GetElementType()); + utils::MLTypeCallDispatcher t_disp(dY->GetElementType()); t_disp.Invoke(N, *dY, mask_data, ratio_data, *dX); return Status::OK(); @@ -88,8 +99,8 @@ ONNX_OPERATOR_KERNEL_EX( 1, kCudaExecutionProvider, KernelDefBuilder() - .TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes()) - .TypeConstraint("T1", DataTypeImpl::AllIEEEFloatTensorTypes()) + .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) + .TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) .TypeConstraint("T2", DataTypeImpl::GetTensorType()) .InputMemoryType(3) .InputMemoryType(4), @@ -161,7 +172,7 @@ Status BiasDropout::ComputeInternal(OpKernelContext* context) const { float ratio_data = default_ratio_; auto ratio = context->Input(3); if (ratio) { - utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); + utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); t_disp.Invoke(ratio, ratio_data); } @@ -182,7 +193,7 @@ Status BiasDropout::ComputeInternal(OpKernelContext* context) const { const fast_divmod fdm_dim(gsl::narrow_cast(dim)); PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default(); - utils::MLTypeCallDispatcherRet t_disp(X->GetElementType()); + utils::MLTypeCallDispatcherRet t_disp(X->GetElementType()); return t_disp.Invoke(GetDeviceProp(), N, fdm_dim, ratio_data, generator, *X, *bias, residual, *Y, mask_data); } diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu index 4bf303d67e..b7960c5151 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu @@ -71,6 +71,9 @@ void DropoutGradientKernelImpl( SPECIALIZED_DROPOUT_GRAD_IMPL(float) SPECIALIZED_DROPOUT_GRAD_IMPL(double) SPECIALIZED_DROPOUT_GRAD_IMPL(half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_DROPOUT_GRAD_IMPL(nv_bfloat16) +#endif constexpr int UNROLL = 4; @@ -168,7 +171,9 @@ void BiasDropoutKernelImpl( SPECIALIZED_BIAS_DROPOUT_IMPL(float) SPECIALIZED_BIAS_DROPOUT_IMPL(double) SPECIALIZED_BIAS_DROPOUT_IMPL(half) - +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_BIAS_DROPOUT_IMPL(nv_bfloat16) +#endif } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc b/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc index b6e9677c56..3695c029c8 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc +++ b/orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc @@ -44,6 +44,9 @@ namespace cuda { REGISTER_GRADIENT_KERNEL_TYPED(float, float) REGISTER_GRADIENT_KERNEL_TYPED(double, double) REGISTER_GRADIENT_KERNEL_TYPED(MLFloat16, float) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_GRADIENT_KERNEL_TYPED(BFloat16, float) +#endif template LayerNormGrad::LayerNormGrad(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) { diff --git a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu index 23f206bb39..cb9be8d6fc 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu @@ -540,6 +540,10 @@ LAYERNORMGRAD_IMPL(half, float, true) LAYERNORMGRAD_IMPL(float, float, false) LAYERNORMGRAD_IMPL(double, double, false) LAYERNORMGRAD_IMPL(half, float, false) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +LAYERNORMGRAD_IMPL(nv_bfloat16, float, true) +LAYERNORMGRAD_IMPL(nv_bfloat16, float, false) +#endif } // namespace cuda -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc index f178522efe..c169955523 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc @@ -47,6 +47,18 @@ REGISTER_ADAM_KERNEL_TYPED(MLFloat16, int64_t, float, MLFloat16, MLFloat16, floa REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, MLFloat16, MLFloat16, MLFloat16, MLFloat16) REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, MLFloat16, MLFloat16, float, MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, float, float, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(BFloat16, int64_t, float, BFloat16, float, float, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, BFloat16, float, float, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, BFloat16, BFloat16, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, BFloat16, float, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(BFloat16, int64_t, float, BFloat16, BFloat16, BFloat16, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(BFloat16, int64_t, float, BFloat16, BFloat16, float, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, BFloat16, BFloat16, BFloat16, BFloat16) +REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, BFloat16, BFloat16, float, BFloat16) +#endif + template Status AdamOptimizer::ComputeInternal(OpKernelContext* ctx) const { typedef typename ToCudaType::MappedType CudaT1; diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu index bde32e2296..bf26206e96 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu @@ -255,5 +255,17 @@ SPECIALIZED_AdamOptimizerImpl(half, int64_t, float, half, half, float, half) SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, half, half, half, half) SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, half, half, float, half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, float, float, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(nv_bfloat16, int64_t, float, nv_bfloat16, float, float, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, nv_bfloat16, float, float, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, nv_bfloat16, nv_bfloat16, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, float, nv_bfloat16, float, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(nv_bfloat16, int64_t, float, nv_bfloat16, nv_bfloat16, nv_bfloat16, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(nv_bfloat16, int64_t, float, nv_bfloat16, nv_bfloat16, float, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, nv_bfloat16, nv_bfloat16, nv_bfloat16, nv_bfloat16) +SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, nv_bfloat16, nv_bfloat16, float, nv_bfloat16) +#endif + } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc index c8883b26ee..569e0210c9 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc @@ -28,6 +28,11 @@ REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(float, float) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(float, MLFloat16) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(MLFloat16, MLFloat16) REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(MLFloat16, float) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(float, BFloat16) +REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(BFloat16, BFloat16) +REGISTER_IN_PLACE_TENSOR_ACCUMULATOR_TYPED(BFloat16, float) +#endif template Status ZeroGradient::ComputeInternal(OpKernelContext* ctx) const { diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu index 6075ab9d51..b6c49a5acb 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu @@ -45,6 +45,11 @@ SPECIALIZED_IMPL_InPlaceAccumulator(float, float) SPECIALIZED_IMPL_InPlaceAccumulator(float, half) SPECIALIZED_IMPL_InPlaceAccumulator(half, half) SPECIALIZED_IMPL_InPlaceAccumulator(half, float) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_IMPL_InPlaceAccumulator(float, nv_bfloat16) +SPECIALIZED_IMPL_InPlaceAccumulator(nv_bfloat16, nv_bfloat16) +SPECIALIZED_IMPL_InPlaceAccumulator(nv_bfloat16, float) +#endif } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc index 945b31245f..0c37031fa9 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc @@ -75,6 +75,17 @@ REGISTER_LAMB_KERNEL_TYPED(MLFloat16, float, MLFloat16, MLFloat16, float, MLFloa REGISTER_LAMB_KERNEL_TYPED(MLFloat16, float, MLFloat16, float, MLFloat16, MLFloat16) REGISTER_LAMB_KERNEL_TYPED(MLFloat16, float, MLFloat16, float, float, MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_LAMB_KERNEL_TYPED(float, float, BFloat16, float, BFloat16, BFloat16) +REGISTER_LAMB_KERNEL_TYPED(float, float, BFloat16, float, float, BFloat16) +REGISTER_LAMB_KERNEL_TYPED(float, float, float, float, float, BFloat16) +REGISTER_LAMB_KERNEL_TYPED(double, double, double, double, double, BFloat16) +REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, BFloat16, BFloat16, BFloat16) +REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, BFloat16, float, BFloat16) +REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, float, BFloat16, BFloat16) +REGISTER_LAMB_KERNEL_TYPED(BFloat16, float, BFloat16, float, float, BFloat16) +#endif + void check_inputs_and_outputs( const Tensor* w, const Tensor* g, diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu index c67ee9c2bb..0e89da0129 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu @@ -173,6 +173,13 @@ SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, half, half, float) SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, half, float, half) SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, half, float, float) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, nv_bfloat16, nv_bfloat16) +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, nv_bfloat16, float) +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, float, nv_bfloat16) +SPECIALIZED_LAMB_COMPUTE_DIRECTION(float, nv_bfloat16, float, float) +#endif + template __device__ __forceinline__ void _LambUpdateRule( const T1 eta, @@ -292,6 +299,13 @@ INSTANTIATE_LAMB_UPDATE(double, double, double, half) INSTANTIATE_LAMB_UPDATE(half, float, half, half) INSTANTIATE_LAMB_UPDATE(float, float, half, half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_LAMB_UPDATE(float, float, float, nv_bfloat16) +INSTANTIATE_LAMB_UPDATE(double, double, double, nv_bfloat16) +INSTANTIATE_LAMB_UPDATE(nv_bfloat16, float, nv_bfloat16, nv_bfloat16) +INSTANTIATE_LAMB_UPDATE(float, float, nv_bfloat16, nv_bfloat16) +#endif + template __global__ void LambMultiTensorComputeDirectionImpl( ChunkGroup<6> chunk_group, @@ -380,6 +394,13 @@ INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, half, half, float) INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, half, float, half) INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, half, float, float) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, nv_bfloat16, nv_bfloat16) +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, nv_bfloat16, float) +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, float, nv_bfloat16) +INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, nv_bfloat16, float, float) +#endif + template __global__ void LambMultiTensorUpdateImpl( ChunkGroup<7> chunk_group, @@ -442,6 +463,13 @@ INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(double, double, double, half) INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(half, float, half, half) INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, half, half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, float, nv_bfloat16) +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(double, double, double, nv_bfloat16) +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(nv_bfloat16, float, nv_bfloat16, nv_bfloat16) +INSTANTIATE_LAMB_MULTI_TENSOR_UPDATE_FUNCTOR(float, float, nv_bfloat16, nv_bfloat16) +#endif + // w_buffer[i], d_buffer[i] is used to store the squared sum of all elements processed by the i-th block. // sync_range_and_lock is used for a well ordered reduction over blocks spanning the same tensor template @@ -613,5 +641,11 @@ INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, half, float, half, float) INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, half, float, float, float) INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(half, half, half, half, float) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, nv_bfloat16, float, nv_bfloat16, float) +INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, nv_bfloat16, float, float, float) +INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(nv_bfloat16, nv_bfloat16, nv_bfloat16, nv_bfloat16, float) +#endif + } // namespace cuda } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc index 939d861229..5d90e9936b 100644 --- a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc +++ b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc @@ -104,6 +104,11 @@ REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, float, float) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, MLFloat16, float) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, float, MLFloat16) REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, MLFloat16, MLFloat16) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, BFloat16, float) +REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, float, BFloat16) +REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, BFloat16, BFloat16) +#endif } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu index b35669d7d2..adde87d307 100644 --- a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu +++ b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu @@ -24,6 +24,10 @@ void ScalarSqrt(Tin* input, Tout* output) { template void ScalarSqrt(float* input, float* output); template void ScalarSqrt(half* input, half* output); template void ScalarSqrt(float* input, half* output); +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template void ScalarSqrt(nv_bfloat16* input, nv_bfloat16* output); +template void ScalarSqrt(float* input, nv_bfloat16* output); +#endif template __launch_bounds__(ChunkGroup<1>::thread_count_per_block) @@ -110,6 +114,11 @@ INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, float) INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, float) INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, half) INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(nv_bfloat16, float) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, nv_bfloat16) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(nv_bfloat16, nv_bfloat16) +#endif } // namespace cuda } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc index 947504e496..41607cb88d 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc @@ -10,6 +10,15 @@ namespace onnxruntime { namespace cuda { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()} +#else +#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() +#endif + ONNX_OPERATOR_KERNEL_EX( GatherGrad, kMSDomain, @@ -18,8 +27,7 @@ ONNX_OPERATOR_KERNEL_EX( KernelDefBuilder() .InputMemoryType(0) .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", {DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType()}) + .TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) .TypeConstraint("Tind", std::vector{ DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), @@ -83,6 +91,11 @@ Status DispatchToGatherGradImpl( } else if (utils::IsPrimitiveDataType(t_data_type)) { return DispatchToGatherGradImplByTindex( tindex_data_type, allocator, num_gathered_per_index, gather_dimension_size, num_batches, dY, gathered_indices, dX); +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 + } else if (utils::IsPrimitiveDataType(t_data_type)) { + return DispatchToGatherGradImplByTindex( + tindex_data_type, allocator, num_gathered_per_index, gather_dimension_size, num_batches, dY, gathered_indices, dX); +#endif } return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "GatherGrad unsupported T type: ", t_data_type); diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu index bc4a49dbc4..9ea6196973 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu @@ -543,6 +543,9 @@ void GatherGradImpl( SPECIALIZED_WITH_IDX(float) SPECIALIZED_WITH_IDX(half) +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_WITH_IDX(nv_bfloat16) +#endif #undef SPECIALIZED_WITH_IDX #undef SPECIALIZED diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc index 8c9fb54fdf..b7dad4963b 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc @@ -8,6 +8,17 @@ namespace onnxruntime { namespace cuda { +#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 +#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()} +#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double, BFloat16 +#else +#define ALL_IEEE_FLOAT_TENSOR_TYPES DataTypeImpl::AllIEEEFloatTensorTypes() +#define ALL_IEEE_FLOAT_DATA_TYPES float, MLFloat16, double +#endif + #define REGISTER_KERNEL_TYPED_GATHER_ND_GRAD(TIndex) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ GatherNDGrad, \ @@ -15,8 +26,7 @@ namespace cuda { 1, \ TIndex, \ kCudaExecutionProvider, \ - KernelDefBuilder().TypeConstraint("T", {DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}) \ + KernelDefBuilder().TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) \ .TypeConstraint("Tind", DataTypeImpl::GetTensorType()) \ .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ .InputMemoryType(0), \ @@ -83,7 +93,7 @@ Status GatherNDGrad::ComputeInternal(OpKernelContext* context) const { const void* const kernel_input_data = update_tensor->DataRaw(); void* const kernel_output_data = output_tensor->MutableDataRaw(); - utils::MLTypeCallDispatcher + utils::MLTypeCallDispatcher t_disp(update_tensor->GetElementType()); t_disp.Invoke(num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get()); diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu index 8abad7e30e..46082ed41c 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu @@ -42,6 +42,9 @@ SPECIALIZED_GRAD_IMPL(float); SPECIALIZED_GRAD_IMPL(half); SPECIALIZED_GRAD_IMPL(double); #endif +#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +SPECIALIZED_GRAD_IMPL(nv_bfloat16); +#endif } // namespace cuda } // namespace onnxruntime