saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Train BERT Using BFloat16 on A100 (#6090)

Browse source 
* 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 <weicwang@microsoft.com>
This commit is contained in:
Vincent Wang 2021-01-14 19:04:32 +08:00 committed by GitHub
parent e35db194e3
commit 4df356d1c9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
64 changed files with 965 additions and 263 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc
View file

@ -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;

40
onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu
View file

@ -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 <unsigned TPB>
__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<const nv_bfloat162*>(input);
const nv_bfloat162* bias2 = reinterpret_cast<const nv_bfloat162*>(bias);
nv_bfloat162* output2 = reinterpret_cast<nv_bfloat162*>(output);
FastGeluKernel2<blockSize><<<gridSize, blockSize, 0, stream>>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2);
} else {
const int gridSize = (input_length + blockSize - 1) / blockSize;
FastGeluKernel<nv_bfloat16, blockSize><<<gridSize, blockSize, 0, stream>>>(A, B, C, input_length, bias_length, input, bias, output);
}
return CUDA_CALL(cudaPeekAtLastError());
}
#endif
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

17
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
View file

@ -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<void>() {
KernelCreateInfo info;
@ -159,7 +166,15 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int8_t_MLFloat16, DequantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, uint8_t_MLFloat16, DequantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int8_t, QAttention)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int8_t, QAttention)>};
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int8_t, QAttention)>,
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FastGelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, TransposeMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FusedMatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, BFloat16_float, LayerNormalization)>,
#endif
};
for (auto& function_table_entry : function_table) {
KernelCreateInfo info = function_table_entry();

3
onnxruntime/contrib_ops/cuda/layer_norm.cc
View file

@ -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 <typename T, typename U, bool simplified>
LayerNorm<T, U, simplified>::LayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {

4
onnxruntime/contrib_ops/cuda/layer_norm_impl.cu
View file

@ -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

5
onnxruntime/contrib_ops/cuda/math/fused_matmul.cc
View file

@ -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

2
onnxruntime/core/optimizer/dropout_elimination.cc
View file

@ -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;
}

7
onnxruntime/core/providers/cuda/activation/activations.cc
View file

@ -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)

13
onnxruntime/core/providers/cuda/activation/activations_impl.cu
View file

@ -93,9 +93,16 @@ struct OP_ThresholdedRelu : public CtxThresholdedRelu {
#define SPECIALIZED_UNARY_ACTIVATION_IMPL(name, T) \
template void Impl_##name<T>(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) \

21
onnxruntime/core/providers/cuda/atomic/common.cuh
View file

@ -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<unsigned int*>(reinterpret_cast<char*>(address) - (reinterpret_cast<size_t>(address) & 2));
unsigned int old = *base_address;
unsigned int assumed;
unsigned short x;
do {
assumed = old;
x = reinterpret_cast<size_t>(address) & 2 ? (old >> 16) : (old & 0xffff);
x = __bfloat16_as_short(__float2bfloat16(__bfloat162float(*reinterpret_cast<const __nv_bfloat16*>(&x)) + __bfloat162float(value)));
old = reinterpret_cast<size_t>(address) & 2 ? (old & 0xffff) | (x << 16) : (old & 0xffff0000) | x;
old = atomicCAS(base_address, assumed, old);
} while (assumed != old);
}
#endif
} // namespace cuda
} // namespace onnxruntime

28
onnxruntime/core/providers/cuda/cu_inc/common.cuh
View file

@ -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 <cuda_bf16.h>
#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<float>(a))); }
template <>
__device__ __inline__ nv_bfloat16 _Exp(nv_bfloat16 a) { return nv_bfloat16(expf(static_cast<float>(a))); }
template <>
__device__ __inline__ nv_bfloat16 _Log(nv_bfloat16 a) { return nv_bfloat16(logf(static_cast<float>(a))); }
template <>
__device__ __inline__ nv_bfloat16 _Tanh(nv_bfloat16 a) { return nv_bfloat16(tanhf(static_cast<float>(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<float>(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.

12
onnxruntime/core/providers/cuda/cuda_common.h
View file

@ -67,6 +67,18 @@ class ToCudaType<MLFloat16> {
}
};
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
template <>
class ToCudaType<BFloat16> {
public:
typedef nv_bfloat16 MappedType;
static MappedType FromFloat(float f) {
uint16_t h = BFloat16(f).val;
return *reinterpret_cast<MappedType*>(&h);
}
};
#endif
inline bool CalculateFdmStrides(gsl::span<fast_divmod> p, const std::vector<int64_t>& dims) {
int stride = 1;
if (dims.empty() || p.size() < dims.size())

31
onnxruntime/core/providers/cuda/cuda_execution_provider.cc
View file

@ -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<void>() {
KernelCreateInfo info;
@ -1702,6 +1717,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, ReduceSumSquare)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int64_t, GatherND)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, Dropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, Resize)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, double, Resize)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, Resize)>,
@ -1715,6 +1731,21 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, LRN)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, Identity)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, ScatterND)>,
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Add)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Sub)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Mul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Div)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Cast)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Softmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, MatMul)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Relu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Sigmoid)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum)>,
#endif
};
for (auto& function_table_entry : function_table) {

10
onnxruntime/core/providers/cuda/cuda_execution_provider.h
View file

@ -122,6 +122,13 @@ class CUDAExecutionProvider : public IExecutionProvider {
constant_ones_half_ = cuda::CreateConstantOnes<half>();
}
return reinterpret_cast<const T*>(constant_ones_half_->GetBuffer(count));
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
} else if (std::is_same<T, nv_bfloat16>::value) {
if (!constant_ones_bfloat16_) {
constant_ones_bfloat16_ = cuda::CreateConstantOnes<nv_bfloat16>();
}
return reinterpret_cast<const T*>(constant_ones_bfloat16_->GetBuffer(count));
#endif
} else {
return nullptr;
}
@ -143,6 +150,9 @@ class CUDAExecutionProvider : public IExecutionProvider {
std::unique_ptr<cuda::IConstantBuffer<float>> constant_ones_float_;
std::unique_ptr<cuda::IConstantBuffer<double>> constant_ones_double_;
std::unique_ptr<cuda::IConstantBuffer<half>> constant_ones_half_;
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
std::unique_ptr<cuda::IConstantBuffer<nv_bfloat16>> constant_ones_bfloat16_;
#endif
AllocatorPtr allocator_;
};

6
onnxruntime/core/providers/cuda/cuda_utils.cu
View file

@ -71,6 +71,9 @@ std::unique_ptr<IConstantBuffer<T>> CreateConstantOnes() {
template std::unique_ptr<IConstantBuffer<float>> CreateConstantOnes<float>();
template std::unique_ptr<IConstantBuffer<double>> CreateConstantOnes<double>();
template std::unique_ptr<IConstantBuffer<half>> CreateConstantOnes<half>();
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
template std::unique_ptr<IConstantBuffer<nv_bfloat16>> CreateConstantOnes<nv_bfloat16>();
#endif
#define SPECIALIZED_FILL(T) \
template void Fill<T>(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

5
onnxruntime/core/providers/cuda/cudnn_common.cc
View file

@ -161,6 +161,11 @@ const float Consts<half>::Zero = 0;
const float Consts<half>::One = 1;
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
const float Consts<nv_bfloat16>::Zero = 0;
const float Consts<nv_bfloat16>::One = 1;
#endif
template <>
const int8_t Consts<int8_t>::Zero = 0;

8
onnxruntime/core/providers/cuda/cudnn_common.h
View file

@ -126,6 +126,14 @@ struct Consts<half> {
static const float One;
};
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
template<>
struct Consts<nv_bfloat16> {
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)

19
onnxruntime/core/providers/cuda/fpgeneric.cu
View file

@ -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<<<dimGrid, dimBlock>>>(x, incx, y, incy, n);
return CUBLAS_STATUS_SUCCESS;
}
}
#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<<<dimGrid, dimBlock>>>(x, incx, y, incy, n);
return CUBLAS_STATUS_SUCCESS;
}
#endif

25
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc
View file

@ -188,9 +188,23 @@ Status BinaryElementwise<ShouldBroadcast>::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<ShouldBroadcast>::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<ShouldBroadcast>::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<ShouldBroadcast>::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)

11
onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu
View file

@ -73,12 +73,21 @@ namespace cuda {
const TArray<int64_t>* rhs_padded_strides, const T2* rhs_data, \
const TArray<fast_divmod>* 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)

3
onnxruntime/core/providers/cuda/math/gemm.cc
View file

@ -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 <typename T>
Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {

3
onnxruntime/core/providers/cuda/math/matmul.cc
View file

@ -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]

27
onnxruntime/core/providers/cuda/math/softmax.cc
View file

@ -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<BFloat16, is_log_softmax>( \
const BFloat16* X, \
const TensorShape& input_shape, \
BFloat16* Y, \
cudnnHandle_t, \
int64_t axis) { \
typedef typename ToCudaType<BFloat16>::MappedType CudaT; \
int64_t N = input_shape.SizeToDimension(axis); \
int64_t D = input_shape.SizeFromDimension(axis); \
auto Y_data = reinterpret_cast<CudaT*>(Y); \
auto X_data = reinterpret_cast<const CudaT*>(X); \
dispatch_softmax_forward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>( \
Y_data, X_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(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<T>::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

3
onnxruntime/core/providers/cuda/math/softmax_impl.cu
View file

@ -214,6 +214,9 @@ template void dispatch_softmax_forward<input_t, output_t, acc_t, true>(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
}
}

23
onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu
View file

@ -91,11 +91,24 @@ struct ViaTypeMap<half> {
typedef float ViaT;
};
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
template <>
struct ViaTypeMap<nv_bfloat16> {
typedef float ViaT;
};
#endif
template <typename InT, typename OutT>
struct OP_Cast {
__device__ __inline__ OutT operator()(const InT& a) const {
const bool any_float16 = std::is_same<half, InT>::value || std::is_same<half, OutT>::value;
#if CUDA_VERSION >= 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
const bool any_bf16 = std::is_same<nv_bfloat16, InT>::value || std::is_same<nv_bfloat16, OutT>::value;
typedef typename std::conditional<any_bf16, nv_bfloat16, OutT>::type T1;
typedef typename std::conditional<any_float16, half, T1>::type T;
#else
typedef typename std::conditional<any_float16, half, OutT>::type T;
#endif
typedef typename ViaTypeMap<T>::ViaT ViaT;
return (OutT)((ViaT)a);
}
@ -115,8 +128,15 @@ void Impl_Cast(
#define SPECIALIZED_CAST_IMPL2(InT, OutT) \
template void Impl_Cast<InT, OutT>(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

16
onnxruntime/core/providers/cuda/math/variadic_elementwise_ops.cc
View file

@ -194,22 +194,28 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
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<uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double>();
BuildKernelDefConstraints<uint32_t, uint64_t, int32_t, int64_t, ALL_IEEE_FLOAT_DATA_TYPES>();
const auto k_hfd_datatypes =
BuildKernelDefConstraints<MLFloat16, float, double>();
BuildKernelDefConstraints<ALL_IEEE_FLOAT_DATA_TYPES>();
} // namespace

7
onnxruntime/core/providers/cuda/math/variadic_elementwise_ops_impl.cu
View file

@ -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)

13
onnxruntime/core/providers/cuda/nn/dropout.cc
View file

@ -6,6 +6,15 @@
namespace onnxruntime {
namespace cuda {
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>()}
#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<bool>())
.InputMemoryType<OrtMemTypeCPUInput>(1)
.InputMemoryType<OrtMemTypeCPUInput>(2),

4
onnxruntime/core/providers/cuda/nn/dropout.h
View file

@ -105,7 +105,11 @@ Status Dropout<trainable_dropout>::ComputeInternal(OpKernelContext* context) con
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
utils::MLTypeCallDispatcher<DropoutComputeImpl, float, MLFloat16, double, BFloat16> t_disp(X->GetElementType());
#else
utils::MLTypeCallDispatcher<DropoutComputeImpl, float, MLFloat16, double> t_disp(X->GetElementType());
#endif
t_disp.Invoke(GetDeviceProp(), N, ratio_data, generator, *X, *Y, mask_data);
return Status::OK();

3
onnxruntime/core/providers/cuda/nn/dropout_impl.cu
View file

@ -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

9
onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
View file

@ -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 <typename TIn, typename TOut>
@ -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

334
onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
View file

@ -667,17 +667,104 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
calculate_log_, calculate_sqt_, log_sum_exp_, fast_reduction);
}
template <>
template <>
Status ReduceKernel<true>::ComputeImpl<int32_t, CUDNN_REDUCE_TENSOR_NO_INDICES>(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const {
typedef typename ToCudaType<int32_t>::MappedType CudaT;
#define SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(T) \
template <> \
template <> \
Status ReduceKernel<true>::ComputeImpl<T, CUDNN_REDUCE_TENSOR_NO_INDICES>( \
OpKernelContext * ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { \
typedef typename ToCudaType<T>::MappedType CudaT; \
const Tensor* X = ctx->Input<Tensor>(0); \
std::vector<int64_t> axes; \
size_t num_inputs = ctx->InputCount(); \
if (num_inputs == 2) { \
const Tensor* axes_tensor = ctx->Input<Tensor>(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<size_t>(axes_tensor->Shape()[0]); \
const auto* data = axes_tensor->template Data<int64_t>(); \
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<T>(), X->template Data<T>(), 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<int64_t>& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn; \
std::vector<int64_t>& 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<T>() != X->template Data<T>()) { \
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData<T>(), X->template Data<T>(), \
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<float> temp_X = GetScratchBuffer<float>(input_count); \
Impl_Cast<CudaT, float>(reinterpret_cast<const CudaT*>(X->template Data<T>()), 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<uint32_t> indices_cuda = GetScratchBuffer<uint32_t>(indices_bytes); \
IAllocatorUniquePtr<CudaT> workspace_cuda = GetScratchBuffer<CudaT>(workspace_bytes); \
\
const auto one = Consts<float>::One; \
const auto zero = Consts<float>::Zero; \
auto temp_Y = GetScratchBuffer<float>(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<float, CudaT>(temp_Y.get(), reinterpret_cast<CudaT*>(Y->template MutableData<T>()), 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<true>::ComputeImpl<BFloat16, CUDNN_REDUCE_TENSOR_NO_INDICES>(
OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const {
typedef typename ToCudaType<BFloat16>::MappedType CudaT;
const Tensor* X = ctx->Input<Tensor>(0);
std::vector<int64_t> 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<Tensor>(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<true>::ComputeImpl<int32_t, CUDNN_REDUCE_TENSOR_NO_INDICES>(
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<int32_t>(), X->template Data<int32_t>(), X->SizeInBytes(), cudaMemcpyDeviceToDevice));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData<BFloat16>(), X->template Data<BFloat16>(),
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<true>::ComputeImpl<int32_t, CUDNN_REDUCE_TENSOR_NO_INDICES>(
std::vector<int64_t>& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn;
std::vector<int64_t>& 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<int32_t>() != X->template Data<int32_t>()) {
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData<int32_t>(), X->template Data<int32_t>(), input_count * sizeof(int32_t), cudaMemcpyDeviceToDevice));
if (Y->template MutableData<BFloat16>() != X->template Data<BFloat16>()) {
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData<BFloat16>(), X->template Data<BFloat16>(),
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<const CudaT*>(X->template Data<BFloat16>()),
reinterpret_cast<CudaT*>(Y->template MutableData<BFloat16>()), m, n);
}
case ApplicableMatrixReduction::Columns: {
const auto buffer_size_bytes = compute_reduce_matrix_columns_buffer_size<CudaT>(m, n);
auto buffer = cuda_ep_->GetScratchBuffer<void>(buffer_size_bytes);
return reduce_matrix_columns(reinterpret_cast<const CudaT*>(X->template Data<BFloat16>()),
reinterpret_cast<CudaT*>(Y->template MutableData<BFloat16>()), 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<true>::ComputeImpl<int32_t, CUDNN_REDUCE_TENSOR_NO_INDICES>(
cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT;
IAllocatorUniquePtr<float> temp_X = GetScratchBuffer<float>(input_count);
Impl_Cast<CudaT, float>(reinterpret_cast<const CudaT*>(X->template Data<int32_t>()), temp_X.get(), X->Shape().Size());
Impl_Cast<CudaT, float>(reinterpret_cast<const CudaT*>(X->template Data<BFloat16>()), 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<uint32_t> indices_cuda = GetScratchBuffer<uint32_t>(indices_bytes);
IAllocatorUniquePtr<CudaT> workspace_cuda = GetScratchBuffer<CudaT>(workspace_bytes);
const auto one = Consts<float>::One;
const auto zero = Consts<float>::Zero;
auto temp_Y = GetScratchBuffer<float>(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<float, int32_t>(temp_Y.get(), Y->template MutableData<int32_t>(), output_count);
return Status::OK();
}
template <>
template <>
Status ReduceKernel<true>::ComputeImpl<int8_t, CUDNN_REDUCE_TENSOR_NO_INDICES>(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const {
typedef typename ToCudaType<int8_t>::MappedType CudaT;
const Tensor* X = ctx->Input<Tensor>(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<int64_t>& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn;
std::vector<int64_t>& 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<int8_t>();
const auto* const src = X->template Data<int8_t>();
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<float> temp_X = GetScratchBuffer<float>(input_count);
Impl_Cast<CudaT, float>(reinterpret_cast<const CudaT*>(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<uint32_t> indices_cuda = GetScratchBuffer<uint32_t>(indices_bytes);
IAllocatorUniquePtr<CudaT> workspace_cuda = GetScratchBuffer<CudaT>(workspace_bytes);
const auto one = Consts<float>::One;
const auto zero = Consts<float>::Zero;
auto temp_Y = GetScratchBuffer<float>(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<float, int8_t>(temp_Y.get(), dst, output_count);
return Status::OK();
}
template <>
template <>
Status ReduceKernel<true>::ComputeImpl<uint8_t, CUDNN_REDUCE_TENSOR_NO_INDICES>(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const {
typedef typename ToCudaType<uint8_t>::MappedType CudaT;
const Tensor* X = ctx->Input<Tensor>(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<int64_t>& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn;
std::vector<int64_t>& 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<uint8_t>();
const auto* const src = X->template Data<uint8_t>();
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<float> temp_X = GetScratchBuffer<float>(input_count);
Impl_Cast<CudaT, float>(reinterpret_cast<const CudaT*>(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<uint32_t> indices_cuda = GetScratchBuffer<uint32_t>(indices_bytes);
IAllocatorUniquePtr<CudaT> workspace_cuda = GetScratchBuffer<CudaT>(workspace_bytes);
const auto one = Consts<float>::One;
const auto zero = Consts<float>::Zero;
auto temp_Y = GetScratchBuffer<float>(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<float, uint8_t>(temp_Y.get(), dst, output_count);
Impl_Cast<float, CudaT>(temp_Y.get(), reinterpret_cast<CudaT*>(Y->template MutableData<BFloat16>()), output_count);
return Status::OK();
}
#endif
namespace ReductionOps {
@ -972,8 +910,15 @@ template Tensor ReduceCompute<double, CUDNN_REDUCE_TENSOR_NO_INDICES>(
} // 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)

4
onnxruntime/core/providers/cuda/shared_inc/accumulation_type.h
View file

@ -17,6 +17,10 @@ template <>
struct AccumulationType<float> { using type = float; };
template <>
struct AccumulationType<double> { using type = double; };
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
template <>
struct AccumulationType<nv_bfloat16> { using type = float; };
#endif
template <typename T>
using AccumulationType_t = typename AccumulationType<T>::type;

100
onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h
View file

@ -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<const uint16_t*>(alpha)).ToFloat();
float h_b = onnxruntime::BFloat16(*reinterpret_cast<const uint16_t*>(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<const uint16_t*>(alpha)).ToFloat();
float h_b = onnxruntime::BFloat16(*reinterpret_cast<const uint16_t*>(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<const uint16_t*>(alpha)).ToFloat();
float h_b = onnxruntime::BFloat16(*reinterpret_cast<const uint16_t*>(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

9
onnxruntime/core/providers/cuda/tensor/cast_op.cc
View file

@ -11,6 +11,9 @@ namespace cuda {
const std::vector<MLDataType> castOpTypeConstraints{
DataTypeImpl::GetTensorType<MLFloat16>(),
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
DataTypeImpl::GetTensorType<BFloat16>(),
#endif
DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<int8_t>(),
@ -76,6 +79,9 @@ Status Cast<SrcT>::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

25
onnxruntime/core/providers/cuda/tensor/gather_nd.cc
View file

@ -109,6 +109,21 @@ Status GatherNDBase::PrepareCompute(
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()), \
GatherND<TIndex>);
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define GATHER_ND_T_TENSOR_TYPES {DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>(), \
DataTypeImpl::GetTensorType<int64_t>()}
#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, BFloat16
#else
#define GATHER_ND_T_TENSOR_TYPES {DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<int64_t>()}
#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<MLDataType>{ \
DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<int64_t>(), \
}) \
.TypeConstraint("T", GATHER_ND_T_TENSOR_TYPES) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()), \
GatherND<TIndex>);
@ -187,7 +196,7 @@ Status GatherND<TIndex>::ComputeInternal(OpKernelContext* context) const {
const void* const kernel_input_data = input_tensor->DataRaw();
void* const kernel_output_data = output_tensor->MutableDataRaw();
utils::MLTypeCallDispatcher<GatherNDComputeImpl, float, MLFloat16, double, int64_t>
utils::MLTypeCallDispatcher<GatherNDComputeImpl, GATHER_ND_T_DATA_TYPES>
t_disp(input_tensor->GetElementType());
t_disp.Invoke(num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get());

3
onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu
View file

@ -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

34
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -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<int64_t>(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<bool>(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);

2
orttraining/orttraining/models/bert/main.cc
View file

@ -110,6 +110,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
("max_eval_steps", "Maximum number of eval steps.", cxxopts::value<int>()->default_value("100"))
("seed", "Random seed.", cxxopts::value<int64_t>()->default_value("-1"))
("use_mixed_precision", "Whether to use a mix of fp32 and fp16 arithmetic on GPU.", cxxopts::value<bool>()->default_value("false"))
("use_bfloat16", "Whether to use BFloat16 arithmetic on GPU.", cxxopts::value<bool>()->default_value("false"))
("enable_adasum", "Whether to use Adasum for allreduction.", cxxopts::value<bool>()->default_value("false"))
("allreduce_in_fp16", "Whether to do AllReduce in fp16. If false, AllReduce will be done in fp32", cxxopts::value<bool>()->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<bool>();
params.use_bfloat16 = flags["use_bfloat16"].as<bool>();
params.allreduce_in_mixed_precision_type = flags["allreduce_in_fp16"].as<bool>() && params.use_mixed_precision;
if (params.use_mixed_precision) {
printf("Mixed precision training is enabled.\n");

2
orttraining/orttraining/models/gpt2/main.cc
View file

@ -64,6 +64,7 @@ Status ParseArguments(int argc, char* argv[], GPT2Parameters& params, OrtParamet
cxxopts::value<int>()->default_value("1"))
("seed", "Random seed.", cxxopts::value<int64_t>()->default_value("-1"))
("use_mixed_precision", "Whether to use a mix of fp32 and fp16 arithmetic on GPU.", cxxopts::value<bool>()->default_value("false"))
("use_bfloat16", "Whether to use BFloat16 arithmetic on GPU.", cxxopts::value<bool>()->default_value("false"))
("enable_adasum", "Whether to use Adasum for allreduction.", cxxopts::value<bool>()->default_value("false"))
("allreduce_in_fp16", "Whether to do AllReduce in fp16. If false, AllReduce will be done in fp32", cxxopts::value<bool>()->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<bool>();
params.use_bfloat16 = flags["use_bfloat16"].as<bool>();
params.allreduce_in_mixed_precision_type = flags["allreduce_in_fp16"].as<bool>() && params.use_mixed_precision;
if (params.use_mixed_precision) {
printf("Mixed precision training is enabled.\n");

12
orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad.cc
View file

@ -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<MLFloat16, float, double>())
.TypeConstraint("T", BuildKernelDefConstraints<ALL_IEEE_FLOAT_DATA_TYPES>())
.MayInplace(0, 0),
BiasGeluGrad_dX<gelu_computation_mode::Default>);
@ -26,7 +32,7 @@ ONNX_OPERATOR_KERNEL_EX(
1,
kCudaExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T", BuildKernelDefConstraints<MLFloat16, float, double>())
.TypeConstraint("T", BuildKernelDefConstraints<ALL_IEEE_FLOAT_DATA_TYPES>())
.MayInplace(0, 0),
BiasGeluGrad_dX<gelu_computation_mode::Approximation>);
@ -70,7 +76,7 @@ Status BiasGeluGrad_dX<GeluComputationMode>::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);

5
orttraining/orttraining/training_ops/cuda/activation/bias_gelu_grad_impl.cu
View file

@ -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

69
orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc
View file

@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Scale)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Scale)>,
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
// Adam
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_float_float_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_float_float_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_BFloat16_float_float_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_BFloat16_BFloat16_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_float_BFloat16_float_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_BFloat16_BFloat16_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_int64_t_float_BFloat16_BFloat16_float_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_BFloat16_BFloat16_BFloat16_BFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_BFloat16_BFloat16_float_BFloat16, AdamOptimizer)>,
// Lamb
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float_float_float_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_BFloat16_float_BFloat16_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_BFloat16_float_float_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double_double_double_double_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_BFloat16_BFloat16_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_BFloat16_float_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_float_BFloat16_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float_BFloat16_float_float_BFloat16, LambOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, InPlaceAccumulator)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, InPlaceAccumulator)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, InPlaceAccumulator)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, SoftmaxGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, MixedPrecisionScale)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, LayerNormalizationGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_float, ReduceAllL2)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_BFloat16, ReduceAllL2)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, ReduceAllL2)>,
#endif
// P2P communication operators.
#if defined(ORT_USE_NCCL) || defined(USE_MPI)
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Send)>,

7
orttraining/orttraining/training_ops/cuda/math/isfinite.cuh
View file

@ -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

24
orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc
View file

@ -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<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>()}
#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<SrcT>()) \
.TypeConstraint("ScaleT", DataTypeImpl::GetTensorType<float>()) \
.TypeConstraint("DstT", DataTypeImpl::AllIEEEFloatTensorTypes()), \
.TypeConstraint("DstT", ALL_IEEE_FLOAT_TENSOR_TYPES), \
MixedPrecisionScale<SrcT>);
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<SrcT>::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<MLFloat16>::ComputeInternal(OpKernelContext* context) const;
template Status MixedPrecisionScale<float>::ComputeInternal(OpKernelContext* context) const;
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
REGISTER_MIXEDPRECISIONSCALE_KERNEL_TYPED(BFloat16)
template Status MixedPrecisionScale<BFloat16>::ComputeInternal(OpKernelContext* context) const;
#endif
} // namespace cuda
} // namespace onnxruntime

8
orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cu
View file

@ -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

30
orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc
View file

@ -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<BFloat16, is_log_softmax>( \
const BFloat16* dY, \
const TensorShape& input_shape, \
const BFloat16* Y, \
BFloat16* dX, \
cudnnHandle_t, \
int64_t axis) { \
typedef typename ToCudaType<BFloat16>::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<const CudaT*>(dY); \
auto Y_data = reinterpret_cast<const CudaT*>(Y); \
auto dX_data = reinterpret_cast<CudaT*>(dX); \
dispatch_softmax_backward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>( \
dX_data, dY_data, Y_data, gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(D), gsl::narrow_cast<int>(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<T>::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

4
orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu
View file

@ -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
}
}

27
orttraining/orttraining/training_ops/cuda/nn/dropout.cc
View file

@ -10,6 +10,17 @@
namespace onnxruntime {
namespace cuda {
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>()}
#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<bool>()) \
.InputMemoryType<OrtMemTypeCPUInput>(2), \
DropoutGrad);
@ -70,13 +81,13 @@ Status DropoutGrad::ComputeInternal(OpKernelContext* context) const {
float ratio_data = default_ratio_;
auto ratio = context->Input<Tensor>(2);
if (ratio) {
utils::MLTypeCallDispatcher<GetRatioDataImpl, float, MLFloat16, double> t_disp(ratio->GetElementType());
utils::MLTypeCallDispatcher<GetRatioDataImpl, ALL_IEEE_FLOAT_DATA_TYPES> t_disp(ratio->GetElementType());
t_disp.Invoke(ratio, ratio_data);
}
auto dX = context->Output(0, shape);
utils::MLTypeCallDispatcher<DropoutGradComputeImpl, float, MLFloat16, double> t_disp(dY->GetElementType());
utils::MLTypeCallDispatcher<DropoutGradComputeImpl, ALL_IEEE_FLOAT_DATA_TYPES> 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<bool>())
.InputMemoryType<OrtMemTypeCPUInput>(3)
.InputMemoryType<OrtMemTypeCPUInput>(4),
@ -161,7 +172,7 @@ Status BiasDropout::ComputeInternal(OpKernelContext* context) const {
float ratio_data = default_ratio_;
auto ratio = context->Input<Tensor>(3);
if (ratio) {
utils::MLTypeCallDispatcher<GetRatioDataImpl, float, MLFloat16, double> t_disp(ratio->GetElementType());
utils::MLTypeCallDispatcher<GetRatioDataImpl, ALL_IEEE_FLOAT_DATA_TYPES> 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<int>(dim));
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
utils::MLTypeCallDispatcherRet<Status, BiasDropoutComputeImpl, float, MLFloat16, double> t_disp(X->GetElementType());
utils::MLTypeCallDispatcherRet<Status, BiasDropoutComputeImpl, ALL_IEEE_FLOAT_DATA_TYPES> t_disp(X->GetElementType());
return t_disp.Invoke(GetDeviceProp(), N, fdm_dim, ratio_data, generator, *X, *bias, residual, *Y, mask_data);
}

7
orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu
View file

@ -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

3
orttraining/orttraining/training_ops/cuda/nn/layer_norm.cc
View file

@ -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 <typename T, typename U, bool simplified>
LayerNormGrad<T, U, simplified>::LayerNormGrad(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {

6
orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu
View file

@ -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
} // namespace onnxruntime

12
orttraining/orttraining/training_ops/cuda/optimizer/adam.cc
View file

@ -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 <typename T1, typename T2, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM, typename T_MIXED_PRECISION_FP>
Status AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>::ComputeInternal(OpKernelContext* ctx) const {
typedef typename ToCudaType<T1>::MappedType CudaT1;

12
orttraining/orttraining/training_ops/cuda/optimizer/adam.cu
View file

@ -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

5
orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc
View file

@ -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 <typename T>
Status ZeroGradient<T>::ComputeInternal(OpKernelContext* ctx) const {

5
orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cu
View file

@ -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

11
orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc
View file

@ -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,

34
orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu
View file

@ -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 <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
__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 <typename T1, typename T2, typename T3, typename T_GRAD_NORM>
__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 <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
__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 <typename TIn1, typename TIn2, typename TOut1, typename TOut2, typename TBuf>
@ -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

5
orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc
View file

@ -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

9
orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu
View file

@ -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 <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
__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

17
orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc
View file

@ -10,6 +10,15 @@
namespace onnxruntime {
namespace cuda {
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>()}
#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<OrtMemTypeCPUInput>(0)
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<MLFloat16>()})
.TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES)
.TypeConstraint("Tind", std::vector<MLDataType>{
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
@ -83,6 +91,11 @@ Status DispatchToGatherGradImpl(
} else if (utils::IsPrimitiveDataType<MLFloat16>(t_data_type)) {
return DispatchToGatherGradImplByTindex<MLFloat16>(
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<BFloat16>(t_data_type)) {
return DispatchToGatherGradImplByTindex<BFloat16>(
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);

3
orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu
View file

@ -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

16
orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad.cc
View file

@ -8,6 +8,17 @@
namespace onnxruntime {
namespace cuda {
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType<float>(), \
DataTypeImpl::GetTensorType<double>(), \
DataTypeImpl::GetTensorType<MLFloat16>(), \
DataTypeImpl::GetTensorType<BFloat16>()}
#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<MLFloat16>(), \
DataTypeImpl::GetTensorType<float>(), DataTypeImpl::GetTensorType<double>()}) \
KernelDefBuilder().TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()) \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()) \
.InputMemoryType<OrtMemTypeCPUInput>(0), \
@ -83,7 +93,7 @@ Status GatherNDGrad<TIndex>::ComputeInternal(OpKernelContext* context) const {
const void* const kernel_input_data = update_tensor->DataRaw();
void* const kernel_output_data = output_tensor->MutableDataRaw();
utils::MLTypeCallDispatcher<GatherNDGradComputeImpl, float, MLFloat16, double>
utils::MLTypeCallDispatcher<GatherNDGradComputeImpl, ALL_IEEE_FLOAT_DATA_TYPES>
t_disp(update_tensor->GetElementType());
t_disp.Invoke(num_slices, slice_size, kernel_input_data, kernel_output_data, input_slice_offsets_buffer.get());

3
orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.cu
View file

@ -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