wehub-resource-sync
563 lines
19 KiB
C++
563 lines
19 KiB
C++
// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// 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.
|
|
|
|
#pragma once
|
|
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
#include "paddle/phi/backends/dynload/cublasLt.h"
|
|
#endif
|
|
|
|
#include "glog/logging.h"
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/kernels/funcs/dropout_impl_util.h"
|
|
#include "paddle/phi/kernels/funcs/functors.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_bias_act_utils.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_dropout_act_bias.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_dropout_common.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_layernorm_residual_dropout_bias.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_residual_dropout_bias.h"
|
|
#include "paddle/phi/kernels/layer_norm_kernel.h"
|
|
|
|
COMMON_DECLARE_bool(use_fast_math);
|
|
|
|
namespace phi {
|
|
namespace fusion {
|
|
|
|
struct DropoutParam {
|
|
uint64_t seed;
|
|
float dropout_prob;
|
|
bool is_upscale_in_train;
|
|
bool is_test;
|
|
bool fix_seed;
|
|
int increment{};
|
|
const DenseTensor* tensor_seed;
|
|
int seed_val;
|
|
|
|
DropoutParam() {
|
|
fix_seed = false;
|
|
seed = 0;
|
|
is_test = false;
|
|
is_upscale_in_train = false;
|
|
dropout_prob = 0.5;
|
|
tensor_seed = nullptr;
|
|
seed_val = 0;
|
|
}
|
|
|
|
DropoutParam(bool fix_seed_,
|
|
uint64_t seed_,
|
|
bool is_test_,
|
|
bool is_upscale_in_train_,
|
|
float dropout_prob_,
|
|
const DenseTensor* tensor_seed_,
|
|
int seed_val_) {
|
|
fix_seed = fix_seed_;
|
|
seed = seed_;
|
|
is_test = is_test_;
|
|
is_upscale_in_train = is_upscale_in_train_;
|
|
dropout_prob = dropout_prob_;
|
|
tensor_seed = tensor_seed_;
|
|
seed_val = seed_val_;
|
|
}
|
|
|
|
uint64_t UpdateSeedAndIncrement(const GPUContext& dev_ctx,
|
|
const uint64_t offset) {
|
|
uint64_t tmp_increment;
|
|
funcs::GetSeedDataAndIncrement(dev_ctx,
|
|
tensor_seed,
|
|
fix_seed,
|
|
seed_val,
|
|
offset,
|
|
&seed,
|
|
&tmp_increment);
|
|
increment = tmp_increment;
|
|
return increment;
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct DataTypeTraits {
|
|
using DataType = T;
|
|
};
|
|
|
|
template <>
|
|
struct DataTypeTraits<phi::float16> {
|
|
// Since LayerNormDirectCUDAFunctor register half type, we need to convert
|
|
// phi::float16 to half.
|
|
using DataType = half;
|
|
};
|
|
|
|
template <typename T,
|
|
typename MaskType,
|
|
typename InType = T,
|
|
typename OutType = T>
|
|
class FusedDropoutHelper {
|
|
private:
|
|
uint64_t GetIncrement(const GPUContext& dev_ctx) {
|
|
const int VecSize = MAX_CACHE_BYTES / sizeof(T);
|
|
const int real_vec_size = cols_ % VecSize == 0 ? VecSize : 1;
|
|
auto config = Get1DBlocksAnd2DGrids(dev_ctx,
|
|
static_cast<uint64_t>(rows_),
|
|
static_cast<uint64_t>(cols_),
|
|
real_vec_size);
|
|
uint64_t increment =
|
|
((cols_ - static_cast<uint64_t>(1)) /
|
|
(static_cast<uint64_t>(config.thread_per_block.x) *
|
|
static_cast<uint64_t>(config.block_per_grid.x) *
|
|
static_cast<uint64_t>(real_vec_size)) +
|
|
static_cast<uint64_t>(1)) *
|
|
static_cast<uint64_t>(real_vec_size);
|
|
increment = dropout_param_.UpdateSeedAndIncrement(dev_ctx, increment);
|
|
return increment;
|
|
}
|
|
|
|
public:
|
|
FusedDropoutHelper() {}
|
|
FusedDropoutHelper(const GPUContext& dev_ctx,
|
|
const int64_t rows,
|
|
const int64_t cols,
|
|
const DropoutParam& dropout_param,
|
|
const float residual_alpha = 1.0) {
|
|
rows_ = rows;
|
|
cols_ = cols;
|
|
dropout_param_ = dropout_param;
|
|
residual_alpha_ = residual_alpha;
|
|
}
|
|
|
|
// out = residual + dropout( src + bias )
|
|
void ResidualDropoutBias(const GPUContext& dev_ctx,
|
|
const InType* src,
|
|
const T* residual,
|
|
const T* bias,
|
|
OutType* out,
|
|
MaskType* mask,
|
|
const float quant_last_in_scale = 1.0,
|
|
const float* dequant_out_scale_data = nullptr,
|
|
const float quant_next_in_scale = 1.0) {
|
|
auto increment = GetIncrement(dev_ctx);
|
|
LaunchResidualDropoutBias<T, MaskType, InType, OutType>(
|
|
rows_,
|
|
cols_,
|
|
increment,
|
|
dropout_param_.seed,
|
|
dropout_param_.dropout_prob,
|
|
dropout_param_.is_test,
|
|
dropout_param_.is_upscale_in_train,
|
|
src,
|
|
residual,
|
|
bias,
|
|
mask,
|
|
out,
|
|
dev_ctx,
|
|
quant_last_in_scale,
|
|
dequant_out_scale_data,
|
|
quant_next_in_scale,
|
|
residual_alpha_);
|
|
}
|
|
|
|
void ResidualDropoutBiasGrad(const GPUContext& dev_ctx,
|
|
const T* d_out,
|
|
const MaskType* mask,
|
|
T* d_src,
|
|
T* d_residual,
|
|
T* d_bias) {
|
|
LaunchResidualDropoutBiasGrad<T, uint8_t>(
|
|
d_out,
|
|
mask,
|
|
dropout_param_.dropout_prob,
|
|
dropout_param_.is_upscale_in_train,
|
|
rows_,
|
|
cols_,
|
|
d_src,
|
|
d_bias,
|
|
dev_ctx);
|
|
if (d_residual) {
|
|
phi::memory_utils::Copy(dev_ctx.GetPlace(),
|
|
d_residual,
|
|
dev_ctx.GetPlace(),
|
|
d_out,
|
|
rows_ * cols_ * sizeof(T),
|
|
dev_ctx.stream());
|
|
}
|
|
}
|
|
|
|
// out = dropout(activation(src + bias))
|
|
void DropoutActBias(const GPUContext& dev_ctx,
|
|
const InType* src,
|
|
const T* bias,
|
|
const std::string& act_method,
|
|
OutType* out,
|
|
MaskType* mask,
|
|
const float quant_last_in_scale = 1.0,
|
|
const float* dequant_out_scale_data = nullptr,
|
|
const float quant_next_in_scale = 1.0,
|
|
const int quant_round_type = 1,
|
|
const float quant_max_bound = 127.0,
|
|
const float quant_min_bound = -127.0) {
|
|
auto increment = GetIncrement(dev_ctx);
|
|
if (act_method == "gelu") {
|
|
if (FLAGS_use_fast_math) {
|
|
phi::fusion::FastGeluFunctor<T> fast_gelu;
|
|
phi::fusion::LaunchDropoutActBias<T,
|
|
MaskType,
|
|
phi::fusion::FastGeluFunctor<T>,
|
|
InType,
|
|
OutType>(
|
|
fast_gelu,
|
|
dropout_param_.seed,
|
|
rows_,
|
|
cols_,
|
|
dropout_param_.increment,
|
|
dropout_param_.dropout_prob,
|
|
dropout_param_.is_upscale_in_train,
|
|
dropout_param_.is_test,
|
|
src,
|
|
bias,
|
|
out,
|
|
mask,
|
|
dev_ctx,
|
|
quant_last_in_scale,
|
|
dequant_out_scale_data,
|
|
quant_next_in_scale,
|
|
quant_round_type,
|
|
quant_max_bound,
|
|
quant_min_bound);
|
|
} else {
|
|
phi::fusion::LayerNormParamTypeGeluFunctor<T> gelu;
|
|
phi::fusion::LaunchDropoutActBias<
|
|
T,
|
|
MaskType,
|
|
phi::fusion::LayerNormParamTypeGeluFunctor<T>,
|
|
InType,
|
|
OutType>(gelu,
|
|
dropout_param_.seed,
|
|
rows_,
|
|
cols_,
|
|
dropout_param_.increment,
|
|
dropout_param_.dropout_prob,
|
|
dropout_param_.is_upscale_in_train,
|
|
dropout_param_.is_test,
|
|
src,
|
|
bias,
|
|
out,
|
|
mask,
|
|
dev_ctx,
|
|
quant_last_in_scale,
|
|
dequant_out_scale_data,
|
|
quant_next_in_scale,
|
|
quant_round_type,
|
|
quant_max_bound,
|
|
quant_min_bound);
|
|
}
|
|
} else if (act_method == "relu") {
|
|
funcs::ReluFunctor<T> relu;
|
|
phi::fusion::LaunchDropoutActBias<T,
|
|
MaskType,
|
|
funcs::ReluFunctor<T>,
|
|
InType,
|
|
OutType>(
|
|
relu,
|
|
dropout_param_.seed,
|
|
rows_,
|
|
cols_,
|
|
increment,
|
|
dropout_param_.dropout_prob,
|
|
dropout_param_.is_upscale_in_train,
|
|
dropout_param_.is_test,
|
|
src,
|
|
bias,
|
|
out,
|
|
mask,
|
|
dev_ctx,
|
|
quant_last_in_scale,
|
|
dequant_out_scale_data,
|
|
quant_next_in_scale,
|
|
quant_round_type,
|
|
quant_max_bound,
|
|
quant_min_bound);
|
|
} else {
|
|
PADDLE_THROW(errors::InvalidArgument(
|
|
"Currently only supports gelu or relu activation functions!"));
|
|
}
|
|
}
|
|
|
|
void DropoutActBiasGrad(const GPUContext& dev_ctx,
|
|
const T* dout,
|
|
const T* src,
|
|
const T* bias,
|
|
const MaskType* mask,
|
|
T* d_src,
|
|
T* d_bias,
|
|
const std::string& act_method) {
|
|
if (act_method == "gelu") {
|
|
phi::fusion::GeluGradFunctor<T> gelu_grad;
|
|
phi::fusion::LaunchDropoutActBiasGrad<T,
|
|
MaskType,
|
|
phi::fusion::GeluGradFunctor<T>>(
|
|
gelu_grad,
|
|
dout,
|
|
mask,
|
|
src,
|
|
bias,
|
|
dropout_param_.dropout_prob,
|
|
dropout_param_.is_upscale_in_train,
|
|
rows_,
|
|
cols_,
|
|
d_src,
|
|
d_bias,
|
|
dev_ctx);
|
|
} else if (act_method == "relu") {
|
|
funcs::ReluGradFunctor<T> relu_grad;
|
|
phi::fusion::
|
|
LaunchDropoutActBiasGrad<T, MaskType, funcs::ReluGradFunctor<T>>(
|
|
relu_grad,
|
|
dout,
|
|
mask,
|
|
src,
|
|
bias,
|
|
dropout_param_.dropout_prob,
|
|
dropout_param_.is_upscale_in_train,
|
|
rows_,
|
|
cols_,
|
|
d_src,
|
|
d_bias,
|
|
dev_ctx);
|
|
} else {
|
|
PADDLE_THROW(errors::InvalidArgument(
|
|
"Currently only supports gelu or relu activation functions!"));
|
|
}
|
|
}
|
|
|
|
protected:
|
|
int64_t rows_;
|
|
int64_t cols_;
|
|
DropoutParam dropout_param_;
|
|
float residual_alpha_;
|
|
};
|
|
|
|
template <typename T,
|
|
typename MaskType,
|
|
typename InType = T,
|
|
typename OutType = T>
|
|
class FusedDropoutLayerNormHelper
|
|
: public FusedDropoutHelper<T, MaskType, InType, OutType> {
|
|
public:
|
|
FusedDropoutLayerNormHelper() {}
|
|
FusedDropoutLayerNormHelper(const int64_t rows,
|
|
const int64_t cols,
|
|
const float epsilon,
|
|
const float residual_alpha = 1.0) {
|
|
using U = funcs::LayerNormParamType<T>;
|
|
this->rows_ = rows;
|
|
this->cols_ = cols;
|
|
epsilon_ = epsilon;
|
|
this->residual_alpha_ = residual_alpha;
|
|
}
|
|
|
|
FusedDropoutLayerNormHelper(const GPUContext& dev_ctx,
|
|
const int64_t rows,
|
|
const int64_t cols,
|
|
const DropoutParam& dropout_param,
|
|
const float epsilon,
|
|
const float residual_alpha = 1.0)
|
|
: FusedDropoutHelper<T, MaskType, InType, OutType>(
|
|
dev_ctx, rows, cols, dropout_param, residual_alpha) {
|
|
using U = funcs::LayerNormParamType<T>;
|
|
epsilon_ = epsilon;
|
|
}
|
|
|
|
// call layer_norm
|
|
void LayerNorm(const GPUContext& dev_ctx,
|
|
const InType* src,
|
|
const funcs::LayerNormParamType<T>* gamma,
|
|
const funcs::LayerNormParamType<T>* beta,
|
|
OutType* out,
|
|
funcs::LayerNormParamType<T>* mean,
|
|
funcs::LayerNormParamType<T>* variance) {
|
|
using InDataType = typename DataTypeTraits<InType>::DataType;
|
|
using OutDataType = typename DataTypeTraits<OutType>::DataType;
|
|
|
|
phi::LayerNormDirectCUDAFunctor<InDataType, funcs::LayerNormParamType<T>>
|
|
layer_norm;
|
|
std::vector<int64_t> src_shape{this->rows_, this->cols_};
|
|
layer_norm(dev_ctx.stream(),
|
|
reinterpret_cast<const InDataType*>(src),
|
|
src_shape,
|
|
beta,
|
|
gamma,
|
|
reinterpret_cast<OutDataType*>(out),
|
|
mean,
|
|
variance,
|
|
1,
|
|
epsilon_);
|
|
}
|
|
|
|
void LayerNormGrad(const GPUContext& dev_ctx,
|
|
const T* dout,
|
|
const T* src,
|
|
const funcs::LayerNormParamType<T>* gamma,
|
|
const funcs::LayerNormParamType<T>* mean,
|
|
const funcs::LayerNormParamType<T>* variance,
|
|
T* d_src,
|
|
funcs::LayerNormParamType<T>* d_scale,
|
|
funcs::LayerNormParamType<T>* d_bias) {
|
|
using U = funcs::LayerNormParamType<T>;
|
|
funcs::LayerNormBackward<T, U>(src,
|
|
dout,
|
|
gamma,
|
|
mean,
|
|
variance,
|
|
d_src,
|
|
d_scale,
|
|
d_bias,
|
|
epsilon_,
|
|
this->rows_,
|
|
this->cols_,
|
|
dev_ctx);
|
|
}
|
|
|
|
// out = layernorm(residual + dropout(src + bias))
|
|
template <typename P = funcs::LayerNormParamType<T>,
|
|
bool is_same_type = false>
|
|
void LayernormResidualDropoutBias(
|
|
const GPUContext& dev_ctx,
|
|
const InType* src,
|
|
const T* residual,
|
|
const T* bias,
|
|
const P* gamma,
|
|
const P* beta,
|
|
T* dropout_out,
|
|
MaskType* mask,
|
|
OutType* out,
|
|
funcs::LayerNormParamType<T>* mean,
|
|
funcs::LayerNormParamType<T>* variance,
|
|
const float quant_last_in_scale = 1.0,
|
|
const float* dequant_out_scale_data = nullptr,
|
|
const float quant_next_in_scale = 1.0,
|
|
const int quant_round_type = 1,
|
|
const float quant_max_bound = 127.0,
|
|
const float quant_min_bound = -127.0) {
|
|
using U = funcs::LayerNormParamType<T>;
|
|
int vec_size = MAX_CACHE_BYTES / sizeof(T);
|
|
if (this->cols_ % vec_size != 0) {
|
|
vec_size = 1;
|
|
}
|
|
int threads = funcs::GetDesiredBlockDim(this->cols_ / vec_size);
|
|
int increment = ((this->cols_ - 1) / (threads * vec_size) + 1) * vec_size;
|
|
increment = this->dropout_param_.UpdateSeedAndIncrement(dev_ctx, increment);
|
|
LaunchLayernormResidualDropoutBias<T,
|
|
MaskType,
|
|
U,
|
|
is_same_type,
|
|
InType,
|
|
OutType>(
|
|
this->rows_,
|
|
this->cols_,
|
|
increment,
|
|
this->dropout_param_.seed,
|
|
this->dropout_param_.dropout_prob,
|
|
epsilon_,
|
|
this->dropout_param_.is_upscale_in_train,
|
|
this->dropout_param_.is_test,
|
|
src,
|
|
residual,
|
|
bias,
|
|
gamma,
|
|
beta,
|
|
mask,
|
|
dropout_out,
|
|
out,
|
|
mean,
|
|
variance,
|
|
dev_ctx,
|
|
quant_last_in_scale,
|
|
dequant_out_scale_data,
|
|
quant_next_in_scale,
|
|
quant_round_type,
|
|
quant_max_bound,
|
|
quant_min_bound,
|
|
this->residual_alpha_);
|
|
}
|
|
|
|
template <typename P = funcs::LayerNormParamType<T>,
|
|
bool is_same_type = false>
|
|
void LayernormResidualDropoutBiasGrad(
|
|
const GPUContext& dev_ctx,
|
|
const T* d_out,
|
|
const T* layernorm_src,
|
|
const MaskType* mask,
|
|
const P* gamma,
|
|
const funcs::LayerNormParamType<T>* mean,
|
|
const funcs::LayerNormParamType<T>* variance,
|
|
T* d_layernorm_src,
|
|
P* d_scale,
|
|
P* d_layernorm_bias,
|
|
T* d_dropout_src,
|
|
T* d_bias,
|
|
T* d_residual) {
|
|
using U = funcs::LayerNormParamType<T>;
|
|
bool can_call_1024_kernel = false;
|
|
// Fast impl for cases when cols is 1024 and linear_bias is nullptr.
|
|
// In fact, linear_bias is not nullptr is also feasible for impl.
|
|
// Here, we do not support it.
|
|
if (this->cols_ == 1024 && d_bias == nullptr && d_scale != nullptr &&
|
|
d_layernorm_bias != nullptr && sizeof(T) <= 4) {
|
|
can_call_1024_kernel = true;
|
|
}
|
|
VLOG(6) << "LaunchLayernormResidualDropoutGrad = " << can_call_1024_kernel;
|
|
|
|
if (can_call_1024_kernel) {
|
|
LaunchLayernormResidualDropoutGrad<T, U, MaskType, is_same_type>(
|
|
dev_ctx,
|
|
this->rows_,
|
|
this->cols_,
|
|
epsilon_,
|
|
this->dropout_param_.dropout_prob,
|
|
this->dropout_param_.is_upscale_in_train,
|
|
d_out,
|
|
layernorm_src,
|
|
gamma,
|
|
mean,
|
|
variance,
|
|
mask,
|
|
d_scale,
|
|
d_layernorm_bias,
|
|
d_residual,
|
|
d_dropout_src);
|
|
} else {
|
|
funcs::LayerNormBackward<T, U, is_same_type>(layernorm_src,
|
|
d_out,
|
|
gamma,
|
|
mean,
|
|
variance,
|
|
d_layernorm_src,
|
|
d_scale,
|
|
d_layernorm_bias,
|
|
epsilon_,
|
|
this->rows_,
|
|
this->cols_,
|
|
dev_ctx);
|
|
this->ResidualDropoutBiasGrad(
|
|
dev_ctx, d_layernorm_src, mask, d_dropout_src, d_residual, d_bias);
|
|
}
|
|
}
|
|
|
|
protected:
|
|
float epsilon_;
|
|
};
|
|
|
|
} // namespace fusion
|
|
} // namespace phi
|