Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

use elementwise to optimize gelu backward implementation on GPU #38263

Merged
merged 6 commits into from
Dec 22, 2021
Merged

use elementwise to optimize gelu backward implementation on GPU #38263

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e2ea903
optimize gelu backward
Zjq9409 Dec 19, 2021
3302e63
optimize gelu backward
Zjq9409 Dec 20, 2021
7c0d65c
optimize code
Zjq9409 Dec 20, 2021
65af9f6
Merge branch 'develop' into gelu_bw_opt
Zjq9409 Dec 21, 2021
33a27c4
Number to expression
Zjq9409 Dec 21, 2021
bf4e1d6
Replacement number
Zjq9409 Dec 22, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 70 additions & 6 deletions paddle/fluid/operators/gelu_op.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ limitations under the License. */
#include "paddle/fluid/operators/amp/fp16_type_traits.h"
#include "paddle/fluid/operators/elementwise/elementwise_op_broadcast.cu.h"
#include "paddle/fluid/operators/gelu_op.h"
#include "paddle/fluid/platform/float16.h"

namespace paddle {
namespace operators {
Expand All @@ -27,9 +26,11 @@ struct GeluWithApproximateFunctor {
// this function is tanh approximation of gelu
MPType x = static_cast<MPType>(arg_x);
MPType one = static_cast<MPType>(1);
MPType out = x * static_cast<MPType>(0.5) *
(one + tanh(static_cast<MPType>(0.79788456) * x *
(one + static_cast<MPType>(0.044715) * x * x)));
MPType half = static_cast<MPType>(0.5);
MPType kAlpha = static_cast<MPType>(M_2_SQRTPI * M_SQRT1_2);
auto tanh_out =
tanh(kAlpha * x * (one + static_cast<MPType>(GELU_CONSTANT) * x * x));
MPType out = x * half * (one + tanh_out);
return static_cast<T>(out);
}
};
Expand All @@ -40,9 +41,10 @@ struct GeluWithoutApproximateFunctor {
inline HOSTDEVICE T operator()(T arg_x) {
// actual gelu with approximation = false
MPType x = static_cast<MPType>(arg_x);
MPType one = static_cast<MPType>(1);
MPType half = static_cast<MPType>(0.5);
MPType erf_out = erf(x * static_cast<MPType>(M_SQRT1_2));
MPType out =
x * static_cast<MPType>(0.5) * (static_cast<MPType>(1) + erf_out);
MPType out = x * half * (one + erf_out);
return static_cast<T>(out);
}
};
Expand Down Expand Up @@ -71,6 +73,68 @@ class GeluKernel<platform::CUDADeviceContext, T>
}
};

template <typename T>
struct GeluWithApproximateGradFunctor {
using MPType = typename details::MPTypeTrait<T>::Type;
inline HOSTDEVICE T operator()(T arg_x, T arg_dout) {
MPType x = static_cast<MPType>(arg_x);
MPType dout = static_cast<MPType>(arg_dout);
MPType one = static_cast<MPType>(1);
MPType half = static_cast<MPType>(0.5);
MPType kAlpha = static_cast<MPType>(M_2_SQRTPI * M_SQRT1_2);
MPType kBeta =
kAlpha * static_cast<MPType>(GELU_CONSTANT) * static_cast<MPType>(3);
auto cube_x = x * x * x;
auto tanh_out =
tanh(kAlpha * ((static_cast<MPType>(GELU_CONSTANT) * cube_x) + x));
auto ans =
half * (one + tanh_out +
(one - tanh_out * tanh_out) * (x * kAlpha + kBeta * cube_x));
return static_cast<T>(ans * dout);
}
};

template <typename T>
struct GeluWithoutApproximateGradFunctor {
using MPType = typename details::MPTypeTrait<T>::Type;
inline HOSTDEVICE T operator()(T arg_x, T arg_dout) {
MPType x = static_cast<MPType>(arg_x);
MPType dout = static_cast<MPType>(arg_dout);
MPType one = static_cast<MPType>(1);
MPType half = static_cast<MPType>(0.5);
MPType kAlpha = static_cast<MPType>(M_2_SQRTPI * M_SQRT1_2);
auto ans = half * (one + erf(x * static_cast<MPType>(M_SQRT1_2))) +
half * kAlpha * x * exp(-half * x * x);
return static_cast<T>(ans * dout);
}
};

template <typename T>
class GeluGradKernel<platform::CUDADeviceContext, T>
: public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& context) const override {
auto* x = context.Input<framework::Tensor>("X");
auto* dout =
context.Input<framework::Tensor>(framework::GradVarName("Out"));
auto* dx = context.Output<framework::Tensor>(framework::GradVarName("X"));
auto approximate = context.Attr<bool>("approximate");
dx->mutable_data<T>(dout->place());

std::vector<const framework::Tensor*> ins = {x, dout};
std::vector<framework::Tensor*> outs = {dx};
const auto& dev_ctx =
context.template device_context<platform::CUDADeviceContext>();
if (approximate) {
LaunchElementwiseCudaKernel<ElementwiseType::kBinary, T, T>(
dev_ctx, ins, &outs, 0, GeluWithApproximateGradFunctor<T>());
} else {
LaunchElementwiseCudaKernel<ElementwiseType::kBinary, T, T>(
dev_ctx, ins, &outs, 0, GeluWithoutApproximateGradFunctor<T>());
}
}
};

} // namespace operators
} // namespace paddle

Expand Down
15 changes: 9 additions & 6 deletions paddle/fluid/operators/gelu_op.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@ limitations under the License. */
namespace paddle {
namespace operators {

#define GELU_CONSTANT 0.044715

template <typename T>
struct GeluFunctor {
template <typename Device, typename X, typename Out>
Expand All @@ -41,14 +43,14 @@ struct GeluFunctor {
auto casted_x = x.template cast<float>();
auto temp =
(static_cast<float>(M_2_SQRTPI * M_SQRT1_2) *
(casted_x + static_cast<float>(0.044715) * casted_x.cube()))
(casted_x + static_cast<float>(GELU_CONSTANT) * casted_x.cube()))
.tanh();
out.device(d) = (casted_x * static_cast<float>(0.5) *
(static_cast<float>(1) + temp))
.template cast<T>();
} else {
auto temp = (static_cast<T>(M_2_SQRTPI * M_SQRT1_2) *
(x + static_cast<T>(0.044715) * x.cube()))
(x + static_cast<T>(GELU_CONSTANT) * x.cube()))
.tanh();
out.device(d) = x * static_cast<T>(0.5) * (static_cast<T>(1) + temp);
}
Expand Down Expand Up @@ -101,10 +103,10 @@ struct GeluGradFunctor {

const float kAlpha = static_cast<float>(M_2_SQRTPI * M_SQRT1_2);
const float kBeta =
kAlpha * static_cast<float>(0.044715) * static_cast<float>(3);
kAlpha * static_cast<float>(GELU_CONSTANT) * static_cast<float>(3);
const auto y =
(kAlpha *
((static_cast<float>(0.044715) * casted_x.cube()) + casted_x))
((static_cast<float>(GELU_CONSTANT) * casted_x.cube()) + casted_x))
.tanh();
dx.device(d) = (static_cast<float>(0.5) * casted_dout *
(static_cast<float>(1) + y +
Expand All @@ -113,9 +115,10 @@ struct GeluGradFunctor {
.template cast<T>();
} else {
const T kAlpha = static_cast<T>(M_2_SQRTPI * M_SQRT1_2);
const T kBeta = kAlpha * static_cast<T>(0.044715) * static_cast<T>(3);
const T kBeta =
kAlpha * static_cast<T>(GELU_CONSTANT) * static_cast<T>(3);
const auto y =
(kAlpha * ((static_cast<T>(0.044715) * x.cube()) + x)).tanh();
(kAlpha * ((static_cast<T>(GELU_CONSTANT) * x.cube()) + x)).tanh();
dx.device(d) = static_cast<T>(0.5) * dout *
(static_cast<T>(1) + y +
(x - x * y.square()) * (kAlpha + kBeta * x.square()));
Expand Down