wehub-resource-sync
319 lines
13 KiB
Plaintext
319 lines
13 KiB
Plaintext
// Copyright (c) 2024 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.
|
|
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/phi/api/include/tensor.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/broadcast_function.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
|
|
#include "paddle/phi/kernels/funcs/layer_norm_impl.cu.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_attention_utils.h"
|
|
#include "paddle/phi/kernels/fusion/gpu/fused_dropout_helper.h"
|
|
#include "paddle/phi/kernels/impl/matmul_grad_kernel_impl.h"
|
|
|
|
namespace phi {
|
|
namespace fusion {
|
|
|
|
template <typename T, typename Context>
|
|
void MatMul(const GPUContext& dev_ctx,
|
|
const DenseTensor& a,
|
|
const DenseTensor& b,
|
|
DenseTensor* c) {
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
auto a_2d = phi::FoldInitDims(a);
|
|
auto b_2d = phi::FoldInitDims(b);
|
|
auto mat_dim_a = funcs::CreateMatrixDescriptor(a_2d.dims(), 0, false);
|
|
auto mat_dim_b = funcs::CreateMatrixDescriptor(b_2d.dims(), 0, false);
|
|
T alpha = static_cast<T>(1.0);
|
|
blas.MatMul(a, mat_dim_a, b, mat_dim_b, alpha, c, T(0));
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FFN(const GPUContext& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& linear1_weight,
|
|
const DenseTensor* linear1_bias,
|
|
const DenseTensor& linear2_weight,
|
|
const DenseTensor* linear2_bias,
|
|
const DenseTensor* ln1_scale,
|
|
const DenseTensor* ln1_bias,
|
|
const DenseTensor* ln2_scale,
|
|
const DenseTensor* ln2_bias,
|
|
DenseTensor* out,
|
|
DenseTensor* dropout1_mask,
|
|
DenseTensor* dropout2_mask,
|
|
DenseTensor* ln1_mean,
|
|
DenseTensor* ln1_variance,
|
|
DenseTensor* ln2_mean,
|
|
DenseTensor* ln2_variance,
|
|
DenseTensor* linear1_out,
|
|
DenseTensor* ln1_out,
|
|
DenseTensor* dropout1_out,
|
|
DenseTensor* dropout2_out,
|
|
const int bsz_seq,
|
|
const int d_model,
|
|
const int dim_feedforward,
|
|
const std::string& act_method,
|
|
const bool pre_layer_norm,
|
|
const float epsilon1,
|
|
const float epsilon2,
|
|
const bool add_residual,
|
|
const int ring_id,
|
|
const fusion::DropoutParam& dropout_param1,
|
|
const fusion::DropoutParam& dropout_param2) {
|
|
fusion::FusedDropoutLayerNormHelper<T, uint8_t> pre_layernorm_helper(
|
|
bsz_seq, d_model, epsilon1);
|
|
fusion::FusedDropoutHelper<T, uint8_t> fused_act_dropout_helper(
|
|
dev_ctx, bsz_seq, dim_feedforward, dropout_param1);
|
|
fusion::FusedDropoutLayerNormHelper<T, uint8_t>
|
|
fused_dropout_layernorm_helper(
|
|
dev_ctx, bsz_seq, d_model, dropout_param2, epsilon2);
|
|
|
|
using U = funcs::LayerNormParamType<T>;
|
|
const DenseTensor* in = &x;
|
|
|
|
const U* ln1_scale_ptr =
|
|
ln1_scale == nullptr ? nullptr : ln1_scale->data<U>();
|
|
const U* ln1_bias_ptr = ln1_bias == nullptr ? nullptr : ln1_bias->data<U>();
|
|
const U* ln2_scale_ptr =
|
|
ln2_scale == nullptr ? nullptr : ln2_scale->data<U>();
|
|
const U* ln2_bias_ptr = ln2_bias == nullptr ? nullptr : ln2_bias->data<U>();
|
|
const T* linear1_bias_ptr =
|
|
linear1_bias == nullptr ? nullptr : linear1_bias->data<T>();
|
|
const T* linear2_bias_ptr =
|
|
linear2_bias == nullptr ? nullptr : linear2_bias->data<T>();
|
|
|
|
if (pre_layer_norm) {
|
|
pre_layernorm_helper.LayerNorm(dev_ctx,
|
|
x.data<T>(),
|
|
ln1_scale_ptr,
|
|
ln1_bias_ptr,
|
|
ln1_out->data<T>(),
|
|
ln1_mean->data<U>(),
|
|
ln1_variance->data<U>());
|
|
in = ln1_out;
|
|
}
|
|
MatMul<T, Context>(dev_ctx, *in, linear1_weight, linear1_out);
|
|
fused_act_dropout_helper.DropoutActBias(dev_ctx,
|
|
linear1_out->data<T>(),
|
|
linear1_bias_ptr,
|
|
act_method,
|
|
dropout1_out->data<T>(),
|
|
dropout1_mask->data<uint8_t>());
|
|
DenseTensor linear2_out;
|
|
linear2_out.Resize({bsz_seq, d_model});
|
|
dev_ctx.template Alloc<T>(&linear2_out, linear2_out.numel() * sizeof(T));
|
|
MatMul<T, Context>(dev_ctx, *dropout1_out, linear2_weight, &linear2_out);
|
|
|
|
// tensor model parallel
|
|
phi::fusion::AllReduce<T>(linear2_out, ring_id, dev_ctx);
|
|
|
|
const T* residual_ptr = add_residual ? x.data<T>() : nullptr;
|
|
if (!pre_layer_norm) {
|
|
// TODO(Xreki): support post layer_norm case when add_residual is false.
|
|
PADDLE_ENFORCE_EQ(add_residual,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Attribute add_residual is expected to be true "
|
|
"when pre_layer_norm is false."));
|
|
|
|
fused_dropout_layernorm_helper.LayernormResidualDropoutBias(
|
|
dev_ctx,
|
|
linear2_out.data<T>(),
|
|
residual_ptr,
|
|
linear2_bias_ptr,
|
|
ln2_scale_ptr,
|
|
ln2_bias_ptr,
|
|
dropout2_out->data<T>(),
|
|
dropout2_mask->data<uint8_t>(),
|
|
out->data<T>(),
|
|
ln2_mean->data<U>(),
|
|
ln2_variance->data<U>());
|
|
} else {
|
|
fused_dropout_layernorm_helper.ResidualDropoutBias(
|
|
dev_ctx,
|
|
linear2_out.data<T>(),
|
|
residual_ptr,
|
|
linear2_bias_ptr,
|
|
out->data<T>(),
|
|
dropout2_mask->data<uint8_t>());
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FusedFeedForwardKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& dropout1_seed,
|
|
const optional<DenseTensor>& dropout2_seed,
|
|
const DenseTensor& linear1_weight,
|
|
const optional<DenseTensor>& linear1_bias,
|
|
const DenseTensor& linear2_weight,
|
|
const optional<DenseTensor>& linear2_bias,
|
|
const optional<DenseTensor>& ln1_scale,
|
|
const optional<DenseTensor>& ln1_bias,
|
|
const optional<DenseTensor>& ln2_scale,
|
|
const optional<DenseTensor>& ln2_bias,
|
|
bool pre_layer_norm,
|
|
float ln1_epsilon,
|
|
float ln2_epsilon,
|
|
const std::string& act_method,
|
|
float dropout1_prob,
|
|
float dropout2_prob,
|
|
const std::string& dropout1_implementation,
|
|
const std::string& dropout2_implementation,
|
|
bool is_test,
|
|
bool dropout1_fix_seed,
|
|
bool dropout2_fix_seed,
|
|
int dropout1_seed_val,
|
|
int dropout2_seed_val,
|
|
bool add_residual,
|
|
int ring_id,
|
|
DenseTensor* out,
|
|
DenseTensor* dropout1_mask,
|
|
DenseTensor* dropout2_mask,
|
|
DenseTensor* ln1_mean,
|
|
DenseTensor* ln1_variance,
|
|
DenseTensor* ln2_mean,
|
|
DenseTensor* ln2_variance,
|
|
DenseTensor* linear1_out,
|
|
DenseTensor* ln1_out,
|
|
DenseTensor* dropout1_out,
|
|
DenseTensor* dropout2_out) {
|
|
auto* x_ptr = &x;
|
|
auto* linear1_weight_ptr = &linear1_weight;
|
|
auto* linear1_bias_ptr = linear1_bias.get_ptr();
|
|
auto* linear2_weight_ptr = &linear2_weight;
|
|
auto* linear2_bias_ptr = linear2_bias.get_ptr();
|
|
|
|
auto* ln1_scale_ptr = pre_layer_norm ? ln1_scale.get_ptr() : nullptr;
|
|
auto* ln1_bias_ptr = pre_layer_norm ? ln1_bias.get_ptr() : nullptr;
|
|
auto* ln2_scale_ptr = !pre_layer_norm ? ln2_scale.get_ptr() : nullptr;
|
|
auto* ln2_bias_ptr = !pre_layer_norm ? ln2_bias.get_ptr() : nullptr;
|
|
|
|
if (!pre_layer_norm) {
|
|
ln1_mean = nullptr;
|
|
ln1_variance = nullptr;
|
|
ln1_out = nullptr;
|
|
} else {
|
|
ln2_mean = nullptr;
|
|
ln2_variance = nullptr;
|
|
}
|
|
|
|
bool is_upscale_in_train1 = dropout1_implementation == "upscale_in_train";
|
|
bool is_upscale_in_train2 = dropout2_implementation == "upscale_in_train";
|
|
auto* dropout1_seed_ptr = dropout1_seed.get_ptr();
|
|
auto* dropout2_seed_ptr = dropout2_seed.get_ptr();
|
|
|
|
fusion::DropoutParam dropout_param1(dropout1_fix_seed,
|
|
0,
|
|
is_test,
|
|
is_upscale_in_train1,
|
|
dropout1_prob,
|
|
dropout1_seed_ptr,
|
|
dropout1_seed_val);
|
|
fusion::DropoutParam dropout_param2(dropout2_fix_seed,
|
|
0,
|
|
is_test,
|
|
is_upscale_in_train2,
|
|
dropout2_prob,
|
|
dropout2_seed_ptr,
|
|
dropout2_seed_val);
|
|
|
|
using U = funcs::LayerNormParamType<T>;
|
|
dev_ctx.template Alloc<T>(out, out->numel() * sizeof(T));
|
|
dev_ctx.template Alloc<uint8_t>(dropout1_mask,
|
|
dropout1_mask->numel() * sizeof(uint8_t));
|
|
dev_ctx.template Alloc<uint8_t>(dropout2_mask,
|
|
dropout2_mask->numel() * sizeof(uint8_t));
|
|
if (pre_layer_norm) {
|
|
dev_ctx.template Alloc<U>(ln1_mean, ln1_mean->numel() * sizeof(U));
|
|
dev_ctx.template Alloc<U>(ln1_variance, ln1_variance->numel() * sizeof(U));
|
|
dev_ctx.template Alloc<T>(ln1_out, ln1_out->numel() * sizeof(T));
|
|
} else {
|
|
dev_ctx.template Alloc<U>(ln2_mean, ln2_mean->numel() * sizeof(U));
|
|
dev_ctx.template Alloc<U>(ln2_variance, ln2_variance->numel() * sizeof(U));
|
|
}
|
|
|
|
dev_ctx.template Alloc<T>(linear1_out, linear1_out->numel() * sizeof(T));
|
|
dev_ctx.template Alloc<T>(dropout1_out, dropout1_out->numel() * sizeof(T));
|
|
dev_ctx.template Alloc<T>(dropout2_out, dropout2_out->numel() * sizeof(T));
|
|
|
|
if (out->numel() == 0) {
|
|
return;
|
|
}
|
|
|
|
auto x_dim = x_ptr->dims();
|
|
auto mat_dim_x =
|
|
funcs::CreateMatrixDescriptor(phi::RowMatrixFromVector(x_dim), 0, false);
|
|
|
|
auto dim = linear1_weight_ptr->dims();
|
|
int d_model = dim[0];
|
|
int dim_feedforward = dim[dim.size() - 1];
|
|
int bsz_seq = mat_dim_x.batch_size_ * mat_dim_x.height_;
|
|
|
|
phi::fusion::FFN<T, Context>(dev_ctx,
|
|
x,
|
|
linear1_weight,
|
|
linear1_bias_ptr,
|
|
linear2_weight,
|
|
linear2_bias_ptr,
|
|
ln1_scale_ptr,
|
|
ln1_bias_ptr,
|
|
ln2_scale_ptr,
|
|
ln2_bias_ptr,
|
|
out,
|
|
dropout1_mask,
|
|
dropout2_mask,
|
|
ln1_mean,
|
|
ln1_variance,
|
|
ln2_mean,
|
|
ln2_variance,
|
|
linear1_out,
|
|
ln1_out,
|
|
dropout1_out,
|
|
dropout2_out,
|
|
bsz_seq,
|
|
d_model,
|
|
dim_feedforward,
|
|
act_method,
|
|
pre_layer_norm,
|
|
ln1_epsilon,
|
|
ln2_epsilon,
|
|
add_residual,
|
|
ring_id,
|
|
dropout_param1,
|
|
dropout_param2);
|
|
}
|
|
|
|
} // namespace fusion
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(fused_feedforward,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::fusion::FusedFeedForwardKernel,
|
|
float,
|
|
double,
|
|
phi::float16) {
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::UINT8);
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::UINT8);
|
|
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
|
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
|
|
}
|
|
}
|