From 63198a6566e8b3703f84477db0ec7a313cc69dd2 Mon Sep 17 00:00:00 2001 From: ytaous <4484531+ytaous@users.noreply.github.com> Date: Thu, 3 Feb 2022 11:31:14 -0800 Subject: [PATCH] [ROCm] BFloat16 support (#10447) * bf16 support * bf16 support * UTs * fix build * fix UTs Co-authored-by: root --- .../contrib_ops/rocm/rocm_contrib_kernels.cc | 20 ++ .../core/providers/rocm/math/softmax.cc | 3 +- .../providers/rocm/rocm_execution_provider.cc | 16 +- .../cpu/activation/activation_op_test.cc | 193 ++++++++++++++++++ .../test/providers/cpu/math/gemm_test.cc | 2 +- .../test/providers/cpu/math/matmul_test.cc | 2 +- .../test/providers/cpu/math/softmax_test.cc | 56 +++++ .../test/providers/provider_test_utils.h | 6 + .../training_ops/rocm/math/softmax_grad.cc | 24 +++ .../rocm/math/softmax_grad_impl.cu | 2 + .../rocm/reduction/reduction_all.cc | 3 + .../rocm/rocm_training_kernels.cc | 26 +++ 12 files changed, 342 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc index bd940e9f52..99c4755c36 100644 --- a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc @@ -17,6 +17,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, BiasGelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, BiasGelu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, BiasGelu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, BiasGelu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, TransposeMatMul); // backward compatibility +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, TransposeMatMul); // backward compatibility +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, TransposeMatMul); // backward compatibility class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, FusedMatMul); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, FusedMatMul); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, FusedMatMul); @@ -76,6 +80,11 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_int8_t, QAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_int8_t, QAttention); +// class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, FastGelu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, TransposeMatMul); // backward compatibility +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, FusedMatMul); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, BFloat16_float, LayerNormalization); + template <> KernelCreateInfo BuildKernelCreateInfo() { KernelCreateInfo info; @@ -93,6 +102,10 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // backward compatibility + BuildKernelCreateInfo, // backward compatibility + BuildKernelCreateInfo, // backward compatibility BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -152,6 +165,13 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) { // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo + + // BuildKernelCreateInfo, + // TransposedMatMul is still here for backward compatibility + BuildKernelCreateInfo, // backward compatibility + BuildKernelCreateInfo, + BuildKernelCreateInfo, + }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/rocm/math/softmax.cc b/onnxruntime/core/providers/rocm/math/softmax.cc index ecbd918d7f..612423e216 100644 --- a/onnxruntime/core/providers/rocm/math/softmax.cc +++ b/onnxruntime/core/providers/rocm/math/softmax.cc @@ -45,7 +45,8 @@ SPECIALIZED_SOFTMAX_HELPER_IMPL(float) // SPECIALIZED_SOFTMAX_HELPER_IMPL(double) SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16) -// miopenSoftmaxForward/Backward doesn't support BFloat16. +// cudnnSoftmaxForward/Backward doesn't support BFloat16. +// apply the same for miopen for now #define SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(is_log_softmax) \ template <> \ Status SoftMaxComputeHelper(hipStream_t stream, const BFloat16* X, \ diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc index 870c0f8180..b024f38831 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc @@ -1127,9 +1127,9 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kO class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Cast); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Softmax); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, MatMul); -// class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Relu); -// class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Sigmoid); -// class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, 13, BFloat16, Relu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Sigmoid); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum); @@ -1192,7 +1192,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 14, BFloat16, Sub); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 14, BFloat16, Mul); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 14, BFloat16, Div); -// class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 14, BFloat16, Relu); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 14, BFloat16, Relu); //OpSet 15 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 15, Pow); @@ -1971,9 +1971,9 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -2035,7 +2035,7 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, // OpSet 15 BuildKernelCreateInfo, diff --git a/onnxruntime/test/providers/cpu/activation/activation_op_test.cc b/onnxruntime/test/providers/cpu/activation/activation_op_test.cc index 75ead657ae..82e9690718 100644 --- a/onnxruntime/test/providers/cpu/activation/activation_op_test.cc +++ b/onnxruntime/test/providers/cpu/activation/activation_op_test.cc @@ -3,6 +3,7 @@ #include "activation_op_test.h" #include "core/providers/cpu/activation/activations.h" +#include "test/common/cuda_op_test_utils.h" namespace onnxruntime { namespace test { @@ -120,6 +121,198 @@ TEST_F(ActivationOpTest, Relu) { /*opset_version= */ 14); } +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST_F(ActivationOpTest, Sigmoid_fp16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } +#endif + OpTester test("Sigmoid", 14); + + auto formula = [](float x) { + auto y = 1.f / (1.f + std::exp(-std::abs(x))); // safe sigmoid + y = x > 0 ? y : 1 - y; + return y; + }; + + std::vector X = input_values.front(); + std::vector Y; + for (unsigned i = 0; i < X.size(); i++) + Y.push_back(formula(X[i])); + std::vector dims{(int64_t)X.size()}; + + std::vector f_X(X.size()); + std::vector f_Y(Y.size()); + ConvertFloatToMLFloat16(X.data(), f_X.data(), static_cast(X.size())); + ConvertFloatToMLFloat16(Y.data(), f_Y.data(), static_cast(Y.size())); + + test.AddInput("X", dims, f_X); + test.AddOutput("Y", dims, f_Y); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +TEST_F(ActivationOpTest, Tanh_fp16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } +#endif + OpTester test("Tanh", 14); + + auto formula = [](float x) { return std::tanh(x); }; + + std::vector X = input_values.front(); + std::vector Y; + for (unsigned i = 0; i < X.size(); i++) + Y.push_back(formula(X[i])); + std::vector dims{(int64_t)X.size()}; + + std::vector f_X(X.size()); + std::vector f_Y(Y.size()); + ConvertFloatToMLFloat16(X.data(), f_X.data(), static_cast(X.size())); + ConvertFloatToMLFloat16(Y.data(), f_Y.data(), static_cast(Y.size())); + + test.AddInput("X", dims, f_X); + test.AddOutput("Y", dims, f_Y); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +TEST_F(ActivationOpTest, Relu_fp16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } +#endif + OpTester test("Relu", 14); + + auto formula = [](float x) { return std::max(x, 0.0f); }; + + std::vector X = input_values.front(); + std::vector Y; + for (unsigned i = 0; i < X.size(); i++) + Y.push_back(formula(X[i])); + std::vector dims{(int64_t)X.size()}; + + std::vector f_X(X.size()); + std::vector f_Y(Y.size()); + ConvertFloatToMLFloat16(X.data(), f_X.data(), static_cast(X.size())); + ConvertFloatToMLFloat16(Y.data(), f_Y.data(), static_cast(Y.size())); + + test.AddInput("X", dims, f_X); + test.AddOutput("Y", dims, f_Y); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} +#endif + +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST_F(ActivationOpTest, Sigmoid_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + OpTester test("Sigmoid", 14); + + auto formula = [](float x) { + auto y = 1.f / (1.f + std::exp(-std::abs(x))); // safe sigmoid + y = x > 0 ? y : 1 - y; + return y; + }; + + std::vector X = input_values.front(); + std::vector Y; + for (unsigned i = 0; i < X.size(); i++) + Y.push_back(formula(X[i])); + std::vector dims{(int64_t)X.size()}; + + std::vector bf_X = FloatsToBFloat16s(X); + std::vector bf_Y = FloatsToBFloat16s(Y); + + test.AddInput("X", dims, bf_X); + test.AddOutput("Y", dims, bf_Y); + std::vector> execution_providers; +#ifdef USE_CUDA + execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif USE_ROCM + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST_F(ActivationOpTest, Tanh_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + OpTester test("Tanh", 14); + + auto formula = [](float x) { return std::tanh(x); }; + + std::vector X = input_values.front(); + std::vector Y; + for (unsigned i = 0; i < X.size(); i++) + Y.push_back(formula(X[i])); + std::vector dims{(int64_t)X.size()}; + + std::vector bf_X = FloatsToBFloat16s(X); + std::vector bf_Y = FloatsToBFloat16s(Y); + + test.AddInput("X", dims, bf_X); + test.AddOutput("Y", dims, bf_Y); + std::vector> execution_providers; +#ifdef USE_CUDA + execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif USE_ROCM + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST_F(ActivationOpTest, Relu_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + OpTester test("Relu", 14); + + auto formula = [](float x) { return std::max(x, 0.0f); }; + + std::vector X = input_values.front(); + std::vector Y; + for (unsigned i = 0; i < X.size(); i++) + Y.push_back(formula(X[i])); + std::vector dims{(int64_t)X.size()}; + + std::vector bf_X = FloatsToBFloat16s(X); + std::vector bf_Y = FloatsToBFloat16s(Y); + + test.AddInput("X", dims, bf_X); + test.AddOutput("Y", dims, bf_Y); + std::vector> execution_providers; +#ifdef USE_CUDA + execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif USE_ROCM + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + TEST_F(ActivationOpTest, Elu) { float alpha = 0.1f; TestActivationOp("Elu", diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc index e96fbc1c9c..8d0018c6e8 100644 --- a/onnxruntime/test/providers/cpu/math/gemm_test.cc +++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc @@ -91,7 +91,7 @@ TEST(GemmOpTest, GemmNoTrans_bfloat16) { #ifdef USE_CUDA int min_cuda_architecture = 530; if (!HasCudaEnvironment(min_cuda_architecture)) { - LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; return; } #endif diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc index 6d2afcdc53..5633d12d75 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc @@ -199,7 +199,7 @@ TEST(MathOpTest, MatMul_BFloat16) { #ifdef USE_CUDA int min_cuda_architecture = 530; if (!HasCudaEnvironment(min_cuda_architecture)) { - LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; return; } #endif diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index eeb22f2271..b821b16bda 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -4,6 +4,7 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/common/cuda_op_test_utils.h" #include namespace onnxruntime { @@ -45,6 +46,61 @@ TEST(SoftmaxOperator, Simple) { RunTest(x_vals, expected_vals, dimensions); } +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST(SoftmaxOperator, Simple_fp16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } +#endif + OpTester test("Softmax", 14); + + int64_t axis = 1; + test.AddAttribute("axis", axis); + + std::vector X = {-1.0f, 0.0f, 1.0f}; + std::vector Y = {0.09003058f, 0.24472848f, 0.66524094f}; + std::vector dimensions = {1, 3}; + + std::vector f_X(3); + std::vector f_Y(3); + ConvertFloatToMLFloat16(X.data(), f_X.data(), 3); + ConvertFloatToMLFloat16(Y.data(), f_Y.data(), 3); + + test.AddInput("X", dimensions, f_X); + test.AddOutput("Y", dimensions, f_Y); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); +} +#endif + +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST(SoftmaxOperator, Simple_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + OpTester test("Softmax", 14); + + int64_t axis = 1; + test.AddAttribute("axis", axis); + + test.AddInput("X", {1, 3}, MakeBFloat16({-1.0f, 0.0f, 1.0f})); + test.AddOutput("Y", {1, 3}, MakeBFloat16({0.09003058f, 0.24472848f, 0.66524094f})); + std::vector> execution_providers; +#ifdef USE_CUDA + execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif USE_ROCM + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + TEST(SoftmaxOperator, LargeNumber) { // x = np.array([[0, 1, 2, 3], [10000, 10001, 10002, 10003]]).astype(np.float32) // expected output[[0.0320586, 0.08714432, 0.23688284, 0.64391428], diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index 6958d44a35..74bae67c39 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -1198,6 +1198,12 @@ inline std::vector MakeBFloat16(const std::initializer_list& in return output; } +inline std::vector FloatsToBFloat16s(const std::vector& input) { + std::vector output; + std::transform(input.begin(), input.end(), std::back_inserter(output), [](float f) { return BFloat16(f); }); + return output; +} + inline CheckParams MakeCheckParams(const OpTester::Data& d) { return CheckParams{d.sort_output_, d.absolute_error_, d.relative_error_}; } diff --git a/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc b/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc index c4a52cf9a8..8481beda36 100644 --- a/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc +++ b/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc @@ -62,6 +62,29 @@ Status SoftMaxGradComputeHelper( return Status::OK(); } +// cudnnSoftmaxForward/Backward doesn't support BFloat16. +// apply the same for miopen for now +#define SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(is_log_softmax) \ + template <> \ + Status SoftMaxGradComputeHelper(hipStream_t stream, const BFloat16* dY, \ + const TensorShape& input_shape, const BFloat16* Y, \ + BFloat16* dX, miopenHandle_t, int64_t axis) { \ + typedef typename ToHipType::MappedType HipT; \ + const int64_t normalized_axis = HandleNegativeAxis(axis, input_shape.NumDimensions()); \ + int64_t N = input_shape.SizeToDimension(normalized_axis); \ + int64_t D = input_shape.SizeFromDimension(normalized_axis); \ + auto dY_data = reinterpret_cast(dY); \ + auto Y_data = reinterpret_cast(Y); \ + auto dX_data = reinterpret_cast(dX); \ + dispatch_softmax_backward, is_log_softmax>( \ + stream, dX_data, dY_data, Y_data, gsl::narrow_cast(D), gsl::narrow_cast(D), \ + gsl::narrow_cast(N)); \ + return Status::OK(); \ + } + +SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(true) +SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(false) + #define REGISTER_GRADIENT_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ SoftmaxGrad, \ @@ -181,6 +204,7 @@ Status SoftmaxGrad::ComputeInternal(OpKernelContext* ctx) const { SPECIALIZED_GRADIENT(float) // SPECIALIZED_GRADIENT(double) SPECIALIZED_GRADIENT(MLFloat16) +SPECIALIZED_GRADIENT(BFloat16) } // namespace rocm } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/rocm/math/softmax_grad_impl.cu b/orttraining/orttraining/training_ops/rocm/math/softmax_grad_impl.cu index 73343cafc9..81a176fc73 100644 --- a/orttraining/orttraining/training_ops/rocm/math/softmax_grad_impl.cu +++ b/orttraining/orttraining/training_ops/rocm/math/softmax_grad_impl.cu @@ -17,6 +17,7 @@ /* Modifications Copyright (c) Microsoft. */ // The code below is mostly copied from Pytorch PersistentSoftmax.cuh +#include "hip/hip_runtime.h" #include "orttraining/training_ops/rocm/math/softmax_grad.h" @@ -190,6 +191,7 @@ template void dispatch_softmax_backward(hipStrea SPECIALIZED_SOFTMAX_GRAD_IMPL(float, float, float) SPECIALIZED_SOFTMAX_GRAD_IMPL(half, half, float) SPECIALIZED_SOFTMAX_GRAD_IMPL(double, double, double) +SPECIALIZED_SOFTMAX_GRAD_IMPL(BFloat16, BFloat16, float) } } diff --git a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc index d1e7985504..1577dafe9d 100644 --- a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc +++ b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc @@ -104,6 +104,9 @@ 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) +REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, BFloat16, float) +REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, float, BFloat16) +REGISTER_REDUCE_ALL_KERNEL_TYPED(ReduceAllL2, BFloat16, BFloat16) } // namespace rocm } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc index b594c78196..65b0bfb070 100644 --- a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc +++ b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc @@ -57,12 +57,16 @@ class ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider class ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 12, 12, float, int64_t, SoftmaxCrossEntropyLoss); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, MLFloat16, int64_t, SoftmaxCrossEntropyLoss); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, float, int64_t, SoftmaxCrossEntropyLoss); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, int64_t, SoftmaxCrossEntropyLoss); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossGrad); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossGrad); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossGrad); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternal); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternal); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternal); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, int64_t, SoftmaxCrossEntropyLossInternalGrad); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad); +class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, int64_t, SoftmaxCrossEntropyLossInternalGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, SoftmaxGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, SoftmaxGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, SoftmaxGrad); @@ -167,6 +171,15 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_MLFloat16_float, BatchNormInternal); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16_float_float, BatchNormInternal); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, SoftmaxGrad); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, SoftmaxGrad_13); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16, MixedPrecisionScale); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16_float, LayerNormalizationGrad); + +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16_float, ReduceAllL2); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float_BFloat16, ReduceAllL2); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BFloat16_BFloat16, ReduceAllL2); + #if defined(ORT_USE_NCCL) || defined(USE_MPI) // P2P communication operators. class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, Send); @@ -263,12 +276,16 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, // BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -349,6 +366,15 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + // P2P communication operators. #if defined(ORT_USE_NCCL) || defined(USE_MPI) BuildKernelCreateInfo,