// 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. #include "paddle/phi/kernels/fused_attention_grad_kernel.h" #include "paddle/phi/backends/xpu/enforce_xpu.h" #include "paddle/phi/backends/xpu/xpu_context.h" #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/kernels/xpu/xpu_api_wrapper.h" #include "paddle/phi/kernels/xpu/xpu_fused_common_function.h" namespace phi { template void FusedAttentionGradKernel( const Context &dev_ctx, const DenseTensor &out_grad, const DenseTensor &x, const DenseTensor &qkv_weight, const optional &qkv_bias, const optional &qkv_bias_out, const optional &src_mask, const optional &src_mask_out, const DenseTensor &out_linear_weight, const optional &out_linear_bias, const optional &ln_scale, const optional &ln_bias, const optional &ln_scale_2, const optional &ln_bias_2, const optional &ln_out, const optional &ln_mean, const optional &ln_var, const optional &ln_mean_2, const optional &ln_var_2, const optional &bias_dropout_residual_out, const DenseTensor &qkv_out, const DenseTensor &transpose_out_2, const DenseTensor &qk_out, const DenseTensor &qktv_out, const DenseTensor &softmax_out, const DenseTensor &attn_dropout_mask, const DenseTensor &attn_dropout_out, const DenseTensor &fmha_out, const DenseTensor &out_linear_out, const DenseTensor &dropout_mask_out, int num_heads_, // unused bool transpose_qkv_wb, bool pre_layer_norm, float epsilon, float attn_dropout_rate, bool is_test, bool attn_dropout_fix_seed, int attn_dropout_seed, const std::string &attn_dropout_implementation, float dropout_rate, bool dropout_fix_seed, int dropout_seed, const std::string &dropout_implementation, float ln_epsilon, bool add_residual, int ring_id, DenseTensor *qkv_bias_grad, DenseTensor *qkv_bias_out_grad, DenseTensor *src_mask_out_grad, DenseTensor *out_linear_bias_grad, DenseTensor *ln_scale_grad, DenseTensor *ln_bias_grad, DenseTensor *ln_scale_2_grad, DenseTensor *ln_bias_2_grad, DenseTensor *x_grad, DenseTensor *qkv_weight_grad, DenseTensor *out_linear_weight_grad, DenseTensor *ln_out_grad, DenseTensor *bias_dropout_residual_out_grad, DenseTensor *qkv_out_grad, DenseTensor *qktv_out_grad, DenseTensor *transpose_out_2_grad, DenseTensor *qk_out_grad, DenseTensor *softmax_out_grad, DenseTensor *attn_dropout_out_grad, DenseTensor *fmha_out_grad, DenseTensor *out_linear_out_grad) { using XPUTypeT = typename XPUTypeTrait::Type; bool is_upscale_in_train_1 = (attn_dropout_implementation == "upscale_in_train"); const DenseTensor *seed_1 = nullptr; phi::XPUDropoutParam attn_dropout_param; attn_dropout_param.initXPUDropoutParam(attn_dropout_rate, is_upscale_in_train_1, is_test, attn_dropout_fix_seed, seed_1, attn_dropout_seed); phi::XPUDropoutParam dropout_param; dropout_param.initXPUDropoutParam(dropout_rate, is_upscale_in_train_1, is_test, dropout_fix_seed, seed_1, dropout_seed); const auto input_x_dims = x.dims(); const auto qkv_w_dims = qkv_weight.dims(); int64_t batch_size = input_x_dims[0]; int64_t seq_len = input_x_dims[1]; if (batch_size == 0 || seq_len == 0) { if (qkv_bias_grad) dev_ctx.template Alloc(qkv_bias_grad); if (qkv_bias_out_grad) dev_ctx.template Alloc(qkv_bias_out_grad); if (src_mask_out_grad) dev_ctx.template Alloc(src_mask_out_grad); if (out_linear_bias_grad) dev_ctx.template Alloc(out_linear_bias_grad); if (ln_scale_grad) dev_ctx.template Alloc(ln_scale_grad); if (ln_bias_grad) dev_ctx.template Alloc(ln_bias_grad); if (ln_scale_2_grad) dev_ctx.template Alloc(ln_scale_2_grad); if (ln_bias_2_grad) dev_ctx.template Alloc(ln_bias_2_grad); if (x_grad) dev_ctx.template Alloc(x_grad); if (qkv_weight_grad) dev_ctx.template Alloc(qkv_weight_grad); if (out_linear_weight_grad) dev_ctx.template Alloc(out_linear_weight_grad); if (ln_out_grad) dev_ctx.template Alloc(ln_out_grad); if (bias_dropout_residual_out_grad) dev_ctx.template Alloc(bias_dropout_residual_out_grad); if (qkv_out_grad) dev_ctx.template Alloc(qkv_out_grad); if (qktv_out_grad) dev_ctx.template Alloc(qktv_out_grad); if (transpose_out_2_grad) dev_ctx.template Alloc(transpose_out_2_grad); if (qk_out_grad) dev_ctx.template Alloc(qk_out_grad); if (softmax_out_grad) dev_ctx.template Alloc(softmax_out_grad); if (attn_dropout_out_grad) dev_ctx.template Alloc(attn_dropout_out_grad); if (fmha_out_grad) dev_ctx.template Alloc(fmha_out_grad); if (out_linear_out_grad) dev_ctx.template Alloc(out_linear_out_grad); return; } // get inputs. const XPUTypeT *d_y_ptr = reinterpret_cast(out_grad.data()); // 前向必要参数 const XPUTypeT *input_x_ptr = reinterpret_cast(x.data()); const XPUTypeT *qkv_transpose_out_ptr = reinterpret_cast(transpose_out_2.data()); const XPUTypeT *qkv_weight_ptr = reinterpret_cast(qkv_weight.data()); const XPUTypeT *softmax_out_ptr = reinterpret_cast(softmax_out.data()); const XPUTypeT *attn_dropout_out_ptr = reinterpret_cast(attn_dropout_out.data()); const XPUTypeT *attn_dropout_mask_ptr = reinterpret_cast(attn_dropout_mask.data()); const XPUTypeT *fmha_out_ptr = reinterpret_cast(fmha_out.data()); const XPUTypeT *out_linear_weight_ptr = reinterpret_cast(out_linear_weight.data()); const XPUTypeT *dropout_mask_out_ptr = reinterpret_cast(dropout_mask_out.data()); // 需要计算的梯度 auto *d_qkv_weight = qkv_weight_grad; XPUTypeT *d_qkv_weight_ptr = reinterpret_cast(dev_ctx.template Alloc(d_qkv_weight)); auto *d_qkv_bias = qkv_bias_grad; XPUTypeT *d_qkv_bias_ptr = reinterpret_cast(dev_ctx.template Alloc(d_qkv_bias)); auto *d_out_linear_weight = out_linear_weight_grad; XPUTypeT *d_out_linear_weight_ptr = reinterpret_cast( dev_ctx.template Alloc(d_out_linear_weight)); auto *d_out_linear_bias = out_linear_bias_grad; XPUTypeT *d_out_linear_bias_ptr = reinterpret_cast( dev_ctx.template Alloc(d_out_linear_bias)); // 有可能需要 auto *d_src_mask_out = src_mask_out_grad; XPUTypeT *d_src_mask_out_ptr = (d_src_mask_out == nullptr) ? (nullptr) : (reinterpret_cast( dev_ctx.template Alloc(d_src_mask_out))); // 输出 dx auto *d_x = x_grad; XPUTypeT *d_x_ptr = reinterpret_cast(dev_ctx.template Alloc(d_x)); const DenseTensor *ln_out_p = ln_out.get_ptr(); const DenseTensor *bias_dropout_residual_out_p = bias_dropout_residual_out.get_ptr(); const DenseTensor *ln_scale_p = nullptr; const DenseTensor *ln_mean_p = nullptr; const DenseTensor *ln_var_p = nullptr; DenseTensor *d_ln_scale = nullptr; DenseTensor *d_ln_bias = nullptr; const XPUTypeT *ln_out_ptr = NULL; const float *ln_scale_ptr = NULL; const float *ln_mean_ptr = NULL; const float *ln_var_ptr = NULL; const XPUTypeT *bias_dropout_residual_out_ptr = NULL; float *d_ln_scale_ptr = nullptr; float *d_ln_bias_ptr = nullptr; if (pre_layer_norm) { ln_out_ptr = reinterpret_cast(ln_out_p->data()); ln_scale_p = ln_scale.get_ptr(); ln_mean_p = ln_mean.get_ptr(); ln_var_p = ln_var.get_ptr(); d_ln_scale = ln_scale_grad; d_ln_bias = ln_bias_grad; } else { ln_scale_p = ln_scale_2.get_ptr(); ln_mean_p = ln_mean_2.get_ptr(); ln_var_p = ln_var_2.get_ptr(); epsilon = ln_epsilon; d_ln_scale = ln_scale_2_grad; d_ln_bias = ln_bias_2_grad; bias_dropout_residual_out_ptr = reinterpret_cast( bias_dropout_residual_out_p->data()); } ln_scale_ptr = ln_scale_p->data(); ln_mean_ptr = ln_mean_p->data(); ln_var_ptr = ln_var_p->data(); d_ln_scale_ptr = dev_ctx.template Alloc(d_ln_scale); d_ln_bias_ptr = dev_ctx.template Alloc(d_ln_bias); int64_t embed_dims = input_x_dims[2]; int64_t num_heads = qkv_w_dims[1]; int64_t head_dims = qkv_w_dims[2]; xpu::Context *xpu_ctx = dev_ctx.x_context(); xpu::ctx_guard RAII_GUARD(xpu_ctx); int r = 0; // int l3_total_size = xpu_ctx->_l3_mgr.get_size(); XPUTypeT *d_ln_grad_ptr = NULL; // dx5 [batch_size, seq_len, hidden] XPUTypeT *d_dropout_grad_ptr = NULL; // dx5 [batch_size, seq_len, hidden] XPUTypeT *d_fmha_out_ptr = NULL; // d_fmha_out [batch_size, seq_len, num_heads, head_dims] XPUTypeT *d_fmha_out_transpose_tmp_ptr = NULL; // d_fmha_out_transpose [batch_size, seq_len, num_heads, // head_dims] XPUTypeT *d_qk_ptr = NULL; // d_qk_ptr[batch_size, num_heads, seq_len, seq_len] XPUTypeT *d_combination_qkv_ptr = NULL; // d_combination_qkv_ptr[3, batch_size, num_heads, seq_len, // head_dims] XPUTypeT *d_transpose_qkv_ptr = NULL; // dx2 [batch_size, seq_len, 3, num_heads, head_dims] XPUTypeT *d_last_layernorm_grad_ptr = NULL; // d_layer_out [batch_size, seq_len, embed_dims] const XPUTypeT *dy_input_ptr = d_y_ptr; d_ln_grad_ptr = RAII_GUARD.alloc(batch_size * seq_len * embed_dims); d_dropout_grad_ptr = RAII_GUARD.alloc_l3_or_gm(batch_size * seq_len * embed_dims); d_fmha_out_ptr = RAII_GUARD.alloc_l3_or_gm(batch_size * seq_len * num_heads * head_dims); d_combination_qkv_ptr = RAII_GUARD.alloc(batch_size * seq_len * embed_dims * 3); d_transpose_qkv_ptr = RAII_GUARD.alloc_l3_or_gm( batch_size * seq_len * embed_dims * 3); d_fmha_out_transpose_tmp_ptr = RAII_GUARD.alloc_l3_or_gm(batch_size * seq_len * embed_dims); d_qk_ptr = RAII_GUARD.alloc_l3_or_gm(batch_size * seq_len * seq_len * num_heads); d_last_layernorm_grad_ptr = RAII_GUARD.alloc_l3_or_gm(batch_size * seq_len * embed_dims); if (pre_layer_norm == false) { r = xpu::layer_norm_grad(xpu_ctx, bias_dropout_residual_out_ptr, d_y_ptr, d_ln_grad_ptr, batch_size * seq_len, embed_dims, epsilon, ln_scale_ptr, ln_mean_ptr, ln_var_ptr, d_ln_scale_ptr, d_ln_bias_ptr); PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm_grad"); dy_input_ptr = d_ln_grad_ptr; } // dropout_grad DropoutGrad(xpu_ctx, dy_input_ptr, dropout_mask_out_ptr, d_dropout_grad_ptr, dropout_param, batch_size * num_heads * seq_len * head_dims); // linear_out phi::XpuFcInfo linear_fc_info; linear_fc_info.InitFcInfo(0, batch_size * seq_len, embed_dims, embed_dims, false, false, nullptr, nullptr, nullptr); const XPUTypeT *a_1 = reinterpret_cast(NULL); const XPUTypeT *b_1 = reinterpret_cast(NULL); const XPUTypeT *a_2 = reinterpret_cast(NULL); const XPUTypeT *b_2 = reinterpret_cast(NULL); XPUTypeT *c_1 = d_fmha_out_ptr; XPUTypeT *c_2 = d_out_linear_weight_ptr; phi::XpuFcInfo info_dfmha; phi::XpuFcInfo info_dlinear_w; std::tuple fc_info = phi::MatmulGradFcInfo(xpu_ctx, &RAII_GUARD, linear_fc_info, false, false, fmha_out_ptr, out_linear_weight_ptr, d_dropout_grad_ptr); std::tie(info_dfmha, info_dlinear_w, a_1, b_1, a_2, b_2) = fc_info; phi::MatMulXPUFunction( xpu_ctx, a_2, b_2, c_2, info_dlinear_w, 1.0f, 0.f, true); phi::MatMulXPUFunction( xpu_ctx, a_1, b_1, c_1, info_dfmha, 1.0f, 0.f, true); // dlinear_bias r = xpu::reduce_sum(xpu_ctx, d_dropout_grad_ptr, d_out_linear_bias_ptr, {(int64_t)batch_size * seq_len, (int64_t)embed_dims}, {0LL}); PADDLE_ENFORCE_XDNN_SUCCESS(r, "reduce_sum"); { int64_t qkv_size = batch_size * seq_len * num_heads * head_dims; const XPUTypeT *q_out_ptr = qkv_transpose_out_ptr; const XPUTypeT *k_out_ptr = q_out_ptr + qkv_size; const XPUTypeT *v_out_ptr = k_out_ptr + qkv_size; XPUTypeT *d_q_out_ptr = d_combination_qkv_ptr; XPUTypeT *d_k_out_ptr = d_q_out_ptr + qkv_size; XPUTypeT *d_v_out_ptr = d_k_out_ptr + qkv_size; r = xpu::transpose(xpu_ctx, d_fmha_out_ptr, d_fmha_out_transpose_tmp_ptr, {(int64_t)batch_size, (int64_t)seq_len, (int64_t)num_heads, (int64_t)head_dims}, {0LL, 2LL, 1LL, 3LL}); PADDLE_ENFORCE_XDNN_SUCCESS(r, "transpose"); phi::XpuFcInfo qktv_fc_info; qktv_fc_info.InitFcInfo(batch_size * num_heads, seq_len, head_dims, seq_len, false, false, nullptr, nullptr, nullptr); const XPUTypeT *a_1 = reinterpret_cast(NULL); const XPUTypeT *b_1 = reinterpret_cast(NULL); const XPUTypeT *a_2 = reinterpret_cast(NULL); const XPUTypeT *b_2 = reinterpret_cast(NULL); XPUTypeT *c_1 = d_qk_ptr; XPUTypeT *c_2 = d_v_out_ptr; phi::XpuFcInfo info_d_qk; phi::XpuFcInfo info_d_v; std::tuple fc_info = phi::MatmulGradFcInfo(xpu_ctx, &RAII_GUARD, qktv_fc_info, false, false, attn_dropout_out_ptr, v_out_ptr, d_fmha_out_transpose_tmp_ptr); std::tie(info_d_qk, info_d_v, a_1, b_1, a_2, b_2) = fc_info; phi::MatMulXPUFunction( xpu_ctx, a_1, b_1, c_1, info_d_qk, 1.0f, 0.f, true); phi::MatMulXPUFunction( xpu_ctx, a_2, b_2, c_2, info_d_v, 1.0f, 0.f, true); DropoutGrad(xpu_ctx, d_qk_ptr, attn_dropout_mask_ptr, d_qk_ptr, attn_dropout_param, batch_size * seq_len * seq_len * num_heads); r = xpu::softmax_grad(xpu_ctx, softmax_out_ptr, d_qk_ptr, d_qk_ptr, {batch_size, num_heads, seq_len, seq_len}, 3); PADDLE_ENFORCE_XDNN_SUCCESS(r, "softmax_grad"); if (d_src_mask_out_ptr) { r = xpu::copy(xpu_ctx, d_qk_ptr, d_src_mask_out_ptr, batch_size * seq_len * seq_len * num_heads); PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy"); } phi::XpuFcInfo qk_fc_info; qk_fc_info.InitFcInfo(batch_size * num_heads, seq_len, seq_len, head_dims, false, true, nullptr, nullptr, nullptr); a_1 = reinterpret_cast(NULL); b_1 = reinterpret_cast(NULL); a_2 = reinterpret_cast(NULL); b_2 = reinterpret_cast(NULL); c_1 = d_q_out_ptr; c_2 = d_k_out_ptr; phi::XpuFcInfo info_d_q; phi::XpuFcInfo info_d_k; fc_info = phi::MatmulGradFcInfo(xpu_ctx, &RAII_GUARD, qk_fc_info, false, true, q_out_ptr, k_out_ptr, d_qk_ptr); std::tie(info_d_q, info_d_k, a_1, b_1, a_2, b_2) = fc_info; phi::MatMulXPUFunction( xpu_ctx, a_1, b_1, c_1, info_d_q, 1.0f / sqrt(head_dims), 0.f, true); phi::MatMulXPUFunction( xpu_ctx, a_2, b_2, c_2, info_d_k, 1.0f, 0.f, true); } // r = xpu::transpose(xpu_ctx, d_combination_qkv_ptr, d_transpose_qkv_ptr, {3, batch_size, num_heads, seq_len, head_dims}, {1, 3, 0, 2, 4}); PADDLE_ENFORCE_XDNN_SUCCESS(r, "transpose"); // dx and d_qkv_w phi::XpuFcInfo qkv_fc_info; qkv_fc_info.InitFcInfo(0, batch_size * seq_len, 3 * num_heads * head_dims, embed_dims, false, true, nullptr, nullptr, nullptr); a_1 = reinterpret_cast(NULL); b_1 = reinterpret_cast(NULL); a_2 = reinterpret_cast(NULL); b_2 = reinterpret_cast(NULL); c_1 = (pre_layer_norm == true) ? d_last_layernorm_grad_ptr : d_x_ptr; c_2 = d_qkv_weight_ptr; phi::XpuFcInfo info_d_x; phi::XpuFcInfo info_d_qkv_w; const XPUTypeT *use_calc_input_x_ptr = (pre_layer_norm == true) ? ln_out_ptr : input_x_ptr; fc_info = phi::MatmulGradFcInfo(xpu_ctx, &RAII_GUARD, qkv_fc_info, false, true, use_calc_input_x_ptr, qkv_weight_ptr, d_transpose_qkv_ptr); std::tie(info_d_x, info_d_qkv_w, a_1, b_1, a_2, b_2) = fc_info; phi::MatMulXPUFunction( xpu_ctx, a_1, b_1, c_1, info_d_x, 1.0f, 0.f, true); phi::MatMulXPUFunction( xpu_ctx, a_2, b_2, c_2, info_d_qkv_w, 1.0f, 0.f, true); // d_qkv_bias r = xpu::reduce_sum(xpu_ctx, d_transpose_qkv_ptr, d_qkv_bias_ptr, {(int64_t)batch_size * seq_len, 3LL * embed_dims}, {0LL}); PADDLE_ENFORCE_XDNN_SUCCESS(r, "reduce_sum"); if (pre_layer_norm) { r = xpu::layer_norm_grad(xpu_ctx, input_x_ptr, c_1, d_x_ptr, batch_size * seq_len, embed_dims, epsilon, ln_scale_ptr, ln_mean_ptr, ln_var_ptr, d_ln_scale_ptr, d_ln_bias_ptr); PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm_grad"); } // add rediaus dy r = xpu::add(xpu_ctx, dy_input_ptr, d_x_ptr, d_x_ptr, batch_size * seq_len * embed_dims); PADDLE_ENFORCE_XDNN_SUCCESS(r, "add"); } } // namespace phi PD_REGISTER_KERNEL(fused_attention_grad, XPU, ALL_LAYOUT, phi::FusedAttentionGradKernel, float, phi::float16) { kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32); kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32); kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32); kernel->OutputAt(7).SetDataType(phi::DataType::FLOAT32); }