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

[Training] Add bf16 support to GatherElementsGrad. (#20796)

Browse source 
### Description
Adding bf16 support to GatherElementsGrad.

---------

Co-authored-by: Adam Louly <adamlouly@microsoft.com@h100vm-ort.kxelwkzfzxguje5bxvwxxs135a.gvxx.internal.cloudapp.net>
This commit is contained in:
Adam Louly 2024-05-24 15:55:14 -07:00 committed by GitHub
parent 76e1a06986
commit ed8275883a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 56 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu
View file

@ -344,6 +344,7 @@ Status GatherElementsGradNonDeterministicImpl(cudaStream_t stream, const TIndex*
GATHER_ELEMENTS_GRAD_SPECIALIZED_TINDEX_IMPL(T, int32_t) \
GATHER_ELEMENTS_GRAD_SPECIALIZED_TINDEX_IMPL(T, int64_t)
GATHER_ELEMENTS_GRAD_SPECIALIZED_SCATTER_ADD_IMPL(BFloat16);
GATHER_ELEMENTS_GRAD_SPECIALIZED_SCATTER_ADD_IMPL(half)
GATHER_ELEMENTS_GRAD_SPECIALIZED_SCATTER_ADD_IMPL(float)
GATHER_ELEMENTS_GRAD_SPECIALIZED_SCATTER_ADD_IMPL(double)

13
onnxruntime/test/common/tensor_op_test_utils.h
View file

@ -120,6 +120,19 @@ inline std::vector<MLFloat16> ValueRange<MLFloat16>(size_t count, MLFloat16 star
return result;
}
template <>
inline std::vector<BFloat16> ValueRange<BFloat16>(size_t count, BFloat16 start, BFloat16 step) {
std::vector<BFloat16> result;
result.reserve(count);
float curr = start.ToFloat();
float f_step = step.ToFloat();
for (size_t i = 0; i < count; ++i) {
result.emplace_back(BFloat16(curr));
curr += f_step;
}
return result;
}
inline std::pair<float, float> MeanStdev(gsl::span<const float> v) {
float sum = std::accumulate(v.begin(), v.end(), 0.0f);
float mean = sum / v.size();

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

@ -1218,7 +1218,7 @@ void RegisterTrainingOpSchemas() {
"Constrain input shape to integer tensors.")
.TypeConstraint(
"T",
{"tensor(float16)", "tensor(float)", "tensor(double)"},
{"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Input and output types can be of any tensor type.")
.TypeConstraint(
"Tind",

34
orttraining/orttraining/test/training_ops/cuda/gather_elements_grad_test.cc
View file

@ -12,6 +12,8 @@
#include "test/providers/kernel_compute_test_utils.h"
#endif
#include "test/common/cuda_op_test_utils.h"
namespace onnxruntime {
namespace cuda {
namespace test {
@ -28,6 +30,11 @@ void Add<MLFloat16>(MLFloat16* a, const MLFloat16* b) {
*a = MLFloat16((*a).ToFloat() + (*b).ToFloat());
}
template <>
void Add<BFloat16>(BFloat16* a, const BFloat16* b) {
*a = BFloat16((*a).ToFloat() + (*b).ToFloat());
}
template <typename T, typename TIndex>
void GetData(const std::vector<int64_t>& input_dims, const std::vector<int64_t>& indices_dims,
const std::vector<int64_t>& indices_strides, int64_t axis, std::vector<T>& dY_data,
@ -186,6 +193,21 @@ TEST(GatherElementsGrad, double) { RunTestWrapper<double>(); }
TEST(GatherElementsGrad, MLFloat16) { RunTestWrapper<MLFloat16>(); }
#if defined(USE_CUDA) || defined(USE_ROCM)
TEST(GatherElementsGrad, BFloat16) {
#ifdef USE_CUDA
int min_cuda_architecture = 530;
if (!onnxruntime::test::HasCudaEnvironment(min_cuda_architecture)) {
LOGS_DEFAULT(WARNING) << "Hardware does not support BFP16";
return;
}
#endif
RunTestWrapper<BFloat16>();
}
#endif
TEST(GatherElementsGrad, IndicesUpdatesDontMatch) {
onnxruntime::test::OpTester test("GatherElementsGrad", 1, kMSDomain);
test.AddAttribute<int64_t>("axis", 1);
@ -203,6 +225,18 @@ TEST(GatherElementsGrad, Strided_float) { RunKernelComputeTestWrapper<float>();
TEST(GatherElementsGrad, Strided_double) { RunKernelComputeTestWrapper<double>(); }
TEST(GatherElementsGrad, Strided_MLFloat16) { RunKernelComputeTestWrapper<MLFloat16>(); }
TEST(GatherElementsGrad, Strided_BFloat16) {
#ifdef USE_CUDA
int min_cuda_architecture = 530;
if (!onnxruntime::test::HasCudaEnvironment(min_cuda_architecture)) {
LOGS_DEFAULT(WARNING) << "Hardware does not support BFP16";
return;
}
RunKernelComputeTestWrapper<BFloat16>();
#endif
}
#endif
} // namespace test

12
orttraining/orttraining/training_ops/cuda/tensor/gather_elements_grad.cc
View file

@ -21,7 +21,10 @@ namespace cuda {
ONNX_OPERATOR_KERNEL_EX(GatherElementsGrad, kMSDomain, 1, kCudaExecutionProvider,
CREATE_GATHER_ELEMENTS_GRAD_KERNEL_DEF
.InputMemoryType(OrtMemTypeCPUInput, 1) // 'GatherElements' data shape needs to be on CPU
.TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes())
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<MLFloat16>(),
DataTypeImpl::GetTensorType<BFloat16>()})
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("Tind", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
@ -94,9 +97,8 @@ Status GatherElementsGrad::ComputeInternal(OpKernelContext* context) const {
#endif
CoalesceDimensions(data_shape_vec, indices_shape_vec, p_indices_strides_vec, axis, args);
// Use element size instead of concrete types so we can specialize less template functions to reduce binary size.
int dtype = GetElementType(dY->DataType()->Size());
// GatherElementsGrad supports half, float and double only for now, it's element size will not but INT8.
const int dtype = dY->GetElementType();
// GatherElementsGrad supports half, bfloat16, float and double only for now, it's element size will not but INT8.
if (dtype == ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED || dtype == ONNX_NAMESPACE::TensorProto_DataType_INT8) {
ORT_THROW("Unsupported element size by the GatherElementsGrad CUDA kernel");
}
@ -108,7 +110,7 @@ Status GatherElementsGrad::ComputeInternal(OpKernelContext* context) const {
});
}
utils::MLTypeCallDispatcher<MLFloat16, float, double> t_disp(dtype);
utils::MLTypeCallDispatcher<MLFloat16, BFloat16, float, double> t_disp(dtype);
return t_disp.InvokeRet<Status, ComputeImpl>(Stream(context), dY->DataRaw(), indices_tensor->DataRaw(),
dX->MutableDataRaw(),
indices_tensor->DataType()->Size(), args);