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

[ROCm] BFloat16 support (#10465)

Browse source 
* 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 <ettao@microsoft.com@orttrainingdev7.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
Co-authored-by: root <root@GCRAMDRR1-MI100-087.redmond.corp.microsoft.com>
This commit is contained in:
ytaous 2022-02-07 22:55:15 -08:00 committed by GitHub
parent c696da36c7
commit 435e14d60a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 471 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -136,6 +136,11 @@ cudnnDataType_t CudnnTensor::GetDataType<half>() {
return CUDNN_DATA_HALF;
}
template <>
cudnnDataType_t CudnnTensor::GetDataType<BFloat16>() {
return CUDNN_DATA_BFLOAT16;
}
template <>
cudnnDataType_t CudnnTensor::GetDataType<int8_t>() {
return CUDNN_DATA_INT8;

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

@ -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<BFloat16, is_log_softmax>(cudaStream_t stream, const BFloat16* X, \
const TensorShape& input_shape, BFloat16* Y, 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_warpwise_softmax_forward<CudaT, CudaT, AccumulationType_t<CudaT>, is_log_softmax>( \
stream, 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)
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<T>()), \
Softmax<T>);
template <typename T>
Status Softmax<T>::ComputeInternal(OpKernelContext* ctx) const {
template <typename T>
Status Softmax<T>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor* X = ctx->Input<Tensor>(0);
const TensorShape& input_shape{X->Shape()};
size_t rank = input_shape.NumDimensions();

20
onnxruntime/core/providers/rocm/math/softmax.cc
View file

@ -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<BFloat16, is_log_softmax>(hipStream_t stream, const BFloat16* X, \
const TensorShape& input_shape, BFloat16* Y, int64_t axis) { \
typedef typename ToHipType<BFloat16>::MappedType HipT; \
int64_t N = input_shape.SizeToDimension(axis); \
int64_t D = input_shape.SizeFromDimension(axis); \
auto Y_data = reinterpret_cast<HipT*>(Y); \
auto X_data = reinterpret_cast<const HipT*>(X); \
dispatch_warpwise_softmax_forward<HipT, HipT, AccumulationType_t<HipT>, is_log_softmax>( \
stream, 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)
SPECIALIZED_SOFTMAX_HELPER_IMPL(BFloat16)
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \

5
onnxruntime/core/providers/rocm/miopen_common.cc
View file

@ -91,6 +91,11 @@ miopenDataType_t MiopenTensor::GetDataType<half>() {
return miopenHalf;
}
template <>
miopenDataType_t MiopenTensor::GetDataType<BFloat16>() {
return miopenBFloat16;
}
template <>
miopenDataType_t MiopenTensor::GetDataType<int32_t>() {
return miopenInt32;

72
onnxruntime/test/contrib_ops/element_wise_ops_test.cc
View file

@ -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<float> A = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> B = {
-0.5f, 0.6f, 1.2f, 2.1f};
std::vector<float> Y = ComputeGeluWithErf(Add_Simple(A, B));
std::vector<MLFloat16> f_A(8);
std::vector<MLFloat16> f_B(4);
std::vector<MLFloat16> 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<MLFloat16>("A", {2, 4}, f_A);
tester.AddInput<MLFloat16>("B", {4}, f_B);
tester.AddOutput<MLFloat16>("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<float> A = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> B = {
-0.5f, 0.6f, 1.2f, 2.1f};
std::vector<float> Y = ComputeGeluWithErf(Add_Simple(A, B));
std::vector<BFloat16> f_A = FloatsToBFloat16s(A);
std::vector<BFloat16> f_B = FloatsToBFloat16s(B);
std::vector<BFloat16> f_Y = FloatsToBFloat16s(Y);
tester.AddInput<BFloat16>("A", {2, 4}, f_A);
tester.AddInput<BFloat16>("B", {4}, f_B);
tester.AddOutput<BFloat16>("Y", {2, 4}, f_Y);
std::vector<std::unique_ptr<IExecutionProvider>> 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;

47
onnxruntime/test/contrib_ops/fastgelu_op_test.cc
View file

@ -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<float> X = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> B = {
-0.5f, 0.6f, 1.2f, 2.1f};
std::vector<float> Y = {
0.1851806640625f, 0.054046630859375f, 1.0615234375f, 3.095703125f,
0, 0.63037109375f, 1.3984375f, 1.3984375f};
std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
std::vector<int64_t> bias_dims = {hidden_size};
std::vector<int64_t> output_dims = input_dims;
std::vector<BFloat16> f_X = FloatsToBFloat16s(X);
std::vector<BFloat16> f_B = FloatsToBFloat16s(B);
std::vector<BFloat16> f_Y = FloatsToBFloat16s(Y);
tester.AddInput<BFloat16>("X", input_dims, f_X);
tester.AddInput<BFloat16>("bias", bias_dims, f_B);
tester.AddOutput<BFloat16>("Y", output_dims, f_Y);
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
#endif
} // namespace test
} // namespace onnxruntime

95
onnxruntime/test/contrib_ops/fused_matmul_op_test.cc
View file

@ -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<float>("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<float> common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
for (auto t : GenerateSimpleTestCases<float>()) {
OpTester test("FusedMatMul", 1, onnxruntime::kMSDomain);
std::vector<int64_t> input0_dims(t.input0_dims);
std::vector<float> input0_vals;
ProcessInputs(t.input0_dims, common_input_vals, false, false, input0_dims, input0_vals);
std::vector<int64_t> input1_dims(t.input1_dims);
std::vector<float> input1_vals;
ProcessInputs(t.input1_dims, common_input_vals, false, false, input1_dims, input1_vals);
std::vector<MLFloat16> f_A(input0_vals.size());
std::vector<MLFloat16> f_B(input1_vals.size());
std::vector<MLFloat16> 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<MLFloat16>("A", input0_dims, f_A);
test.AddInput<MLFloat16>("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<MLFloat16>("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<float> common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
for (auto t : GenerateSimpleTestCases<float>()) {
OpTester test("FusedMatMul", 1, onnxruntime::kMSDomain);
std::vector<int64_t> input0_dims(t.input0_dims);
std::vector<float> input0_vals;
ProcessInputs(t.input0_dims, common_input_vals, false, false, input0_dims, input0_vals);
std::vector<int64_t> input1_dims(t.input1_dims);
std::vector<float> input1_vals;
ProcessInputs(t.input1_dims, common_input_vals, false, false, input1_dims, input1_vals);
std::vector<BFloat16> f_A = FloatsToBFloat16s(input0_vals);
std::vector<BFloat16> f_B = FloatsToBFloat16s(input1_vals);
std::vector<BFloat16> f_Y = FloatsToBFloat16s(t.expected_vals);
test.AddInput<BFloat16>("A", input0_dims, f_A);
test.AddInput<BFloat16>("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<BFloat16>("Y", t.expected_dims, f_Y);
std::vector<std::unique_ptr<IExecutionProvider>> 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

101
orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc
View file

@ -6,6 +6,7 @@
#include <type_traits>
#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<float> data0 = {1.0f, 2.0f, 3.0f};
std::vector<BFloat16> data0_bf16 = FloatsToBFloat16s(data0);
std::vector<float> data1 = {-1.0f, -2.0f};
std::vector<BFloat16> data1_bf16 = FloatsToBFloat16s(data1);
std::vector<float> result = {4.358898943540674f};
std::vector<BFloat16> result_bf16 = FloatsToBFloat16s(result);
test.AddInput<BFloat16>("data0", {3}, data0_bf16);
test.AddInput<BFloat16>("data1", {2}, data1_bf16);
test.AddOutput<BFloat16>("reduced", {}, result_bf16);
std::vector<std::unique_ptr<IExecutionProvider>> 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<float> data0 = {1.0f, 2.0f, 3.0f};
std::vector<BFloat16> data0_bf16 = FloatsToBFloat16s(data0);
std::vector<float> data1 = {-1.0f, -2.0f};
std::vector<BFloat16> data1_bf16 = FloatsToBFloat16s(data1);
std::vector<float> result = {4.358898943540674f};
test.AddInput<BFloat16>("data0", {3}, data0_bf16);
test.AddInput<BFloat16>("data1", {2}, data1_bf16);
test.AddOutput<float>("reduced", {}, result);
std::vector<std::unique_ptr<IExecutionProvider>> 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<float> data0 = {1.0f, 2.0f, 3.0f};
std::vector<float> data1 = {-1.0f, -2.0f};
std::vector<float> result = {4.358898943540674f};
std::vector<BFloat16> result_bf16 = FloatsToBFloat16s(result);
test.AddInput<float>("data0", {3}, data0);
test.AddInput<float>("data1", {2}, data1);
test.AddOutput<BFloat16>("reduced", {}, result_bf16);
std::vector<std::unique_ptr<IExecutionProvider>> 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,

140
orttraining/orttraining/test/training_ops/cuda/mixed_precision_scale_test.cc
View file

@ -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<BFloat16> input1_bf16 = FloatsToBFloat16s(input1);
std::vector<BFloat16> output1_bf16 = FloatsToBFloat16s(output1);
input2_bf16.resize(input2.size());
output2_bf16.resize(output2.size());
std::vector<BFloat16> input2_bf16 = FloatsToBFloat16s(input2);
std::vector<BFloat16> output2_bf16 = FloatsToBFloat16s(output2);
}
// Fp32 Inputs/Output
@ -32,6 +43,12 @@ struct MixedPrecisionScaleInputOutput {
std::vector<MLFloat16> input2_half;
std::vector<MLFloat16> output1_half;
std::vector<MLFloat16> output2_half;
// BF16 Inputs/Output
std::vector<BFloat16> input1_bf16;
std::vector<BFloat16> input2_bf16;
std::vector<BFloat16> output1_bf16;
std::vector<BFloat16> 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<float>("scale", {1}, data.scale);
test.AddInput<BFloat16>("input1", {3}, data.input1_bf16);
test.AddOutput<BFloat16>("output1", {3}, data.output1_bf16);
std::vector<std::unique_ptr<IExecutionProvider>> 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<float>("scale", {1}, data.scale);
test.AddInput<float>("input1", {3}, data.input1);
test.AddOutput<BFloat16>("output1", {3}, data.output1_bf16);
std::vector<std::unique_ptr<IExecutionProvider>> 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<float>("scale", {1}, data.scale);
test.AddInput<BFloat16>("input1", {3}, data.input1_bf16);
test.AddOutput<float>("output1", {3}, data.output1);
std::vector<std::unique_ptr<IExecutionProvider>> 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<float>("scale", {1}, data.scale);
test.AddInput<MLFloat16>("input1", {3}, data.input1_half);
test.AddOutput<BFloat16>("output1", {3}, data.output1_bf16);
std::vector<std::unique_ptr<IExecutionProvider>> 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<float>("scale", {1}, data.scale);
test.AddInput<BFloat16>("input1", {3}, data.input1_bf16);
test.AddOutput<MLFloat16>("output1", {3}, data.output1_half);
std::vector<std::unique_ptr<IExecutionProvider>> 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

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

@ -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<BFloat16, is_log_softmax>(cudaStream_t stream, 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>( \
stream, 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)
#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<T>()), \
SoftmaxGrad<T>);
template <typename T>
Status SoftmaxGrad<T>::ComputeInternal(OpKernelContext* ctx) const {
template <typename T>
Status SoftmaxGrad<T>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor* dY = ctx->Input<Tensor>(0);
const TensorShape& input_shape{dY->Shape()};
const Tensor* Y = ctx->Input<Tensor>(1);

2
orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc
View file

@ -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<BFloat16, is_log_softmax>(hipStream_t stream, const BFloat16* dY, \