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

Promote BiasDropout from orttraining to onnxruntime (#7116)

Browse source 
* Promote BiasDropout from orttraining to onnxruntime

Co-authored-by: Sherlock Huang <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Sherlock 2021-03-24 20:42:42 -07:00 committed by GitHub
parent cd67f12add
commit 1c8d874412
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 523 additions and 452 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -34,6 +34,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMulConj);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasSoftmax);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasDropout);
// These ops were experimental ops in onnx domain which have been removed now. We add them here as
// contrib ops to maintain backward compatibility
@ -161,6 +162,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16_float, SimplifiedLayerNormalization)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Inverse)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasSoftmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasDropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int8_t_MLFloat16, QuantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, uint8_t_MLFloat16, QuantizeLinear)>,

85
orttraining/orttraining/training_ops/cuda/nn/dropout.cc → onnxruntime/contrib_ops/cuda/math/bias_dropout.cc
View file

@ -1,89 +1,15 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/framework/random_seed.h"
#include "orttraining/training_ops/cuda/nn/dropout.h"
#include "core/providers/cuda/cuda_common.h"
#include "contrib_ops/cuda/math/bias_dropout.h"
#include "core/providers/cuda/nn/dropout.h"
#include "core/providers/common.h"
namespace onnxruntime {
namespace contrib {
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_GRADIENT_KERNEL(OpName) \
ONNX_OPERATOR_KERNEL_EX( \
OpName, \
kMSDomain, \
1, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) \
.TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<bool>()) \
.InputMemoryType<OrtMemTypeCPUInput>(2), \
DropoutGrad);
REGISTER_GRADIENT_KERNEL(DropoutGrad)
template <typename T>
struct DropoutGradComputeImpl {
void operator()(cudaStream_t stream,
const int64_t N,
const Tensor& dY,
const bool* mask_data,
const float ratio_data,
Tensor& dX) const {
typedef typename ToCudaType<T>::MappedType CudaT;
const CudaT* dY_data = reinterpret_cast<const CudaT*>(dY.template Data<T>());
CudaT* dX_data = reinterpret_cast<CudaT*>(dX.template MutableData<T>());
DropoutGradientKernelImpl<CudaT>(stream, N, dY_data, mask_data, ratio_data, dX_data);
}
};
// REVIEW(codemzs): Common out this structure because it is also used in Dropout forward op.
template <typename T>
struct GetRatioDataImpl {
void operator()(const Tensor* ratio, float& ratio_data) const {
ratio_data = static_cast<float>(*(ratio->template Data<T>()));
ORT_ENFORCE(ratio_data >= 0.0f && ratio_data < 1.0f, "ratio_data is outside range [0, 1)");
}
};
Status DropoutGrad::ComputeInternal(OpKernelContext* context) const {
auto dY = context->Input<Tensor>(0);
const TensorShape& shape = dY->Shape();
const int64_t N = shape.Size();
auto mask = context->Input<Tensor>(1);
ORT_ENFORCE(mask->Shape().Size() == N);
const bool* mask_data = mask->template Data<bool>();
//Get the ratio_data
float ratio_data = default_ratio_;
auto ratio = context->Input<Tensor>(2);
if (ratio) {
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(ratio->GetElementType());
t_disp.Invoke<GetRatioDataImpl>(ratio, ratio_data);
}
auto dX = context->Output(0, shape);
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(dY->GetElementType());
t_disp.Invoke<DropoutGradComputeImpl>(Stream(), N, *dY, mask_data, ratio_data, *dX);
return Status::OK();
}
ONNX_OPERATOR_KERNEL_EX(
BiasDropout,
@ -192,4 +118,5 @@ Status BiasDropout::ComputeInternal(OpKernelContext* context) const {
}
} // namespace cuda
} // namespace onnxruntime
} // namespace contrib
} // namespace onnxruntime

31
orttraining/orttraining/training_ops/cuda/nn/dropout.h → onnxruntime/contrib_ops/cuda/math/bias_dropout.h
View file

@ -3,22 +3,30 @@
#pragma once
#include "gsl/gsl"
#include "core/providers/cuda/cuda_kernel.h"
#include "orttraining/training_ops/cuda/nn/dropout_impl.h"
#include "core/providers/cuda/cuda_common.h"
#include "core/framework/random_generator.h"
using namespace onnxruntime::cuda;
namespace onnxruntime {
namespace contrib {
namespace cuda {
class DropoutGrad final : public CudaKernel {
public:
DropoutGrad(const OpKernelInfo& info) : CudaKernel(info) {
}
Status ComputeInternal(OpKernelContext* context) const override;
private:
static constexpr float default_ratio_ = 0.5f;
};
template <typename T>
void BiasDropoutKernelImpl(
const cudaDeviceProp& prop,
cudaStream_t stream,
const int64_t N,
const fast_divmod fdm_dim,
const float ratio,
PhiloxGenerator& generator,
const T* X_data,
const T* bias_data,
const T* residual_data,
T* Y_data,
bool* mask_data);
class BiasDropout final : public CudaKernel {
public:
@ -37,4 +45,5 @@ class BiasDropout final : public CudaKernel {
};
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

58
orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu → onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu
View file

@ -17,66 +17,15 @@
/* Modifications Copyright (c) Microsoft. */
#include "core/providers/cuda/cu_inc/common.cuh"
#include "orttraining/training_ops/cuda/nn/dropout_impl.h"
#include "contrib_ops/cuda/math/bias_dropout.h"
#include <curand_kernel.h>
#include <algorithm>
namespace onnxruntime {
namespace contrib {
namespace cuda {
template <typename T, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void DropoutGradientKernel(
const int64_t N,
const T* dY_data,
const bool* mask_data,
const float scale,
T* dX_data) {
CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
dX_data[id] = T(float(dY_data[id]) * mask_data[id] * scale);
id += NumThreadsPerBlock;
}
}
}
template <typename T>
void DropoutGradientKernelImpl(
cudaStream_t stream,
const int64_t N,
const T* dY_data,
const bool* mask_data,
const float ratio,
T* dX_data) {
if (ratio == 0.0f) {
if (dY_data != dX_data) {
CUDA_CALL_THROW(cudaMemcpyAsync(dX_data, dY_data, N * sizeof(T), cudaMemcpyDeviceToDevice, stream));
}
} else {
const float scale = 1.f / (1.f - ratio);
const int blocksPerGrid = static_cast<int>(CeilDiv(N, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
DropoutGradientKernel<T, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(N, dY_data, mask_data, scale, dX_data);
}
}
#define SPECIALIZED_DROPOUT_GRAD_IMPL(T) \
template void DropoutGradientKernelImpl( \
cudaStream_t stream, \
const int64_t N, \
const T* dY_data, \
const bool* mask_data, \
const float scale, \
T* dX_data);
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;
template <typename T, bool has_residual>
@ -180,4 +129,5 @@ SPECIALIZED_BIAS_DROPOUT_IMPL(nv_bfloat16)
#endif
} // namespace cuda
} // namespace contrib {
} // namespace onnxruntime

2
onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc
View file

@ -31,6 +31,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, ComplexMulConj);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BiasSoftmax);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BiasDropout);
// These ops were experimental ops in onnx domain which have been removed now. We add them here as
// contrib ops to maintain backward compatibility
@ -137,6 +138,7 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16_float, LayerNormalization)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, Inverse)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BiasSoftmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BiasDropout)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, int8_t_MLFloat16, QuantizeLinear)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, uint8_t_MLFloat16, QuantizeLinear)>,

53
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -275,7 +275,6 @@ void FusedMatMulShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
updateOutputShape(ctx, 0, resultShape);
}
void AttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int past_input_index) {
// Type inference
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 2, 0);
@ -441,7 +440,7 @@ and present state are optional. Present state could appear in output even when p
AttentionTypeAndShapeInference(ctx, past_input_index);
});
static const char* Longformer_Attention_doc = R"DOC(
static const char* Longformer_Attention_doc = R"DOC(
Longformer Self Attention with a local context and a global context. Tokens attend locally: Each token
attends to its W previous tokens and W succeding tokens with W being the window length. A selected few tokens
attend globally to all other tokens.
@ -2377,6 +2376,56 @@ It's an extension of Gelu. It takes the sum of input A and bias input B as the i
"Constrain input and output types to float tensors.")
.TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput);
ONNX_CONTRIB_OPERATOR_SCHEMA(BiasDropout)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetDoc(
"output, dropout_mask = Dropout(data + bias, ratio) + residual, "
"Intended to specialize the dropout pattern commonly found in transformer models.")
.Attr("seed", "(Optional) Seed to the random generator, if not specified we will auto generate one.", AttributeProto::INT, OPTIONAL_VALUE)
.AllowUncheckedAttributes()
.Input(0, "data", "The input data as Tensor.", "T")
.Input(1, "bias", "The bias input, a vector with the same shape as last dim of data", "T")
.Input(2, "residual", "The residual input, must have the same shape as data", "T", OpSchema::Optional)
.Input(3, "ratio",
"The ratio of random dropout, with value in [0, 1). If this input was not set, "
"or if it was set to 0, the output would be a simple copy of the input. "
"If it's non-zero, output will be a random dropout of input, which is typically "
"the case during training.",
"T1",
OpSchema::Optional)
.Input(4, "training_mode",
"If set to true then it indicates dropout is being used for "
"training. It is an optional value hence unless specified explicitly, it is false. "
"If it is false, ratio is ignored and the operation mimics inference mode where nothing "
"will be dropped from the input data and if mask is requested as output it will contain "
"all ones.",
"T2",
OpSchema::Optional)
.Output(0, "output", "The output.", "T")
.Output(1, "mask", "The output mask of dropout.", "T2", OpSchema::Optional)
.TypeConstraint(
"T",
{"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(bfloat16)"},
"Constrain input 'ratio' types to float tensors.")
.TypeConstraint(
"T2",
{"tensor(bool)"},
"Constrain output 'mask' types to boolean tensors.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateShapeAndTypeFromFirstInput(ctx);
if (ctx.getNumOutputs() == 2) {
updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto::BOOL);
if (hasNInputShapes(ctx, 1)) {
propagateShapeFromInputToOutput(ctx, 0, 1);
}
}
});
// Register the NCHWc schemas if supported by the platform.
if (MlasNchwcGetBlockSize() > 1) {
RegisterNchwcSchemas();

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

@ -6,15 +6,6 @@
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,
@ -41,5 +32,59 @@ ONNX_OPERATOR_KERNEL_EX(
.InputMemoryType<OrtMemTypeCPUInput>(2),
Dropout);
Status Dropout::ComputeInternal(OpKernelContext* context) const {
//Get X_data
const Tensor* X = context->Input<Tensor>(0);
if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available.");
const TensorShape& shape = X->Shape();
const int64_t N = shape.Size();
//Get Y_data
auto Y = context->Output(0, shape);
//Get mask_data
auto mask = context->Output(1, shape);
ORT_ENFORCE(!mask || mask->Shape().Size() == N);
//Get the ratio_data
float ratio_data = default_ratio_;
auto ratio = context->Input<Tensor>(1);
if (ratio) {
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(ratio->GetElementType());
t_disp.Invoke<GetRatioDataImpl>(ratio, ratio_data);
}
const Tensor* training_mode = context->Input<Tensor>(2);
//Check for inference mode.
if ((0 == ratio_data) ||(training_mode == nullptr || *(training_mode->Data<bool>()) == false)) {
const void* X_data = X->DataRaw();
void* Y_data = Y->MutableDataRaw();
if (Y_data != X_data) {
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y_data, X_data, X->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream()));
}
// If mask is requested, return all 1s.
if (mask != nullptr) {
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(mask->MutableData<bool>(), true, N * sizeof(bool), Stream()));
}
return Status::OK();
}
IAllocatorUniquePtr<bool> temp_mask_buffer{}; // buffer to use if mask is not provided
bool* const mask_data = [this, N, mask, &temp_mask_buffer]() {
if (mask) return mask->MutableData<bool>();
temp_mask_buffer = GetScratchBuffer<bool>(N);
return temp_mask_buffer.get();
}();
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(X->GetElementType());
t_disp.Invoke<DropoutComputeImpl>(GetDeviceProp(), Stream(), N, ratio_data, generator, *X, *Y, mask_data);
return Status::OK();
}
} // namespace cuda
} // namespace onnxruntime

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

@ -6,11 +6,22 @@
#include "core/providers/cuda/cuda_kernel.h"
#include "core/providers/cuda/nn/dropout_impl.h"
#include "core/providers/common.h"
#include "core/framework/random_seed.h"
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
template <typename T>
struct GetRatioDataImpl {
void operator()(const Tensor* ratio, float& ratio_data) const {
@ -54,67 +65,5 @@ class Dropout final : public CudaKernel {
static constexpr float default_ratio_ = 0.5f;
};
Status Dropout::ComputeInternal(OpKernelContext* context) const {
//Get X_data
const Tensor* X = context->Input<Tensor>(0);
if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available.");
const TensorShape& shape = X->Shape();
const int64_t N = shape.Size();
//Get Y_data
auto Y = context->Output(0, shape);
//Get mask_data
auto mask = context->Output(1, shape);
ORT_ENFORCE(!mask || mask->Shape().Size() == N);
//Get the ratio_data
float ratio_data = default_ratio_;
auto ratio = context->Input<Tensor>(1);
if (ratio) {
utils::MLTypeCallDispatcher<float, MLFloat16, double> t_disp(ratio->GetElementType());
t_disp.Invoke<GetRatioDataImpl>(ratio, ratio_data);
}
const Tensor* training_mode = context->Input<Tensor>(2);
//Check for inference mode.
if ((0 == ratio_data) ||(training_mode == nullptr || *(training_mode->Data<bool>()) == false)) {
const void* X_data = X->DataRaw();
void* Y_data = Y->MutableDataRaw();
if (Y_data != X_data) {
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y_data, X_data, X->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream()));
}
// If mask is requested, return all 1s.
if (mask != nullptr) {
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(mask->MutableData<bool>(), true, N * sizeof(bool), Stream()));
}
return Status::OK();
}
IAllocatorUniquePtr<bool> temp_mask_buffer{}; // buffer to use if mask is not provided
bool* const mask_data = [this, N, mask, &temp_mask_buffer]() {
if (mask) return mask->MutableData<bool>();
temp_mask_buffer = GetScratchBuffer<bool>(N);
return temp_mask_buffer.get();
}();
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
using SupportedTypes = onnxruntime::TypeList<
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
float, MLFloat16, double, BFloat16
#else
float, MLFloat16, double
#endif
>;
utils::MLTypeCallDispatcherFromTypeList<SupportedTypes> t_disp(X->GetElementType());
t_disp.Invoke<DropoutComputeImpl>(GetDeviceProp(), Stream(), N, ratio_data, generator, *X, *Y, mask_data);
return Status::OK();
}
} // namespace cuda
} // namespace onnxruntime

171
onnxruntime/test/contrib_ops/bias_dropout_op_test.cc Normal file
View file

@ -0,0 +1,171 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#ifdef _MSC_VER
#pragma warning(disable : 4389)
#endif
#include <algorithm>
#include <memory>
#include <numeric>
#include <random>
#include "gtest/gtest.h"
#include "test/common/tensor_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
#include "test/util/include/default_providers.h"
namespace onnxruntime {
namespace contrib {
namespace test {
using namespace onnxruntime::test;
enum TrainingMode { TrainingFalse, TrainingTrue, NoTraining };
// BiasDropout kernel is only implemented for CUDA
#ifdef USE_CUDA
namespace {
void RunBiasDropoutTest(const bool use_mask, const std::vector<int64_t>& input_shape, float ratio = -1.0f,
TrainingMode training_mode = TrainingTrue, bool use_float16_ratio = false, bool has_residual = true) {
OpTester t{"BiasDropout", 1, kMSDomain};
const int64_t seed = 42;
t.AddAttribute("seed", seed);
const auto input_size = std::accumulate(
input_shape.begin(), input_shape.end(), static_cast<int64_t>(1), std::multiplies<>{});
const std::vector<float> input = ValueRange(input_size, 1.0f, 1.0f);
t.AddInput("data", input_shape, input);
std::vector<int64_t> bias_shape{input_shape.back()};
const auto bias_size = input_shape.back();
const std::vector<float> bias = ValueRange(bias_size, 2.0f, 1.0f);
t.AddInput("bias", bias_shape, bias);
float residual_value = 0.0f;
if (has_residual) {
residual_value = 1.0f;
const auto residual_size = input_size;
const std::vector<float> residual(residual_size, residual_value);
t.AddInput("residual", input_shape, residual);
} else {
t.AddMissingOptionalInput<float>();
}
if (ratio == -1.0f) {
if (use_float16_ratio) {
t.AddMissingOptionalInput<MLFloat16>();
} else {
t.AddMissingOptionalInput<float>();
}
// set ratio to default value
ratio = 0.5f;
} else {
if (use_float16_ratio) {
t.AddInput("ratio", {}, {MLFloat16(math::floatToHalf(ratio))});
} else {
t.AddInput("ratio", {}, {ratio});
}
}
if (training_mode != NoTraining) {
if (training_mode == TrainingTrue) {
t.AddInput("training_mode", {}, {true});
} else {
t.AddInput("training_mode", {}, {false});
}
}
t.AddOutput<float>("output", input_shape, input); // we'll do our own output verification
std::unique_ptr<bool[]> mask_buffer{};
if (use_mask) {
mask_buffer = onnxruntime::make_unique<bool[]>(input_size);
t.AddOutput<bool>("mask", input_shape, mask_buffer.get(), input_size);
} else {
t.AddMissingOptionalOutput<bool>();
}
auto output_verifier = [&](const std::vector<OrtValue>& fetches, const std::string& provider_type) {
ASSERT_GE(fetches.size(), 1);
const auto& output_tensor = FetchTensor(fetches[0]);
auto output_span = output_tensor.DataAsSpan<float>();
const auto num_dropped_values = std::count(output_span.begin(), output_span.end(), residual_value);
if (ratio == 1.0f) {
ASSERT_EQ(num_dropped_values, static_cast<size_t>(output_span.size())) << "provider: " << provider_type;
} else {
ASSERT_NEAR(static_cast<float>(num_dropped_values) / static_cast<size_t>(output_span.size()), training_mode == TrainingTrue ? ratio : 0.0f, 0.1f)
<< "provider: " << provider_type;
for (decltype(output_span.size()) i = 0; i < output_span.size(); ++i) {
if (output_span[i] == residual_value) continue;
const auto expected_value = (bias[i % bias_size] + i + 1.0f) / (1 - (training_mode == TrainingTrue ? ratio : 0.0f)) + residual_value;
ASSERT_NEAR(output_span[i], expected_value, 0.01f)
<< "unexpected output value at index " << i << ", provider: " << provider_type;
}
}
if (use_mask) {
ASSERT_GE(fetches.size(), 2);
const auto& mask_tensor = FetchTensor(fetches[1]);
auto mask_span = mask_tensor.DataAsSpan<bool>();
ASSERT_EQ(mask_span.size(), output_span.size()) << "provider: " << provider_type;
const auto num_mask_zeros = std::count(mask_span.begin(), mask_span.end(), false);
ASSERT_EQ(num_dropped_values, num_mask_zeros) << "provider: " << provider_type;
for (decltype(mask_span.size()) i = 0; i < mask_span.size(); ++i) {
ASSERT_TRUE(
(mask_span[i] && output_span[i] != residual_value) || (!mask_span[i] && output_span[i] == residual_value))
<< "output and mask mismatch at index " << i << ", output[i]: " << output_span[i]
<< ", mask[i]: " << mask_span[i] << ", provider: " << provider_type;
}
}
};
t.SetCustomOutputVerifier(output_verifier);
t.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
} // namespace
TEST(BiasDropoutTest, Basic) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f);
}
TEST(BiasDropoutTest, BasicTrainingFalse) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f, TrainingFalse);
}
TEST(BiasDropoutTest, BasicNoTraining) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f, NoTraining);
}
TEST(BiasDropoutTest, BasicWithoutResidual) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f, TrainingTrue, false, false);
}
TEST(BiasDropoutTest, Mask) {
RunBiasDropoutTest(true, {3, 5, 768}, 0.25f);
}
TEST(BiasDropoutTest, MaskTrainingFalse) {
RunBiasDropoutTest(true, {3, 5, 768}, 0.25f, TrainingFalse);
}
TEST(BiasDropoutTest, MaskNoTraining) {
RunBiasDropoutTest(true, {3, 5, 768}, 0.25f, NoTraining);
}
TEST(BiasDropoutTest, RatioLimit) {
RunBiasDropoutTest(true, {4, 8, 1024}, 0.0f, TrainingFalse);
}
TEST(BiasDropoutTest, EmptyRatio) {
RunBiasDropoutTest(true, {2, 7, 1024});
}
#endif
} // namespace test
} // namespace contrib
} // namespace onnxruntime

26
onnxruntime/test/optimizer/graph_transform_test.cc
View file

@ -17,6 +17,7 @@
#include "core/optimizer/attention_fusion.h"
#include "core/optimizer/bias_gelu_fusion.h"
#include "core/optimizer/bias_softmax_fusion.h"
#include "core/optimizer/bias_dropout_fusion.h"
#include "core/optimizer/computation_reduction.h"
#include "core/optimizer/cast_elimination.h"
#include "core/optimizer/common_subexpression_elimination.h"
@ -2690,6 +2691,31 @@ TEST_F(GraphTransformationTests, BiasSoftmaxFusionTest_NoLeadingOnes) {
tester.TestFusionOccurs(0);
}
static void TestBiasDropoutFusion(const PathString& file_path, const logging::Logger& logger, const int add_count = 0) {
std::shared_ptr<Model> p_model;
ASSERT_TRUE(Model::Load(file_path, p_model, nullptr, logger).IsOK());
Graph& graph = p_model->MainGraph();
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<BiasDropoutFusion>(), TransformerLevel::Level2);
auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, logger);
ASSERT_STATUS_OK(ret);
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_EQ(op_to_count["Add"], add_count);
ASSERT_EQ(op_to_count["Dropout"], 0);
ASSERT_EQ(op_to_count["com.microsoft.BiasDropout"], 1);
}
TEST_F(GraphTransformationTests, BiasDropoutFusionTest) {
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_fusion1.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_fusion2.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_residual_fusion1.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_residual_fusion2.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_residual_fusion_mismatch.onnx", *logger_, 1);
}
TEST_F(GraphTransformationTests, LayerNormFusionTest) {
auto model_uri = MODEL_FOLDER "fusion/layer_norm.onnx";
std::shared_ptr<Model> p_model;

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

@ -1120,54 +1120,6 @@ Example 4:
})
.SetDoc(R"DOC(SoftmaxCrossEntropyLossGrad)DOC");
ONNX_CONTRIB_OPERATOR_SCHEMA(BiasDropout)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetDoc("BiasDropout")
.Attr("seed", "(Optional) Seed to the random generator, if not specified we will auto generate one.", AttributeProto::INT, OPTIONAL_VALUE)
.AllowUncheckedAttributes()
.Input(0, "data", "The input data as Tensor.", "T")
.Input(1, "bias", "The bias input, a vector with the same shape as last dim of data", "T")
.Input(2, "residual", "The residual input, must have the same shape as data", "T", OpSchema::Optional)
.Input(3, "ratio",
"The ratio of random dropout, with value in [0, 1). If this input was not set, "
"or if it was set to 0, the output would be a simple copy of the input. "
"If it's non-zero, output will be a random dropout of input, which is typically "
"the case during training.",
"T1",
OpSchema::Optional)
.Input(4, "training_mode",
"If set to true then it indicates dropout is being used for "
"training. It is an optional value hence unless specified explicitly, it is false. "
"If it is false, ratio is ignored and the operation mimics inference mode where nothing "
"will be dropped from the input data and if mask is requested as output it will contain "
"all ones.",
"T2",
OpSchema::Optional)
.Output(0, "output", "The output.", "T")
.Output(1, "mask", "The output mask of dropout.", "T2", OpSchema::Optional)
.TypeConstraint(
"T",
{"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(bfloat16)"},
"Constrain input 'ratio' types to float tensors.")
.TypeConstraint(
"T2",
{"tensor(bool)"},
"Constrain output 'mask' types to boolean tensors.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateShapeAndTypeFromFirstInput(ctx);
if (ctx.getNumOutputs() == 2) {
updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto::BOOL);
if (hasNInputShapes(ctx, 1)) {
propagateShapeFromInputToOutput(ctx, 0, 1);
}
}
});
ONNX_CONTRIB_OPERATOR_SCHEMA(ReduceSumTraining)
.SetDomain(kMSDomain)
.SinceVersion(1)

26
orttraining/orttraining/test/optimizer/graph_transform_test.cc
View file

@ -13,7 +13,6 @@
#include "core/optimizer/bias_gelu_fusion.h"
#include "core/optimizer/gelu_fusion.h"
#include "core/optimizer/dropout_elimination.h"
#include "core/optimizer/bias_dropout_fusion.h"
#include "orttraining/core/optimizer/gist_encode_decode.h"
#include "orttraining/core/optimizer/nonzero_shape_setter.h"
#include "orttraining/core/optimizer/megatron_transformer.h"
@ -74,31 +73,6 @@ TEST_F(GraphTransformationTests, GistEncodeDecode) {
ASSERT_TRUE(op_to_count["GistBinarizeEncoder"] == op_to_count["GistBinarizeEncoder"]);
}
static void TestBiasDropoutFusion(const PathString& file_path, const logging::Logger& logger, const int add_count = 0) {
std::shared_ptr<Model> p_model;
ASSERT_TRUE(Model::Load(file_path, p_model, nullptr, logger).IsOK());
Graph& graph = p_model->MainGraph();
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<BiasDropoutFusion>(), TransformerLevel::Level2);
auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, logger);
ASSERT_STATUS_OK(ret);
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_EQ(op_to_count["Add"], add_count);
ASSERT_EQ(op_to_count["Dropout"], 0);
ASSERT_EQ(op_to_count["com.microsoft.BiasDropout"], 1);
}
TEST_F(GraphTransformationTests, BiasDropoutFusionTest) {
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_fusion1.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_fusion2.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_residual_fusion1.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_residual_fusion2.onnx", *logger_);
TestBiasDropoutFusion(MODEL_FOLDER "fusion/bias_dropout_residual_fusion_mismatch.onnx", *logger_, 1);
}
Node* GetNodeByName(Graph& graph, std::string node_name) {
GraphViewer graph_viewer(graph);
const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder();

141
orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc
View file

@ -151,147 +151,6 @@ TEST(DropoutTest, EmptyRatio) {
RunDropoutTest(true, {1000});
}
// BiasDropout kernel is only implemented for CUDA
#ifdef USE_CUDA
namespace {
void RunBiasDropoutTest(const bool use_mask, const std::vector<int64_t>& input_shape, float ratio = -1.0f,
TrainingMode training_mode = TrainingTrue, bool use_float16_ratio = false, bool has_residual = true) {
OpTester t{"BiasDropout", 1, kMSDomain};
const int64_t seed = 42;
t.AddAttribute("seed", seed);
const auto input_size = std::accumulate(
input_shape.begin(), input_shape.end(), static_cast<int64_t>(1), std::multiplies<>{});
const std::vector<float> input = ValueRange(input_size, 1.0f, 1.0f);
t.AddInput("data", input_shape, input);
std::vector<int64_t> bias_shape{input_shape.back()};
const auto bias_size = input_shape.back();
const std::vector<float> bias = ValueRange(bias_size, 2.0f, 1.0f);
t.AddInput("bias", bias_shape, bias);
float residual_value = 0.0f;
if (has_residual) {
residual_value = 1.0f;
const auto residual_size = input_size;
const std::vector<float> residual(residual_size, residual_value);
t.AddInput("residual", input_shape, residual);
} else {
t.AddMissingOptionalInput<float>();
}
if (ratio == -1.0f) {
if (use_float16_ratio) {
t.AddMissingOptionalInput<MLFloat16>();
} else {
t.AddMissingOptionalInput<float>();
}
// set ratio to default value
ratio = 0.5f;
} else {
if (use_float16_ratio) {
t.AddInput("ratio", {}, {MLFloat16(math::floatToHalf(ratio))});
} else {
t.AddInput("ratio", {}, {ratio});
}
}
if (training_mode != NoTraining) {
if (training_mode == TrainingTrue) {
t.AddInput("training_mode", {}, {true});
} else {
t.AddInput("training_mode", {}, {false});
}
}
t.AddOutput<float>("output", input_shape, input); // we'll do our own output verification
std::unique_ptr<bool[]> mask_buffer{};
if (use_mask) {
mask_buffer = onnxruntime::make_unique<bool[]>(input_size);
t.AddOutput<bool>("mask", input_shape, mask_buffer.get(), input_size);
} else {
t.AddMissingOptionalOutput<bool>();
}
auto output_verifier = [&](const std::vector<OrtValue>& fetches, const std::string& provider_type) {
ASSERT_GE(fetches.size(), 1);
const auto& output_tensor = FetchTensor(fetches[0]);
auto output_span = output_tensor.DataAsSpan<float>();
const auto num_dropped_values = std::count(output_span.begin(), output_span.end(), residual_value);
if (ratio == 1.0f) {
ASSERT_EQ(num_dropped_values, static_cast<size_t>(output_span.size())) << "provider: " << provider_type;
} else {
ASSERT_NEAR(static_cast<float>(num_dropped_values) / static_cast<size_t>(output_span.size()), training_mode == TrainingTrue ? ratio : 0.0f, 0.1f)
<< "provider: " << provider_type;
for (decltype(output_span.size()) i = 0; i < output_span.size(); ++i) {
if (output_span[i] == residual_value) continue;
const auto expected_value = (bias[i % bias_size] + i + 1.0f) / (1 - (training_mode == TrainingTrue ? ratio : 0.0f)) + residual_value;
ASSERT_NEAR(output_span[i], expected_value, 0.01f)
<< "unexpected output value at index " << i << ", provider: " << provider_type;
}
}
if (use_mask) {
ASSERT_GE(fetches.size(), 2);
const auto& mask_tensor = FetchTensor(fetches[1]);
auto mask_span = mask_tensor.DataAsSpan<bool>();
ASSERT_EQ(mask_span.size(), output_span.size()) << "provider: " << provider_type;
const auto num_mask_zeros = std::count(mask_span.begin(), mask_span.end(), false);
ASSERT_EQ(num_dropped_values, num_mask_zeros) << "provider: " << provider_type;
for (decltype(mask_span.size()) i = 0; i < mask_span.size(); ++i) {
ASSERT_TRUE(
(mask_span[i] && output_span[i] != residual_value) || (!mask_span[i] && output_span[i] == residual_value))
<< "output and mask mismatch at index " << i << ", output[i]: " << output_span[i]
<< ", mask[i]: " << mask_span[i] << ", provider: " << provider_type;
}
}
};
t.SetCustomOutputVerifier(output_verifier);
t.Run();
}
} // namespace
TEST(BiasDropoutTest, Basic) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f);
}
TEST(BiasDropoutTest, BasicTrainingFalse) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f, TrainingFalse);
}
TEST(BiasDropoutTest, BasicNoTraining) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f, NoTraining);
}
TEST(BiasDropoutTest, BasicWithoutResidual) {
RunBiasDropoutTest(false, {10, 10, 10}, 0.75f, TrainingTrue, false, false);
}
TEST(BiasDropoutTest, Mask) {
RunBiasDropoutTest(true, {3, 5, 768}, 0.25f);
}
TEST(BiasDropoutTest, MaskTrainingFalse) {
RunBiasDropoutTest(true, {3, 5, 768}, 0.25f, TrainingFalse);
}
TEST(BiasDropoutTest, MaskNoTraining) {
RunBiasDropoutTest(true, {3, 5, 768}, 0.25f, NoTraining);
}
TEST(BiasDropoutTest, RatioLimit) {
RunBiasDropoutTest(true, {4, 8, 1024}, 0.0f, TrainingFalse);
}
TEST(BiasDropoutTest, EmptyRatio) {
RunBiasDropoutTest(true, {2, 7, 1024});
}
#endif
namespace {
void RunDropoutGradTest(float ratio, const std::vector<int64_t>& input_dims, bool default_ratio = true) {
const auto input_shape = TensorShape(input_dims);

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

@ -65,7 +65,6 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, BatchNormalizationGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, BatchNormalizationGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, GatherGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasDropout);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, DropoutGrad);
// TODO: decprecate GatherND-1 after updating training models to opset-12
@ -214,8 +213,7 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ZeroGradient)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ZeroGradient)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasDropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, DropoutGrad)>,
// TODO: decprecate GatherND-1 after updating training models to opset-12

69
orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc Normal file
View file

@ -0,0 +1,69 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_ops/cuda/nn/dropout_grad.h"
#include "core/providers/cuda/nn/dropout.h"
#include "core/providers/cuda/cuda_common.h"
#include "core/providers/common.h"
namespace onnxruntime {
namespace cuda {
#define REGISTER_GRADIENT_KERNEL(OpName) \
ONNX_OPERATOR_KERNEL_EX( \
OpName, \
kMSDomain, \
1, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", ALL_IEEE_FLOAT_TENSOR_TYPES) \
.TypeConstraint("T1", ALL_IEEE_FLOAT_TENSOR_TYPES) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<bool>()) \
.InputMemoryType<OrtMemTypeCPUInput>(2), \
DropoutGrad);
REGISTER_GRADIENT_KERNEL(DropoutGrad)
template <typename T>
struct DropoutGradComputeImpl {
void operator()(cudaStream_t stream,
const int64_t N,
const Tensor& dY,
const bool* mask_data,
const float ratio_data,
Tensor& dX) const {
typedef typename ToCudaType<T>::MappedType CudaT;
const CudaT* dY_data = reinterpret_cast<const CudaT*>(dY.template Data<T>());
CudaT* dX_data = reinterpret_cast<CudaT*>(dX.template MutableData<T>());
DropoutGradientKernelImpl<CudaT>(stream, N, dY_data, mask_data, ratio_data, dX_data);
}
};
Status DropoutGrad::ComputeInternal(OpKernelContext* context) const {
auto dY = context->Input<Tensor>(0);
const TensorShape& shape = dY->Shape();
const int64_t N = shape.Size();
auto mask = context->Input<Tensor>(1);
ORT_ENFORCE(mask->Shape().Size() == N);
const bool* mask_data = mask->template Data<bool>();
//Get the ratio_data
float ratio_data = default_ratio_;
auto ratio = context->Input<Tensor>(2);
if (ratio) {
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(ratio->GetElementType());
t_disp.Invoke<GetRatioDataImpl>(ratio, ratio_data);
}
auto dX = context->Output(0, shape);
utils::MLTypeCallDispatcher<ALL_IEEE_FLOAT_DATA_TYPES> t_disp(dY->GetElementType());
t_disp.Invoke<DropoutGradComputeImpl>(Stream(), N, *dY, mask_data, ratio_data, *dX);
return Status::OK();
}
} // namespace cuda
} // namespace onnxruntime

24
orttraining/orttraining/training_ops/cuda/nn/dropout_grad.h Normal file
View file

@ -0,0 +1,24 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/cuda/cuda_kernel.h"
#include "orttraining/training_ops/cuda/nn/dropout_grad_impl.h"
namespace onnxruntime {
namespace cuda {
class DropoutGrad final : public CudaKernel {
public:
DropoutGrad(const OpKernelInfo& info) : CudaKernel(info) {
}
Status ComputeInternal(OpKernelContext* context) const override;
private:
static constexpr float default_ratio_ = 0.5f;
};
} // namespace cuda
} // namespace onnxruntime

81
orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu Normal file
View file

@ -0,0 +1,81 @@
/**
* Copyright (c) 2016-present, Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/* Modifications Copyright (c) Microsoft. */
#include "core/providers/cuda/cu_inc/common.cuh"
#include "orttraining/training_ops/cuda/nn/dropout_grad_impl.h"
#include <curand_kernel.h>
#include <algorithm>
namespace onnxruntime {
namespace cuda {
template <typename T, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void DropoutGradientKernel(
const int64_t N,
const T* dY_data,
const bool* mask_data,
const float scale,
T* dX_data) {
CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
dX_data[id] = T(float(dY_data[id]) * mask_data[id] * scale);
id += NumThreadsPerBlock;
}
}
}
template <typename T>
void DropoutGradientKernelImpl(
cudaStream_t stream,
const int64_t N,
const T* dY_data,
const bool* mask_data,
const float ratio,
T* dX_data) {
if (ratio == 0.0f) {
if (dY_data != dX_data) {
CUDA_CALL_THROW(cudaMemcpyAsync(dX_data, dY_data, N * sizeof(T), cudaMemcpyDeviceToDevice, stream));
}
} else {
const float scale = 1.f / (1.f - ratio);
const int blocksPerGrid = static_cast<int>(CeilDiv(N, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
DropoutGradientKernel<T, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(N, dY_data, mask_data, scale, dX_data);
}
}
#define SPECIALIZED_DROPOUT_GRAD_IMPL(T) \
template void DropoutGradientKernelImpl( \
cudaStream_t stream, \
const int64_t N, \
const T* dY_data, \
const bool* mask_data, \
const float scale, \
T* dX_data);
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
} // namespace cuda
} // namespace onnxruntime

14
orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h → orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.h
View file

@ -17,19 +17,5 @@ void DropoutGradientKernelImpl(
const float ratio,
T* dX_data);
template <typename T>
void BiasDropoutKernelImpl(
const cudaDeviceProp& prop,
cudaStream_t stream,
const int64_t N,
const fast_divmod fdm_dim,
const float ratio,
PhiloxGenerator& generator,
const T* X_data,
const T* bias_data,
const T* residual_data,
T* Y_data,
bool* mask_data);
} // namespace cuda
} // namespace onnxruntime

2
orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc
View file

@ -63,7 +63,6 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, BatchNormalizationGrad);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, double, BatchNormalizationGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, GatherGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BiasDropout);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, DropoutGrad);
// TODO: decprecate GatherND-1 after updating training models to opset-12
@ -174,7 +173,6 @@ Status RegisterRocmTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, ZeroGradient)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, ZeroGradient)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, BiasDropout)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, DropoutGrad)>,
// TODO: decprecate GatherND-1 after updating training models to opset-12