From 435e14d60aa7fa7e53f25bc73b65d41642b951e4 Mon Sep 17 00:00:00 2001 From: ytaous <4484531+ytaous@users.noreply.github.com> Date: Mon, 7 Feb 2022 22:55:15 -0800 Subject: [PATCH] [ROCm] BFloat16 support (#10465) * bf16 support * minor clean up * UTs * fix build * UTs * UTs * merge commit 6b5504c * minor * ROCm code cleanup * fix build * fix build * minor Co-authored-by: Ethan Tao Co-authored-by: root --- .../core/providers/cuda/cudnn_common.cc | 5 + .../core/providers/cuda/math/softmax.cc | 23 +-- .../core/providers/rocm/math/softmax.cc | 20 +-- .../core/providers/rocm/miopen_common.cc | 5 + .../test/contrib_ops/element_wise_ops_test.cc | 72 +++++++++ .../test/contrib_ops/fastgelu_op_test.cc | 47 ++++++ .../test/contrib_ops/fused_matmul_op_test.cc | 95 ++++++++++++ .../cpu/reduction/reduction_ops_test.cc | 101 +++++++++++++ .../cuda/mixed_precision_scale_test.cc | 140 ++++++++++++++++++ .../training_ops/cuda/math/softmax_grad.cc | 26 +--- .../training_ops/rocm/math/softmax_grad.cc | 2 - 11 files changed, 471 insertions(+), 65 deletions(-) diff --git a/onnxruntime/core/providers/cuda/cudnn_common.cc b/onnxruntime/core/providers/cuda/cudnn_common.cc index ed1d792a52..0b2543ec15 100644 --- a/onnxruntime/core/providers/cuda/cudnn_common.cc +++ b/onnxruntime/core/providers/cuda/cudnn_common.cc @@ -136,6 +136,11 @@ cudnnDataType_t CudnnTensor::GetDataType() { return CUDNN_DATA_HALF; } +template <> +cudnnDataType_t CudnnTensor::GetDataType() { + return CUDNN_DATA_BFLOAT16; +} + template <> cudnnDataType_t CudnnTensor::GetDataType() { return CUDNN_DATA_INT8; diff --git a/onnxruntime/core/providers/cuda/math/softmax.cc b/onnxruntime/core/providers/cuda/math/softmax.cc index 9116178dda..b2a7d4dc6a 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.cc +++ b/onnxruntime/core/providers/cuda/math/softmax.cc @@ -43,24 +43,7 @@ Status SoftMaxComputeHelper( SPECIALIZED_SOFTMAX_HELPER_IMPL(float) SPECIALIZED_SOFTMAX_HELPER_IMPL(double) SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16) - -// cudnnSoftmaxForward/Backward doesn't support BFloat16. -#define SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(is_log_softmax) \ - template <> \ - Status SoftMaxComputeHelper(cudaStream_t stream, const BFloat16* X, \ - const TensorShape& input_shape, BFloat16* Y, int64_t axis) { \ - typedef typename ToCudaType::MappedType CudaT; \ - int64_t N = input_shape.SizeToDimension(axis); \ - int64_t D = input_shape.SizeFromDimension(axis); \ - auto Y_data = reinterpret_cast(Y); \ - auto X_data = reinterpret_cast(X); \ - dispatch_warpwise_softmax_forward, is_log_softmax>( \ - stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); \ - return Status::OK(); \ - } - -SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(true) -SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(false) +SPECIALIZED_SOFTMAX_HELPER_IMPL(BFloat16) #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ @@ -112,8 +95,8 @@ SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(false) (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ Softmax); - template - Status Softmax::ComputeInternal(OpKernelContext* ctx) const { +template +Status Softmax::ComputeInternal(OpKernelContext* ctx) const { const Tensor* X = ctx->Input(0); const TensorShape& input_shape{X->Shape()}; size_t rank = input_shape.NumDimensions(); diff --git a/onnxruntime/core/providers/rocm/math/softmax.cc b/onnxruntime/core/providers/rocm/math/softmax.cc index 612423e216..59e0e54049 100644 --- a/onnxruntime/core/providers/rocm/math/softmax.cc +++ b/onnxruntime/core/providers/rocm/math/softmax.cc @@ -44,25 +44,7 @@ SPECIALIZED_SOFTMAX_HELPER_IMPL(float) // MIOpen double data type not supported // SPECIALIZED_SOFTMAX_HELPER_IMPL(double) SPECIALIZED_SOFTMAX_HELPER_IMPL(MLFloat16) - -// 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, \ - const TensorShape& input_shape, BFloat16* Y, int64_t axis) { \ - typedef typename ToHipType::MappedType HipT; \ - int64_t N = input_shape.SizeToDimension(axis); \ - int64_t D = input_shape.SizeFromDimension(axis); \ - auto Y_data = reinterpret_cast(Y); \ - auto X_data = reinterpret_cast(X); \ - dispatch_warpwise_softmax_forward, is_log_softmax>( \ - stream, Y_data, X_data, gsl::narrow_cast(D), gsl::narrow_cast(D), gsl::narrow_cast(N)); \ - return Status::OK(); \ - } - -SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(true) -SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(false) +SPECIALIZED_SOFTMAX_HELPER_IMPL(BFloat16) #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ diff --git a/onnxruntime/core/providers/rocm/miopen_common.cc b/onnxruntime/core/providers/rocm/miopen_common.cc index 3de6c408cb..88c92df144 100644 --- a/onnxruntime/core/providers/rocm/miopen_common.cc +++ b/onnxruntime/core/providers/rocm/miopen_common.cc @@ -91,6 +91,11 @@ miopenDataType_t MiopenTensor::GetDataType() { return miopenHalf; } +template <> +miopenDataType_t MiopenTensor::GetDataType() { + return miopenBFloat16; +} + template <> miopenDataType_t MiopenTensor::GetDataType() { return miopenInt32; diff --git a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc index 03a67b539b..1c556890be 100644 --- a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc +++ b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc @@ -112,6 +112,78 @@ TEST(BiasGeluTest, Two_One_Dim) { RunBiasGeluTest(input_a_data, input_b_data, {2, 4}, {4}); } +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST(BiasGeluTest, Two_One_Dim_fp16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } +#endif + OpTester tester("BiasGelu", 1, onnxruntime::kMSDomain); + + std::vector A = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector B = { + -0.5f, 0.6f, 1.2f, 2.1f}; + + std::vector Y = ComputeGeluWithErf(Add_Simple(A, B)); + + std::vector f_A(8); + std::vector f_B(4); + std::vector f_Y(8); + ConvertFloatToMLFloat16(A.data(), f_A.data(), 8); + ConvertFloatToMLFloat16(B.data(), f_B.data(), 4); + ConvertFloatToMLFloat16(Y.data(), f_Y.data(), 8); + + tester.AddInput("A", {2, 4}, f_A); + tester.AddInput("B", {4}, f_B); + tester.AddOutput("Y", {2, 4}, f_Y); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: fp16 is not supported +} +#endif + +// failed test for CUDA (therefore ROCM as well) to be investigated +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST(BiasGeluTest, DISABLED_Two_One_Dim_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + OpTester tester("BiasGelu", 1, onnxruntime::kMSDomain); + + std::vector A = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector B = { + -0.5f, 0.6f, 1.2f, 2.1f}; + + std::vector Y = ComputeGeluWithErf(Add_Simple(A, B)); + + std::vector f_A = FloatsToBFloat16s(A); + std::vector f_B = FloatsToBFloat16s(B); + std::vector f_Y = FloatsToBFloat16s(Y); + + tester.AddInput("A", {2, 4}, f_A); + tester.AddInput("B", {4}, f_B); + tester.AddOutput("Y", {2, 4}, f_Y); + std::vector> execution_providers; +#ifdef USE_CUDA + execution_providers.push_back(DefaultCudaExecutionProvider()); +#elif USE_ROCM + execution_providers.push_back(DefaultRocmExecutionProvider()); +#endif + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + TEST(MathOpTest, ComplexMul) { if (DefaultCudaExecutionProvider() == nullptr) return; diff --git a/onnxruntime/test/contrib_ops/fastgelu_op_test.cc b/onnxruntime/test/contrib_ops/fastgelu_op_test.cc index 5614e065ac..60b9bec02b 100644 --- a/onnxruntime/test/contrib_ops/fastgelu_op_test.cc +++ b/onnxruntime/test/contrib_ops/fastgelu_op_test.cc @@ -197,5 +197,52 @@ TEST(FastGeluTest, FastGeluWithoutBiasFloat16) { RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, false, true); } + +// failed with device error, disabled for now +// CUDA only, ROCM has not been supported yet +#ifdef USE_CUDA +TEST(FastGeluTest, DISABLED_FastGeluWithBias_BFloat16) { + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } + OpTester tester("FastGelu", 1, onnxruntime::kMSDomain); + + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + + std::vector X = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector B = { + -0.5f, 0.6f, 1.2f, 2.1f}; + + std::vector Y = { + 0.1851806640625f, 0.054046630859375f, 1.0615234375f, 3.095703125f, + 0, 0.63037109375f, 1.3984375f, 1.3984375f}; + + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector bias_dims = {hidden_size}; + std::vector output_dims = input_dims; + + std::vector f_X = FloatsToBFloat16s(X); + std::vector f_B = FloatsToBFloat16s(B); + std::vector f_Y = FloatsToBFloat16s(Y); + + tester.AddInput("X", input_dims, f_X); + tester.AddInput("bias", bias_dims, f_B); + tester.AddOutput("Y", output_dims, f_Y); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + + + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc b/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc index c1c7a75a18..75857c61b8 100644 --- a/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc +++ b/onnxruntime/test/contrib_ops/fused_matmul_op_test.cc @@ -3,6 +3,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/common/cuda_op_test_utils.h" namespace onnxruntime { namespace test { @@ -269,6 +270,100 @@ TEST(FusedMatMulOpTest, FloatTypeTransposeBatch) { RunFusedMatMulTest("FusedMatMul", 1, true, true, true, true); } +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST(FusedMatMulOpTest, Float16_NoTranspose) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } +#endif + std::vector common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; + for (auto t : GenerateSimpleTestCases()) { + + OpTester test("FusedMatMul", 1, onnxruntime::kMSDomain); + + std::vector input0_dims(t.input0_dims); + std::vector input0_vals; + ProcessInputs(t.input0_dims, common_input_vals, false, false, input0_dims, input0_vals); + + std::vector input1_dims(t.input1_dims); + std::vector input1_vals; + ProcessInputs(t.input1_dims, common_input_vals, false, false, input1_dims, input1_vals); + + std::vector f_A(input0_vals.size()); + std::vector f_B(input1_vals.size()); + std::vector f_Y(t.expected_vals.size()); + ConvertFloatToMLFloat16(input0_vals.data(), f_A.data(), (int)input0_vals.size()); + ConvertFloatToMLFloat16(input1_vals.data(), f_B.data(), (int)input1_vals.size()); + ConvertFloatToMLFloat16(t.expected_vals.data(), f_Y.data(), (int)t.expected_vals.size()); + + test.AddInput("A", input0_dims, f_A); + test.AddInput("B", input1_dims, f_B, false); + + test.AddAttribute("transA", (int64_t)0); + test.AddAttribute("transB", (int64_t)0); + test.AddAttribute("transBatchA", (int64_t)0); + test.AddAttribute("transBatchB", (int64_t)0); + test.AddAttribute("alpha", 1.0f); + + test.AddOutput("Y", t.expected_dims, f_Y); + + // Disable TensorRT because of unsupported data type + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + } +} +#endif + +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST(FusedMatMulOpTest, BFloat16_NoTranspose) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support FP16"; + return; + } +#endif + std::vector common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; + for (auto t : GenerateSimpleTestCases()) { + + OpTester test("FusedMatMul", 1, onnxruntime::kMSDomain); + + std::vector input0_dims(t.input0_dims); + std::vector input0_vals; + ProcessInputs(t.input0_dims, common_input_vals, false, false, input0_dims, input0_vals); + + std::vector input1_dims(t.input1_dims); + std::vector input1_vals; + ProcessInputs(t.input1_dims, common_input_vals, false, false, input1_dims, input1_vals); + + std::vector f_A = FloatsToBFloat16s(input0_vals); + std::vector f_B = FloatsToBFloat16s(input1_vals); + std::vector f_Y = FloatsToBFloat16s(t.expected_vals); + + test.AddInput("A", input0_dims, f_A); + test.AddInput("B", input1_dims, f_B, false); + + test.AddAttribute("transA", (int64_t)0); + test.AddAttribute("transB", (int64_t)0); + test.AddAttribute("transBatchA", (int64_t)0); + test.AddAttribute("transBatchB", (int64_t)0); + test.AddAttribute("alpha", 1.0f); + + test.AddOutput("Y", t.expected_dims, f_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 + } // namespace transpose_matmul } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc b/orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc index eff6cd567c..46d775173c 100644 --- a/orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc @@ -6,6 +6,7 @@ #include #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/common/cuda_op_test_utils.h" namespace onnxruntime { namespace test { @@ -163,6 +164,106 @@ TEST_P(ReductionOpTest, ReduceAllL2HalfFloat) { } #endif +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST_P(ReductionOpTest, ReduceAllL2_BFloat16_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("ReduceAllL2", 1, onnxruntime::kMSDomain, true); + test.SetDeterminism(GetParam()); + + std::vector data0 = {1.0f, 2.0f, 3.0f}; + std::vector data0_bf16 = FloatsToBFloat16s(data0); + + std::vector data1 = {-1.0f, -2.0f}; + std::vector data1_bf16 = FloatsToBFloat16s(data1); + + std::vector result = {4.358898943540674f}; + std::vector result_bf16 = FloatsToBFloat16s(result); + + test.AddInput("data0", {3}, data0_bf16); + test.AddInput("data1", {2}, data1_bf16); + + test.AddOutput("reduced", {}, result_bf16); + + 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_P(ReductionOpTest, ReduceAllL2_BFloat16_Float) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + OpTester test("ReduceAllL2", 1, onnxruntime::kMSDomain, true); + test.SetDeterminism(GetParam()); + + std::vector data0 = {1.0f, 2.0f, 3.0f}; + std::vector data0_bf16 = FloatsToBFloat16s(data0); + + std::vector data1 = {-1.0f, -2.0f}; + std::vector data1_bf16 = FloatsToBFloat16s(data1); + + std::vector result = {4.358898943540674f}; + + test.AddInput("data0", {3}, data0_bf16); + test.AddInput("data1", {2}, data1_bf16); + + test.AddOutput("reduced", {}, result); + + 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_P(ReductionOpTest, ReduceAllL2_Float_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("ReduceAllL2", 1, onnxruntime::kMSDomain, true); + test.SetDeterminism(GetParam()); + + std::vector data0 = {1.0f, 2.0f, 3.0f}; + std::vector data1 = {-1.0f, -2.0f}; + + std::vector result = {4.358898943540674f}; + std::vector result_bf16 = FloatsToBFloat16s(result); + + test.AddInput("data0", {3}, data0); + test.AddInput("data1", {2}, data1); + + test.AddOutput("reduced", {}, result_bf16); + + 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 + void TestMultiTensorReduce( const int tensor_count, const int min_tensor_size, diff --git a/orttraining/orttraining/test/training_ops/cuda/mixed_precision_scale_test.cc b/orttraining/orttraining/test/training_ops/cuda/mixed_precision_scale_test.cc index f2b515a017..57b8d4e3af 100644 --- a/orttraining/orttraining/test/training_ops/cuda/mixed_precision_scale_test.cc +++ b/orttraining/orttraining/test/training_ops/cuda/mixed_precision_scale_test.cc @@ -3,6 +3,7 @@ #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" +#include "test/common/cuda_op_test_utils.h" namespace onnxruntime { namespace test { @@ -18,6 +19,16 @@ struct MixedPrecisionScaleInputOutput { output2_half.resize(output2.size()); ConvertFloatToMLFloat16(input2.data(), input2_half.data(), int(input2.size())); ConvertFloatToMLFloat16(output2.data(), output2_half.data(), int(output2.size())); + + input1_bf16.resize(input1.size()); + output1_bf16.resize(output1.size()); + std::vector input1_bf16 = FloatsToBFloat16s(input1); + std::vector output1_bf16 = FloatsToBFloat16s(output1); + + input2_bf16.resize(input2.size()); + output2_bf16.resize(output2.size()); + std::vector input2_bf16 = FloatsToBFloat16s(input2); + std::vector output2_bf16 = FloatsToBFloat16s(output2); } // Fp32 Inputs/Output @@ -32,6 +43,12 @@ struct MixedPrecisionScaleInputOutput { std::vector input2_half; std::vector output1_half; std::vector output2_half; + + // BF16 Inputs/Output + std::vector input1_bf16; + std::vector input2_bf16; + std::vector output1_bf16; + std::vector output2_bf16; }; TEST(CudaKernelTest, MixedPrecisionScaleF2F) { @@ -130,5 +147,128 @@ TEST(CudaKernelTest, MixedPrecisionScaleH2H) { test.Run(); } +#if defined(USE_CUDA) || defined(USE_ROCM) +TEST(CudaKernelTest, MixedPrecisionScale_bfloat16_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + MixedPrecisionScaleInputOutput data; + OpTester test("MixedPrecisionScale", 1, onnxruntime::kMSDomain); + test.AddAttribute("to", int64_t(ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16)); + test.AddInput("scale", {1}, data.scale); + test.AddInput("input1", {3}, data.input1_bf16); + test.AddOutput("output1", {3}, data.output1_bf16); + + 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); +} + +// failed with data error, disabled for now +TEST(CudaKernelTest, DISABLED_MixedPrecisionScale_float_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + MixedPrecisionScaleInputOutput data; + OpTester test("MixedPrecisionScale", 1, onnxruntime::kMSDomain); + test.AddAttribute("to", int64_t(ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16)); + test.AddInput("scale", {1}, data.scale); + test.AddInput("input1", {3}, data.input1); + test.AddOutput("output1", {3}, data.output1_bf16); + + 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(CudaKernelTest, DISABLED_MixedPrecisionScale_bfloat16_float) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + MixedPrecisionScaleInputOutput data; + OpTester test("MixedPrecisionScale", 1, onnxruntime::kMSDomain); + test.AddAttribute("to", int64_t(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT)); + test.AddInput("scale", {1}, data.scale); + test.AddInput("input1", {3}, data.input1_bf16); + test.AddOutput("output1", {3}, data.output1); + + 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(CudaKernelTest, DISABLED_MixedPrecisionScale_half_bfloat16) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + MixedPrecisionScaleInputOutput data; + OpTester test("MixedPrecisionScale", 1, onnxruntime::kMSDomain); + test.AddAttribute("to", int64_t(ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16)); + test.AddInput("scale", {1}, data.scale); + test.AddInput("input1", {3}, data.input1_half); + test.AddOutput("output1", {3}, data.output1_bf16); + + 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(CudaKernelTest, DISABLED_MixedPrecisionScale_bfloat16_half) { +#ifdef USE_CUDA + int min_cuda_architecture = 530; + if (!HasCudaEnvironment(min_cuda_architecture)) { + LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16"; + return; + } +#endif + MixedPrecisionScaleInputOutput data; + OpTester test("MixedPrecisionScale", 1, onnxruntime::kMSDomain); + test.AddAttribute("to", int64_t(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16)); + test.AddInput("scale", {1}, data.scale); + test.AddInput("input1", {3}, data.input1_bf16); + test.AddOutput("output1", {3}, data.output1_half); + + 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 + } // namespace test } // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc b/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc index 4d4c4bdc76..1a8af5045b 100644 --- a/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc @@ -62,28 +62,6 @@ Status SoftMaxGradComputeHelper( return Status::OK(); } -// cudnnSoftmaxForward/Backward doesn't support BFloat16. -#define SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(is_log_softmax) \ - template <> \ - Status SoftMaxGradComputeHelper(cudaStream_t stream, const BFloat16* dY, \ - const TensorShape& input_shape, const BFloat16* Y, \ - BFloat16* dX, cudnnHandle_t, int64_t axis) { \ - typedef typename ToCudaType::MappedType CudaT; \ - const int64_t normalized_axis = HandleNegativeAxis(axis, input_shape.NumDimensions()); \ - int64_t N = input_shape.SizeToDimension(normalized_axis); \ - int64_t D = input_shape.SizeFromDimension(normalized_axis); \ - auto dY_data = reinterpret_cast(dY); \ - auto Y_data = reinterpret_cast(Y); \ - auto dX_data = reinterpret_cast(dX); \ - dispatch_softmax_backward, is_log_softmax>( \ - 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, \ @@ -121,8 +99,8 @@ SPECIALIZED_SOFTMAXGRAD_HELPER_IMPL_BFloat16(false) (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ SoftmaxGrad); - template - Status SoftmaxGrad::ComputeInternal(OpKernelContext* ctx) const { +template +Status SoftmaxGrad::ComputeInternal(OpKernelContext* ctx) const { const Tensor* dY = ctx->Input(0); const TensorShape& input_shape{dY->Shape()}; const Tensor* Y = ctx->Input(1); diff --git a/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc b/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc index 8481beda36..b39fdf0e9c 100644 --- a/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc +++ b/orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc @@ -62,8 +62,6 @@ 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, \