diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 1f88aac942..26fd54bc8c 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout.cc b/onnxruntime/contrib_ops/cuda/math/bias_dropout.cc similarity index 54% rename from orttraining/orttraining/training_ops/cuda/nn/dropout.cc rename to onnxruntime/contrib_ops/cuda/math/bias_dropout.cc index 8d970bdd90..a13974ad99 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout.cc +++ b/onnxruntime/contrib_ops/cuda/math/bias_dropout.cc @@ -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(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType()} -#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()) \ - .InputMemoryType(2), \ - DropoutGrad); - -REGISTER_GRADIENT_KERNEL(DropoutGrad) - -template -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::MappedType CudaT; - - const CudaT* dY_data = reinterpret_cast(dY.template Data()); - CudaT* dX_data = reinterpret_cast(dX.template MutableData()); - DropoutGradientKernelImpl(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 -struct GetRatioDataImpl { - void operator()(const Tensor* ratio, float& ratio_data) const { - ratio_data = static_cast(*(ratio->template Data())); - 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(0); - const TensorShape& shape = dY->Shape(); - const int64_t N = shape.Size(); - - auto mask = context->Input(1); - ORT_ENFORCE(mask->Shape().Size() == N); - const bool* mask_data = mask->template Data(); - - //Get the ratio_data - float ratio_data = default_ratio_; - auto ratio = context->Input(2); - if (ratio) { - utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); - t_disp.Invoke(ratio, ratio_data); - } - - auto dX = context->Output(0, shape); - - utils::MLTypeCallDispatcher t_disp(dY->GetElementType()); - t_disp.Invoke(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 \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout.h b/onnxruntime/contrib_ops/cuda/math/bias_dropout.h similarity index 58% rename from orttraining/orttraining/training_ops/cuda/nn/dropout.h rename to onnxruntime/contrib_ops/cuda/math/bias_dropout.h index 58cf470c50..a6e2166997 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout.h +++ b/onnxruntime/contrib_ops/cuda/math/bias_dropout.h @@ -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 +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 diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu b/onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu similarity index 70% rename from orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu rename to onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu index eed291e1d7..fb6cc7c7b1 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu +++ b/onnxruntime/contrib_ops/cuda/math/bias_dropout_impl.cu @@ -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 #include namespace onnxruntime { +namespace contrib { namespace cuda { -template -__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 -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(CeilDiv(N, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); - DropoutGradientKernel - <<>>(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 @@ -180,4 +129,5 @@ SPECIALIZED_BIAS_DROPOUT_IMPL(nv_bfloat16) #endif } // namespace cuda +} // namespace contrib { } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc index 71fa8b3278..a15d9eea03 100644 --- a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc @@ -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, // BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index de257885ed..3248dd688e 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -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(); diff --git a/onnxruntime/core/providers/cuda/nn/dropout.cc b/onnxruntime/core/providers/cuda/nn/dropout.cc index 0293df46a9..42ec6597ac 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout.cc +++ b/onnxruntime/core/providers/cuda/nn/dropout.cc @@ -6,15 +6,6 @@ namespace onnxruntime { namespace cuda { -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 -#define ALL_IEEE_FLOAT_TENSOR_TYPES {DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType()} -#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(2), Dropout); +Status Dropout::ComputeInternal(OpKernelContext* context) const { + //Get X_data + const Tensor* X = context->Input(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(1); + if (ratio) { + utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); + t_disp.Invoke(ratio, ratio_data); + } + + const Tensor* training_mode = context->Input(2); + //Check for inference mode. + if ((0 == ratio_data) ||(training_mode == nullptr || *(training_mode->Data()) == 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(), true, N * sizeof(bool), Stream())); + } + + return Status::OK(); + } + + IAllocatorUniquePtr 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(); + temp_mask_buffer = GetScratchBuffer(N); + return temp_mask_buffer.get(); + }(); + + PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default(); + + utils::MLTypeCallDispatcher t_disp(X->GetElementType()); + t_disp.Invoke(GetDeviceProp(), Stream(), N, ratio_data, generator, *X, *Y, mask_data); + + return Status::OK(); +} + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/dropout.h b/onnxruntime/core/providers/cuda/nn/dropout.h index bea0bfbb42..2a14f08cfa 100644 --- a/onnxruntime/core/providers/cuda/nn/dropout.h +++ b/onnxruntime/core/providers/cuda/nn/dropout.h @@ -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(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType() } +#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 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(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(1); - if (ratio) { - utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); - t_disp.Invoke(ratio, ratio_data); - } - - const Tensor* training_mode = context->Input(2); - //Check for inference mode. - if ((0 == ratio_data) ||(training_mode == nullptr || *(training_mode->Data()) == 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(), true, N * sizeof(bool), Stream())); - } - - return Status::OK(); - } - - IAllocatorUniquePtr 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(); - temp_mask_buffer = GetScratchBuffer(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 t_disp(X->GetElementType()); - t_disp.Invoke(GetDeviceProp(), Stream(), N, ratio_data, generator, *X, *Y, mask_data); - - return Status::OK(); -} - } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/bias_dropout_op_test.cc b/onnxruntime/test/contrib_ops/bias_dropout_op_test.cc new file mode 100644 index 0000000000..f28e313be1 --- /dev/null +++ b/onnxruntime/test/contrib_ops/bias_dropout_op_test.cc @@ -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 +#include +#include +#include + +#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& 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(1), std::multiplies<>{}); + const std::vector input = ValueRange(input_size, 1.0f, 1.0f); + t.AddInput("data", input_shape, input); + + std::vector bias_shape{input_shape.back()}; + const auto bias_size = input_shape.back(); + const std::vector 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 residual(residual_size, residual_value); + t.AddInput("residual", input_shape, residual); + } else { + t.AddMissingOptionalInput(); + } + + if (ratio == -1.0f) { + if (use_float16_ratio) { + t.AddMissingOptionalInput(); + } else { + t.AddMissingOptionalInput(); + } + // 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("output", input_shape, input); // we'll do our own output verification + + std::unique_ptr mask_buffer{}; + if (use_mask) { + mask_buffer = onnxruntime::make_unique(input_size); + t.AddOutput("mask", input_shape, mask_buffer.get(), input_size); + } else { + t.AddMissingOptionalOutput(); + } + + auto output_verifier = [&](const std::vector& fetches, const std::string& provider_type) { + ASSERT_GE(fetches.size(), 1); + const auto& output_tensor = FetchTensor(fetches[0]); + auto output_span = output_tensor.DataAsSpan(); + + 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(output_span.size())) << "provider: " << provider_type; + } else { + ASSERT_NEAR(static_cast(num_dropped_values) / static_cast(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(); + 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 diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 049616570b..19dacf1b1c 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -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 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(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, logger); + ASSERT_STATUS_OK(ret); + + std::map 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 p_model; diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 2aae8b1524..cf78dc7e50 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -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) diff --git a/orttraining/orttraining/test/optimizer/graph_transform_test.cc b/orttraining/orttraining/test/optimizer/graph_transform_test.cc index bd9d99272c..b33408bd73 100644 --- a/orttraining/orttraining/test/optimizer/graph_transform_test.cc +++ b/orttraining/orttraining/test/optimizer/graph_transform_test.cc @@ -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 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(), TransformerLevel::Level2); - auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, logger); - ASSERT_STATUS_OK(ret); - - std::map 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(); diff --git a/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc b/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc index ded07dc7c5..71ca986628 100644 --- a/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc @@ -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& 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(1), std::multiplies<>{}); - const std::vector input = ValueRange(input_size, 1.0f, 1.0f); - t.AddInput("data", input_shape, input); - - std::vector bias_shape{input_shape.back()}; - const auto bias_size = input_shape.back(); - const std::vector 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 residual(residual_size, residual_value); - t.AddInput("residual", input_shape, residual); - } else { - t.AddMissingOptionalInput(); - } - - if (ratio == -1.0f) { - if (use_float16_ratio) { - t.AddMissingOptionalInput(); - } else { - t.AddMissingOptionalInput(); - } - // 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("output", input_shape, input); // we'll do our own output verification - - std::unique_ptr mask_buffer{}; - if (use_mask) { - mask_buffer = onnxruntime::make_unique(input_size); - t.AddOutput("mask", input_shape, mask_buffer.get(), input_size); - } else { - t.AddMissingOptionalOutput(); - } - - auto output_verifier = [&](const std::vector& fetches, const std::string& provider_type) { - ASSERT_GE(fetches.size(), 1); - const auto& output_tensor = FetchTensor(fetches[0]); - auto output_span = output_tensor.DataAsSpan(); - - 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(output_span.size())) << "provider: " << provider_type; - } else { - ASSERT_NEAR(static_cast(num_dropped_values) / static_cast(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(); - 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& input_dims, bool default_ratio = true) { const auto input_shape = TensorShape(input_dims); diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 78b0cf0bd7..21fde3adcf 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -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, BuildKernelCreateInfo, - - BuildKernelCreateInfo, + BuildKernelCreateInfo, // TODO: decprecate GatherND-1 after updating training models to opset-12 diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc new file mode 100644 index 0000000000..5bbfefb80e --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.cc @@ -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()) \ + .InputMemoryType(2), \ + DropoutGrad); + +REGISTER_GRADIENT_KERNEL(DropoutGrad) + +template +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::MappedType CudaT; + + const CudaT* dY_data = reinterpret_cast(dY.template Data()); + CudaT* dX_data = reinterpret_cast(dX.template MutableData()); + DropoutGradientKernelImpl(stream, N, dY_data, mask_data, ratio_data, dX_data); + } +}; + +Status DropoutGrad::ComputeInternal(OpKernelContext* context) const { + auto dY = context->Input(0); + const TensorShape& shape = dY->Shape(); + const int64_t N = shape.Size(); + + auto mask = context->Input(1); + ORT_ENFORCE(mask->Shape().Size() == N); + const bool* mask_data = mask->template Data(); + + //Get the ratio_data + float ratio_data = default_ratio_; + auto ratio = context->Input(2); + if (ratio) { + utils::MLTypeCallDispatcher t_disp(ratio->GetElementType()); + t_disp.Invoke(ratio, ratio_data); + } + + auto dX = context->Output(0, shape); + + utils::MLTypeCallDispatcher t_disp(dY->GetElementType()); + t_disp.Invoke(Stream(), N, *dY, mask_data, ratio_data, *dX); + + return Status::OK(); +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.h b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.h new file mode 100644 index 0000000000..8155329221 --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad.h @@ -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 diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu new file mode 100644 index 0000000000..c8780b7b26 --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.cu @@ -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 +#include + +namespace onnxruntime { +namespace cuda { + +template +__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 +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(CeilDiv(N, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + DropoutGradientKernel + <<>>(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 diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.h similarity index 58% rename from orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h rename to orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.h index 8dbf3f9655..e36878e259 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_grad_impl.h @@ -17,19 +17,5 @@ void DropoutGradientKernelImpl( const float ratio, T* dX_data); -template -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 diff --git a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc index 616a98964d..fece1d6ece 100644 --- a/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc +++ b/orttraining/orttraining/training_ops/rocm/rocm_training_kernels.cc @@ -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, BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, // TODO: decprecate GatherND-1 after updating training models to opset-12