From fdcdbec5330efbe850d648f9444d60ce7881f4dc Mon Sep 17 00:00:00 2001 From: crystal <62974595+Zjq9409@users.noreply.github.com> Date: Mon, 30 May 2022 17:56:51 +0800 Subject: [PATCH] Implement fused_gate_attention operator for AlphaFold. (#42018) --- paddle/fluid/operators/fused/CMakeLists.txt | 4 +- paddle/fluid/operators/fused/attn_gemm.h | 164 +++-- paddle/fluid/operators/fused/fmha_ref.h | 4 +- .../operators/fused/fused_gate_attention.h | 647 ++++++++++++++++++ .../fused/fused_gate_attention_op.cc | 317 +++++++++ .../fused/fused_gate_attention_op.cu | 488 +++++++++++++ paddle/fluid/platform/device/gpu/gpu_info.cc | 4 +- paddle/fluid/pybind/op_function_generator.h | 7 + paddle/phi/kernels/gpudnn/softmax_gpudnn.h | 42 +- .../fluid/tests/unittests/CMakeLists.txt | 1 + .../unittests/test_fused_gate_attention_op.py | 252 +++++++ 11 files changed, 1821 insertions(+), 109 deletions(-) create mode 100644 paddle/fluid/operators/fused/fused_gate_attention.h create mode 100644 paddle/fluid/operators/fused/fused_gate_attention_op.cc create mode 100644 paddle/fluid/operators/fused/fused_gate_attention_op.cu create mode 100644 python/paddle/fluid/tests/unittests/test_fused_gate_attention_op.py diff --git a/paddle/fluid/operators/fused/CMakeLists.txt b/paddle/fluid/operators/fused/CMakeLists.txt index 03351dbca09..a86d26bcd58 100644 --- a/paddle/fluid/operators/fused/CMakeLists.txt +++ b/paddle/fluid/operators/fused/CMakeLists.txt @@ -23,7 +23,8 @@ register_operators(EXCLUDES fused_feedforward_op fused_multi_transformer_op resnet_unit_op - fused_gemm_epilogue_op) + fused_gemm_epilogue_op + fused_gate_attention_op) # fusion_gru_op does not have CUDA kernel op_library(fusion_gru_op) @@ -58,6 +59,7 @@ if (WITH_GPU OR WITH_ROCM) op_library(yolo_box_head_op) op_library(yolo_box_post_op) op_library(fused_embedding_eltwise_layernorm_op) + op_library(fused_gate_attention_op) # fusion_group if(NOT APPLE AND NOT WIN32) op_library(fusion_group_op DEPS device_code) diff --git a/paddle/fluid/operators/fused/attn_gemm.h b/paddle/fluid/operators/fused/attn_gemm.h index 9542f0742ea..304aad16ad0 100644 --- a/paddle/fluid/operators/fused/attn_gemm.h +++ b/paddle/fluid/operators/fused/attn_gemm.h @@ -1,8 +1,11 @@ /* Copyright (c) 2021 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. @@ -13,6 +16,7 @@ limitations under the License. */ #include "paddle/fluid/platform/float16.h" #include "paddle/phi/kernels/funcs/blas/blas.h" +#include "paddle/phi/kernels/funcs/elementwise_functor.h" #include "paddle/fluid/operators/kernel_primitives/kernel_primitives.h" #include "paddle/fluid/operators/reduce_ops/reduce_op.cu.h" @@ -21,6 +25,8 @@ limitations under the License. */ namespace paddle { namespace operators { + +using Tensor = framework::Tensor; // support gemm-nt and gemm-nn, which is used in fused_attention_op. template class AttnMatMul { @@ -45,31 +51,21 @@ class AttnMatMul { framework::Tensor* bias_out) { // Note: for blas.GEMM API in Paddle, it treats all inputs as row-major. // here: (transa, transb): nt, input * weight. - CBLAS_TRANSPOSE transA = CblasNoTrans; - CBLAS_TRANSPOSE transB = CblasNoTrans; - if (transA_) { - transA = CblasTrans; - } - if (transB_) { - transB = CblasTrans; - } + CBLAS_TRANSPOSE transA = transA_ ? CblasTrans : CblasNoTrans; + CBLAS_TRANSPOSE transB = transB_ ? CblasTrans : CblasNoTrans; T alpha = static_cast(1.0); T beta = static_cast(0.0); - // here: (m, n, k) = bsz_seq, output_size, input_size, (input, weight, out) + // (m, n, k) = bsz_seq, output_size, input_size, (input, weight, out) auto blas = phi::funcs::GetBlas(dev_ctx_); blas.GEMM(transA, transB, bsz_seq_, output_size_, input_size_, alpha, input->data(), weight->data(), beta, output->data()); if (compute_bias_) { - // compute output + bias - std::vector ins; - std::vector outs; - ins.emplace_back(output); - ins.emplace_back(bias); - outs.emplace_back(bias_out); - int elewise_add_axis = -1; + // bias_out = output + bias + std::vector ins = {output, bias}; + std::vector outs = {bias_out}; phi::funcs::BroadcastKernel( - dev_ctx_, ins, &outs, elewise_add_axis, phi::funcs::AddFunctor()); + dev_ctx_, ins, &outs, -1, phi::funcs::AddFunctor()); } } @@ -77,82 +73,71 @@ class AttnMatMul { const framework::Tensor* weight, const framework::Tensor* d_output, framework::Tensor* d_input, framework::Tensor* d_weight, - framework::Tensor* d_bias) { + framework::Tensor* d_bias, bool use_addto = false) { T alpha = static_cast(1.0); - T beta = static_cast(0.0); - auto blas = phi::funcs::GetBlas(dev_ctx_); - - CBLAS_TRANSPOSE dB_transA = CblasNoTrans; - CBLAS_TRANSPOSE dB_transB = CblasNoTrans; - CBLAS_TRANSPOSE dA_transA = CblasNoTrans; - CBLAS_TRANSPOSE dA_transB = CblasNoTrans; - int dB_m = 1; - int dB_n = 1; - int dB_k = 1; - int dA_m = 1; - int dA_n = 1; - int dA_k = 1; - - T* dB_input_1_ptr = nullptr; - T* dB_input_2_ptr = nullptr; - T* dB_output_ptr = d_weight->data(); - - T* dA_input_1_ptr = nullptr; - T* dA_input_2_ptr = nullptr; - T* dA_output_ptr = d_input->data(); + T beta_dA = use_addto ? static_cast(1.0) : static_cast(0.0); + T beta_dB = static_cast(0.0); + auto blas = phi::funcs::GetBlas(dev_ctx_); if (!transA_) { - // fw: gemm-nt + // forward: gemm-nt if (transB_) { - // bw: gemm-tn, dB = (dC)^t * A - dB_transA = CblasTrans; - dB_transB = CblasNoTrans; - dB_m = output_size_; - dB_n = input_size_; - dB_k = bsz_seq_; - - // bw: gemm-nn, dA = dC * B - dA_transA = CblasNoTrans; - dA_transB = CblasNoTrans; - dA_m = bsz_seq_; - dA_n = input_size_; - dA_k = output_size_; - - blas.GEMM(dB_transA, dB_transB, dB_m, dB_n, dB_k, alpha, - d_output->data(), input->data(), beta, dB_output_ptr); - blas.GEMM(dA_transA, dA_transB, dA_m, dA_n, dA_k, alpha, - d_output->data(), weight->data(), beta, dA_output_ptr); + // backward: gemm-tn, dB = (dC)^T * A + if (d_weight) { + int dB_m = output_size_; + int dB_n = input_size_; + int dB_k = bsz_seq_; + + T* dB_output_ptr = d_weight->data(); + blas.GEMM(CblasTrans, CblasNoTrans, dB_m, dB_n, dB_k, alpha, + d_output->data(), input->data(), beta_dB, + dB_output_ptr); + } + + // backward: gemm-nn, dA = dC * B + if (d_input) { + int dA_m = bsz_seq_; + int dA_n = input_size_; + int dA_k = output_size_; + + T* dA_output_ptr = d_input->data(); + blas.GEMM(CblasNoTrans, CblasNoTrans, dA_m, dA_n, dA_k, alpha, + d_output->data(), weight->data(), beta_dA, + dA_output_ptr); + } } else { // fw: gemm-nn - // bw: gemm-tn, dB = A^t * dC - dB_transA = CblasTrans; - dB_transB = CblasNoTrans; - dB_m = input_size_; - dB_n = output_size_; - dB_k = bsz_seq_; - - // bw: gemm-nt, dA = dC * B^t - dA_transA = CblasNoTrans; - dA_transB = CblasTrans; - dA_m = bsz_seq_; - dA_n = input_size_; - dA_k = output_size_; - - blas.GEMM(dB_transA, dB_transB, dB_m, dB_n, dB_k, alpha, - input->data(), d_output->data(), beta, dB_output_ptr); - blas.GEMM(dA_transA, dA_transB, dA_m, dA_n, dA_k, alpha, - d_output->data(), weight->data(), beta, dA_output_ptr); + // backward: gemm-tn, dB = A^T * dC + if (d_weight) { + int dB_m = input_size_; + int dB_n = output_size_; + int dB_k = bsz_seq_; + + T* dB_output_ptr = d_weight->data(); + blas.GEMM(CblasTrans, CblasNoTrans, dB_m, dB_n, dB_k, alpha, + input->data(), d_output->data(), beta_dB, + dB_output_ptr); + } + + // backward: gemm-nt, dA = dC * B^T + if (d_input) { + int dA_m = bsz_seq_; + int dA_n = input_size_; + int dA_k = output_size_; + + T* dA_output_ptr = d_input->data(); + blas.GEMM(CblasNoTrans, CblasTrans, dA_m, dA_n, dA_k, alpha, + d_output->data(), weight->data(), beta_dA, + dA_output_ptr); + } } - } else if (transB_) { - PADDLE_THROW(platform::errors::InvalidArgument( - "AttnMatMul wrapper do not support (transA=T, transB=T)" - "parameters.")); } else { PADDLE_THROW(platform::errors::InvalidArgument( - "AttnMatMul wrapper do not support (transA=T, transB=N)" + "AttnMatMul wrapper do not support (transA=T, transB=T/N)" "parameters.")); } - if (compute_bias_) { - // reduce: {0, 1, 2, 3, 4} -> {2, 3, 4} or {0, 1, 2} -> {2} + if (compute_bias_ && d_bias) { + // reduce: {0, 1, 2, 3, 4} -> {2, 3, 4} or {0, 1, 2} -> {2} or {0,1,2,3} + // -> {3} or {0,1,2,3,4} -> {3,4} const auto input_dims = d_output->dims(); const auto output_dims = d_bias->dims(); bool support_case_1 = @@ -163,11 +148,22 @@ class AttnMatMul { bool support_case_2 = (input_dims.size() == 3 && output_dims.size() == 1 && (input_dims[2] == output_dims[0])); + bool support_case_3 = + (input_dims.size() == 4 && output_dims.size() == 1 && + input_dims[3] == output_dims[0]); + bool support_case_4 = + (input_dims.size() == 5 && output_dims.size() == 2 && + input_dims[3] == output_dims[0] && input_dims[4] == output_dims[1]); + + gpuStream_t stream = dev_ctx_.stream(); if (support_case_1 || support_case_2) { - gpuStream_t stream = dev_ctx_.stream(); TensorReduceImpl>( dev_ctx_, *d_output, d_bias, kps::IdentityFunctor(), {0, 1}, stream); + } else if (support_case_3 || support_case_4) { + TensorReduceImpl>( + dev_ctx_, *d_output, d_bias, kps::IdentityFunctor(), {0, 1, 2}, + stream); } else { PADDLE_THROW(platform::errors::InvalidArgument( "Only support reduce when the input dims are [0,1,2,3,4] and " diff --git a/paddle/fluid/operators/fused/fmha_ref.h b/paddle/fluid/operators/fused/fmha_ref.h index 0e9fba73933..38f9aff226e 100644 --- a/paddle/fluid/operators/fused/fmha_ref.h +++ b/paddle/fluid/operators/fused/fmha_ref.h @@ -1,8 +1,11 @@ /* Copyright (c) 2021 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. @@ -297,7 +300,6 @@ class FMHARef { phi::SoftmaxBackwardCUDAKernelDriver( dev_ctx_, softmax_out_tensor, *softmax_out_grad_tensor, softmax_axis, src_mask_out_grad_tensor); - // recall LaunchElementwiseCudaKernel fw: src_mask_out = qk_out + // src_mask // Special case when dy is not needed and dx doesn't reduce diff --git a/paddle/fluid/operators/fused/fused_gate_attention.h b/paddle/fluid/operators/fused/fused_gate_attention.h new file mode 100644 index 00000000000..cda33987d68 --- /dev/null +++ b/paddle/fluid/operators/fused/fused_gate_attention.h @@ -0,0 +1,647 @@ +/* Copyright (c) 2022 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 + +#include "paddle/fluid/operators/elementwise/elementwise_op_broadcast.cu.h" +#include "paddle/fluid/operators/transpose_op.cu.h" +#include "paddle/phi/kernels/funcs/broadcast_function.h" +#include "paddle/phi/kernels/funcs/elementwise_base.h" +#include "paddle/phi/kernels/funcs/elementwise_functor.h" +#include "paddle/phi/kernels/gpudnn/softmax_gpudnn.h" + +namespace paddle { +namespace operators { + +using Tensor = framework::Tensor; + +inline std::string MemoryDebugString(const Tensor& t) { + std::stringstream ss; + ss << "shape=[" << t.dims() + << "], size=" << static_cast(t.memory_size()) / (1 << 20) + << " MB, ptr=" << t.data(); + + size_t total = 0; + size_t available = 0; + platform::GpuMemoryUsage(&available, &total); + ss << "; memory allocated=" + << static_cast(total - available) / (1 << 20) << " MB"; + return ss.str(); +} + +template +struct TernaryAddFunctor { + inline HOSTDEVICE T operator()(T a, T b, T c) const { return a + b + c; } +}; + +template +struct GateAttentionConfig { + public: + int64_t batch_size; + int64_t seq_len_m; + int64_t seq_len_r; + int64_t q_dim; + int64_t kv_dim; + int64_t key_dim; + int64_t m_size; + int64_t num_heads; + + phi::DDim qkv_out_dims; + phi::DDim qkv_transpose_out_dims; + + phi::DDim q_out_dims; + phi::DDim kv_out_dims; + phi::DDim q_transpose_out_dims; + phi::DDim kv_transpose_out_dims; + + phi::DDim qk_out_dims; + phi::DDim softmax_out_dims; + phi::DDim qktv_out_dims; + phi::DDim gate_out_dims; + + GateAttentionConfig(const Tensor* query, const Tensor* key, + const Tensor* query_weight, const Tensor* qkv_weight, + bool merge_qkv) { + // query: shape=[batch_size, seq_len_m, seq_len_r, q_dim] + batch_size = query->dims()[0]; + seq_len_m = query->dims()[1]; + seq_len_r = query->dims()[2]; + q_dim = query->dims()[3]; + + if (merge_qkv) { + PADDLE_ENFORCE_NOT_NULL( + qkv_weight, + platform::errors::NotFound("The input qkv_weight can not be nullptr " + "when merge_qkv is true.")); + + // When q_dim == kv_dim, QKV matmul can be computed merged. + // qkv_weight: shape=[3, num_heads, key_dim, q_dim] + num_heads = qkv_weight->dims()[1]; + key_dim = qkv_weight->dims()[2]; + m_size = seq_len_r; + kv_dim = q_dim; + + qkv_out_dims = {batch_size, seq_len_m, seq_len_r, 3, num_heads, key_dim}; + qkv_transpose_out_dims = {3, batch_size, seq_len_m, + num_heads, seq_len_r, key_dim}; + } else { + PADDLE_ENFORCE_NOT_NULL( + key, + platform::errors::NotFound( + "The input key can not be nullptr when merge_qkv is false.")); + PADDLE_ENFORCE_NOT_NULL( + query_weight, + platform::errors::NotFound("The input query_weight can not be " + "nullptr when merge_qkv is false.")); + + // When q_dim != kv_dim, QKV matmul must be computed saparately. + // key: shape=[batch_size, seq_len_m, m_size, kv_dim] + // query_w: shape=[q_dim, num_heads, key_dim] + num_heads = query_weight->dims()[1]; + key_dim = query_weight->dims()[2]; + m_size = key->dims()[2]; + kv_dim = key->dims()[3]; + + q_out_dims = {batch_size, seq_len_m, seq_len_r, num_heads, key_dim}; + kv_out_dims = {batch_size, seq_len_m, m_size, num_heads, key_dim}; + q_transpose_out_dims = {batch_size, seq_len_m, num_heads, seq_len_r, + key_dim}; + kv_transpose_out_dims = {batch_size, seq_len_m, num_heads, m_size, + key_dim}; + } + + qk_out_dims = {batch_size, seq_len_m, num_heads, seq_len_r, m_size}; + softmax_out_dims = {batch_size, seq_len_m, num_heads, seq_len_r, m_size}; + qktv_out_dims = {batch_size, seq_len_m, num_heads, seq_len_r, key_dim}; + gate_out_dims = {batch_size, seq_len_m, seq_len_r, num_heads, key_dim}; + } + + int64_t GetQuerySize() const { + return batch_size * seq_len_m * seq_len_r * num_heads * key_dim; + } + + Tensor* GetQKVOut(const platform::CUDADeviceContext& dev_ctx) { + if (!qkv_out.IsInitialized()) { + qkv_out.Resize(qkv_out_dims); + qkv_out.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "qkv_out: " << MemoryDebugString(qkv_out); + } + return &qkv_out; + } + + Tensor* GetQueryOut(const platform::CUDADeviceContext& dev_ctx) { + if (!query_out.IsInitialized()) { + query_out.Resize(q_out_dims); + query_out.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "query_out: " << MemoryDebugString(query_out); + } + return &query_out; + } + + Tensor* GetKeyOut(const platform::CUDADeviceContext& dev_ctx) { + if (!key_out.IsInitialized()) { + key_out.Resize(kv_out_dims); + key_out.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "key_out: " << MemoryDebugString(key_out); + } + return &key_out; + } + + Tensor* GetValueOut(const platform::CUDADeviceContext& dev_ctx) { + if (!value_out.IsInitialized()) { + value_out.Resize(kv_out_dims); + value_out.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "value_out: " << MemoryDebugString(value_out); + } + return &value_out; + } + + Tensor* GetQKOut(const platform::CUDADeviceContext& dev_ctx, + Tensor* softmax_out) { + // softmax_dim = qk_out_dim[-1] = qk_out_dim[rank - 1] + int softmax_dim = m_size; + if (!softmax_out || phi::UseCudnnSoftmax(dev_ctx, softmax_dim, true)) { + // Not sure whether cudnn softmax can execute inplace. + if (!qkv_out.IsInitialized()) { + qk_out.Resize(qk_out_dims); + qk_out.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "qk_out: " << MemoryDebugString(qk_out); + } + return &qk_out; + } else { + return softmax_out; + } + } + + void ClearQKVOut() { + if (qkv_out.IsInitialized()) { + qkv_out.clear(); + } + } + + void ClearQKOut() { + if (qk_out.IsInitialized()) { + qk_out.clear(); + } + } + + protected: + Tensor qkv_out; + // QKV is not merged + Tensor query_out; + Tensor key_out; + Tensor value_out; + // qk_out = BatchedGEMM(Q, K^T) + // qk_out: shape=[batch_size, seq_len_m, num_heads, seq_len_r, m_size] + // softmax_out = softmax(qk_out + nonbatched_bias + src_mask) + // The shape of qk_out, softmax_out is the same, thus can be called inplace. + Tensor qk_out; +}; + +template +struct GateAttentionGradConfig : public GateAttentionConfig { + public: + GateAttentionGradConfig(const Tensor* query, const Tensor* key, + const Tensor* query_weight, const Tensor* qkv_weight, + bool merge_qkv) + : GateAttentionConfig(query, key, query_weight, qkv_weight, + merge_qkv) {} + + Tensor* GetQKVOutGrad(const platform::CUDADeviceContext& dev_ctx) { + if (!qkv_out_grad.IsInitialized()) { + qkv_out_grad.Resize(this->qkv_out_dims); + qkv_out_grad.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "qkv_out_grad: " << MemoryDebugString(qkv_out_grad); + } + return &qkv_out_grad; + } + + Tensor* GetQueryOutGrad(const platform::CUDADeviceContext& dev_ctx) { + if (!query_out_grad.IsInitialized()) { + query_out_grad.Resize(this->q_out_dims); + query_out_grad.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "query_out_grad: " << MemoryDebugString(query_out_grad); + } + return &query_out_grad; + } + + Tensor* GetKeyOutGrad(const platform::CUDADeviceContext& dev_ctx) { + if (!key_out_grad.IsInitialized()) { + key_out_grad.Resize(this->kv_out_dims); + key_out_grad.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "key_out_grad: " << MemoryDebugString(key_out_grad); + } + return &key_out_grad; + } + + Tensor* GetValueOutGrad(const platform::CUDADeviceContext& dev_ctx) { + if (!value_out_grad.IsInitialized()) { + value_out_grad.Resize(this->kv_out_dims); + value_out_grad.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "value_out_grad: " << MemoryDebugString(value_out_grad); + } + return &value_out_grad; + } + + Tensor* GetQKOutGrad(const platform::CUDADeviceContext& dev_ctx, + Tensor* softmax_out_grad) { + // softmax_dim = qk_out_dim[-1] = qk_out_dim[rank - 1] + int softmax_dim = this->m_size; + if (!softmax_out_grad || + phi::UseCudnnSoftmax(dev_ctx, softmax_dim, true)) { + if (!qk_out_grad.IsInitialized()) { + qk_out_grad.Resize(this->qk_out_dims); + qk_out_grad.mutable_data(dev_ctx.GetPlace()); + VLOG(4) << "qk_out_grad: " << MemoryDebugString(qk_out_grad); + } + return &qk_out_grad; + } else { + return softmax_out_grad; + } + } + + protected: + Tensor qkv_out_grad; + Tensor query_out_grad; + Tensor key_out_grad; + Tensor value_out_grad; + Tensor qk_out_grad; +}; + +template +class FMHAGateRef { + public: + FMHAGateRef(const platform::CUDADeviceContext& dev_ctx, bool merge_qkv) + : dev_ctx_(dev_ctx), merge_qkv_(merge_qkv) {} + + void ComputeForward(const Tensor* nonbatched_bias, const Tensor* src_mask, + Tensor* q_transpose_out, Tensor* k_transpose_out, + Tensor* v_transpose_out, Tensor* qkv_transpose_out, + Tensor* softmax_out, Tensor* fmha_out, + GateAttentionConfig* config) { + T* q_ptr = nullptr; + T* k_ptr = nullptr; + T* v_ptr = nullptr; + if (merge_qkv_) { + // qkv_transpose_out = transpose(qkv_out) + PADDLE_ENFORCE_NOT_NULL( + qkv_transpose_out, + platform::errors::NotFound("The input qkv_transpose_out can not be " + "nullptr when merge_qkv is true.")); + + Tensor* qkv_out = config->GetQKVOut(dev_ctx_); + ComputeQKVTransposeForward(*qkv_out, qkv_transpose_out); + config->ClearQKVOut(); + + // q_size == k_size + int64_t q_size = config->GetQuerySize(); + q_ptr = qkv_transpose_out->data(); + k_ptr = q_ptr + q_size; + v_ptr = k_ptr + q_size; + } else { + PADDLE_ENFORCE_NOT_NULL( + q_transpose_out, + platform::errors::NotFound("The input q_transpose_out can not be " + "nullptr when merge_qkv is false.")); + PADDLE_ENFORCE_NOT_NULL( + k_transpose_out, + platform::errors::NotFound("The input k_transpose_out can not be " + "nullptr when merge_qkv is false.")); + PADDLE_ENFORCE_NOT_NULL( + v_transpose_out, + platform::errors::NotFound("The input v_transpose_out can not be " + "nullptr when merge_qkv is false.")); + + Tensor* query_out = config->GetQueryOut(dev_ctx_); + Tensor* key_out = config->GetKeyOut(dev_ctx_); + Tensor* value_out = config->GetValueOut(dev_ctx_); + ComputeQKVTransposeForward(*query_out, *key_out, *value_out, + q_transpose_out, k_transpose_out, + v_transpose_out); + + // q_size != k_size + q_ptr = q_transpose_out->data(); + k_ptr = k_transpose_out->data(); + v_ptr = v_transpose_out->data(); + } + + // qk_out = BatchedGEMM(Q, K^T) + // [batch_size, seq_len_m, num_heads, seq_len_r, key_dim] * + // [batch_size, seq_len_m, num_heads, m_size, key_dim] + // -> [batch_size, seq_len_m, num_heads, seq_len_r, m_size] + Tensor* qk_out = config->GetQKOut(dev_ctx_, softmax_out); + T* qk_out_ptr = qk_out->data(); + + int64_t gemm_batch_size = + config->batch_size * config->seq_len_m * config->num_heads; + int64_t gemm_m = config->seq_len_r; + int64_t gemm_n = config->m_size; + int64_t gemm_k = config->key_dim; + + T alpha = static_cast(1.0 / sqrt(config->key_dim)); + ComputeBatchedGEMM(q_ptr, k_ptr, qk_out_ptr, false, true, gemm_m, gemm_n, + gemm_k, gemm_batch_size, alpha); + + // softmax_out = softmax(qk_out + nonbatched_bias + src_mask) + ComputeBiasMaskSoftmaxForward(nonbatched_bias, src_mask, qk_out, + softmax_out); + config->ClearQKOut(); + + // qktv_out = BatchedGEMM(softmax_out, V) + // [batch_size, seq_len_m, num_heads, seq_len_r, m_size] * + // [batch_size, seq_len_m, num_heads, m_size, key_dim] + // -> [batch_size, seq_len_m, num_heads, seq_len_r, key_dim] + Tensor qktv_out; + qktv_out.Resize(config->qktv_out_dims); + T* qktv_out_ptr = qktv_out.mutable_data(dev_ctx_.GetPlace()); + + gemm_m = config->seq_len_r; + gemm_n = config->key_dim; + gemm_k = config->m_size; + + T* softmax_out_ptr = softmax_out->data(); + ComputeBatchedGEMM(softmax_out_ptr, v_ptr, qktv_out_ptr, false, false, + gemm_m, gemm_n, gemm_k, gemm_batch_size); + + // fmha_out = transpose(qktv_out) + ComputeQKTVTransposeForward(qktv_out, fmha_out); + } + + void ComputeBackward(const Tensor* q_transpose_out, + const Tensor* k_transpose_out, + const Tensor* v_transpose_out, + const Tensor* qkv_transpose_out, + const Tensor* softmax_out, const Tensor* fmha_out_grad, + Tensor* src_mask_grad, Tensor* nonbatched_bias_grad, + GateAttentionGradConfig* config) { + const T* q_ptr = nullptr; + const T* k_ptr = nullptr; + const T* v_ptr = nullptr; + + T* q_grad_ptr = nullptr; + T* k_grad_ptr = nullptr; + T* v_grad_ptr = nullptr; + + Tensor q_transpose_out_grad; + Tensor k_transpose_out_grad; + Tensor v_transpose_out_grad; + Tensor qkv_transpose_out_grad; + if (merge_qkv_) { + PADDLE_ENFORCE_NOT_NULL( + qkv_transpose_out, + platform::errors::NotFound("The input qkv_transpose_out can not be " + "nullptr when merge_qkv is true.")); + + int64_t q_size = config->GetQuerySize(); + q_ptr = qkv_transpose_out->data(); + k_ptr = q_ptr + q_size; + v_ptr = k_ptr + q_size; + + qkv_transpose_out_grad.Resize(config->qkv_transpose_out_dims); + + q_grad_ptr = qkv_transpose_out_grad.mutable_data(dev_ctx_.GetPlace()); + k_grad_ptr = q_grad_ptr + q_size; + v_grad_ptr = k_grad_ptr + q_size; + } else { + PADDLE_ENFORCE_NOT_NULL( + q_transpose_out, + platform::errors::NotFound("The input q_transpose_out can not be " + "nullptr when merge_qkv is false.")); + PADDLE_ENFORCE_NOT_NULL( + k_transpose_out, + platform::errors::NotFound("The input k_transpose_out can not be " + "nullptr when merge_qkv is false.")); + PADDLE_ENFORCE_NOT_NULL( + v_transpose_out, + platform::errors::NotFound("The input v_transpose_out can not be " + "nullptr when merge_qkv is false.")); + + q_ptr = q_transpose_out->data(); + k_ptr = k_transpose_out->data(); + v_ptr = v_transpose_out->data(); + + q_transpose_out_grad.Resize(config->q_transpose_out_dims); + k_transpose_out_grad.Resize(config->kv_transpose_out_dims); + v_transpose_out_grad.Resize(config->kv_transpose_out_dims); + + q_grad_ptr = q_transpose_out_grad.mutable_data(dev_ctx_.GetPlace()); + k_grad_ptr = k_transpose_out_grad.mutable_data(dev_ctx_.GetPlace()); + v_grad_ptr = v_transpose_out_grad.mutable_data(dev_ctx_.GetPlace()); + } + + Tensor softmax_out_grad; + softmax_out_grad.Resize(config->softmax_out_dims); + softmax_out_grad.mutable_data(dev_ctx_.GetPlace()); + + int64_t gemm_batch_size = + config->batch_size * config->seq_len_m * config->num_heads; + { + // Forward: fmha_out = transpose(qktv_out) + Tensor qktv_out_grad; + qktv_out_grad.Resize(config->qktv_out_dims); + T* qktv_out_grad_ptr = qktv_out_grad.mutable_data(dev_ctx_.GetPlace()); + ComputeQKTVTransposeBackward(*fmha_out_grad, &qktv_out_grad); + + // Forward: qktv_out = BatchedGEMM(softmax_out, V) + // Backward: + // V_grad = BatchedGEMM(softmax_out^T, qktv_out_grad) (dy = x^T * dout) + int64_t gemm_m = config->m_size; + int64_t gemm_n = config->key_dim; + int64_t gemm_k = config->seq_len_r; + + const T* softmax_out_ptr = softmax_out->data(); + ComputeBatchedGEMM(softmax_out_ptr, qktv_out_grad_ptr, v_grad_ptr, true, + false, gemm_m, gemm_n, gemm_k, gemm_batch_size); + + // Backward: softmax_out_grad = qktv_out_grad * V^T (dx = dout * y^T) + gemm_m = config->seq_len_r; + gemm_n = config->m_size; + gemm_k = config->key_dim; + + T* softmax_out_grad_ptr = softmax_out_grad.data(); + ComputeBatchedGEMM(qktv_out_grad_ptr, v_ptr, softmax_out_grad_ptr, false, + true, gemm_m, gemm_n, gemm_k, gemm_batch_size); + } + + Tensor* qk_out_grad = config->GetQKOutGrad(dev_ctx_, &softmax_out_grad); + ComputeBiasMaskSoftmaxBackward(&softmax_out_grad, softmax_out, + src_mask_grad, qk_out_grad, + nonbatched_bias_grad); + + // Forward: qk_out = BatchedGEMM(Q, K^T) + // Backward: k_grad = BatchedGEMM(qk_out_grad^T, Q) (dy = dout^t * x) + int64_t gemm_m = config->m_size; + int64_t gemm_n = config->key_dim; + int64_t gemm_k = config->seq_len_r; + T alpha = static_cast(1.0 / sqrt(config->key_dim)); + + T* qk_out_grad_ptr = qk_out_grad->data(); + ComputeBatchedGEMM(qk_out_grad_ptr, q_ptr, k_grad_ptr, true, false, gemm_m, + gemm_n, gemm_k, gemm_batch_size, alpha); + + // Backward: q_grad = BatchedGEMM(qk_out_grad, K) (dx = dout * y) + gemm_m = config->seq_len_r; + gemm_n = config->key_dim; + gemm_k = config->m_size; + ComputeBatchedGEMM(qk_out_grad_ptr, k_ptr, q_grad_ptr, false, false, gemm_m, + gemm_n, gemm_k, gemm_batch_size, alpha); + + if (merge_qkv_) { + Tensor* qkv_out_grad = config->GetQKVOutGrad(dev_ctx_); + ComputeQKVTransposeBackward(qkv_transpose_out_grad, qkv_out_grad); + } else { + Tensor* q_out_grad = config->GetQueryOutGrad(dev_ctx_); + Tensor* k_out_grad = config->GetKeyOutGrad(dev_ctx_); + Tensor* v_out_grad = config->GetValueOutGrad(dev_ctx_); + ComputeQKVTransposeBackward(q_transpose_out_grad, k_transpose_out_grad, + v_transpose_out_grad, q_out_grad, k_out_grad, + v_out_grad); + } + } + + void ComputeQKVTransposeForward(const Tensor& q_out, const Tensor& k_out, + const Tensor& v_out, Tensor* q_transpose_out, + Tensor* k_transpose_out, + Tensor* v_transpose_out) { + int ndims = 5; + std::vector perm = {0, 1, 3, 2, 4}; + TransposeGPUKernelDriver(dev_ctx_, ndims, q_out, perm, q_transpose_out); + TransposeGPUKernelDriver(dev_ctx_, ndims, k_out, perm, k_transpose_out); + TransposeGPUKernelDriver(dev_ctx_, ndims, v_out, perm, v_transpose_out); + } + + void ComputeQKVTransposeBackward(const Tensor& q_transpose_out_grad, + const Tensor& k_transpose_out_grad, + const Tensor& v_transpose_out_grad, + Tensor* q_out_grad, Tensor* k_out_grad, + Tensor* v_out_grad) { + int ndims = 5; + std::vector perm = {0, 1, 3, 2, 4}; + TransposeGPUKernelDriver(dev_ctx_, ndims, q_transpose_out_grad, perm, + q_out_grad); + TransposeGPUKernelDriver(dev_ctx_, ndims, k_transpose_out_grad, perm, + k_out_grad); + TransposeGPUKernelDriver(dev_ctx_, ndims, v_transpose_out_grad, perm, + v_out_grad); + } + + // [batch_size, seq_len_m, seq_len_r, 3, num_heads, key_dim] -> + // [3, batch_size, seq_len_m, num_heads, seq_len_r, key_dim] + void ComputeQKVTransposeForward(const Tensor& qkv_out, + Tensor* qkv_transpose_out) { + int ndims = 6; + std::vector perm = {3, 0, 1, 4, 2, 5}; + TransposeGPUKernelDriver(dev_ctx_, ndims, qkv_out, perm, + qkv_transpose_out); + } + + void ComputeQKVTransposeBackward(const Tensor& qkv_transpose_out_grad, + Tensor* qkv_out_grad) { + int ndims = 6; + std::vector perm = {1, 2, 4, 0, 3, 5}; + TransposeGPUKernelDriver(dev_ctx_, ndims, qkv_transpose_out_grad, perm, + qkv_out_grad); + } + + // [batch_size, seq_len_m, num_head, seq_len_r, c] -> + // [batch_size, seq_len_m, seq_len_r, num_head, c] + void ComputeQKTVTransposeForward(const Tensor& qktv_out, Tensor* fmha_out) { + int ndims = 5; + std::vector perm = {0, 1, 3, 2, 4}; + TransposeGPUKernelDriver(dev_ctx_, ndims, qktv_out, perm, fmha_out); + } + + void ComputeQKTVTransposeBackward(const Tensor& fmha_out_grad, + Tensor* qktv_out_grad) { + int ndims = 5; + std::vector perm = {0, 1, 3, 2, 4}; + TransposeGPUKernelDriver(dev_ctx_, ndims, fmha_out_grad, perm, + qktv_out_grad); + } + + // qk_out = qk_out + nonbatched_bias + src_mask + // softmax_out = softmax(src_mask_out) + void ComputeBiasMaskSoftmaxForward(const Tensor* nonbatched_bias, + const Tensor* src_mask, Tensor* qk_out, + Tensor* softmax_out) { + if (nonbatched_bias) { + std::vector ins = {qk_out, nonbatched_bias, src_mask}; + std::vector outs = {qk_out}; + phi::funcs::BroadcastKernel( + dev_ctx_, ins, &outs, -1, TernaryAddFunctor()); + } else { + std::vector ins = {qk_out, src_mask}; + std::vector outs = {qk_out}; + phi::funcs::BroadcastKernel( + dev_ctx_, ins, &outs, -1, phi::funcs::AddFunctor()); + } + phi::SoftmaxForwardCUDAKernelDriver(dev_ctx_, *qk_out, -1, softmax_out); + } + + // src_mask_out = qk_out + nonbatched_bias + src_mask + // softmax_out = softmax(src_mask_out) + void ComputeBiasMaskSoftmaxBackward(const Tensor* softmax_out_grad, + const Tensor* softmax_out, + Tensor* src_mask_grad, + Tensor* qk_out_grad, + Tensor* nonbatched_bias_grad) { + PADDLE_ENFORCE_NOT_NULL( + qk_out_grad, + platform::errors::NotFound("The qk_out_grad can not be nullptr.")); + + PADDLE_ENFORCE_EQ(qk_out_grad->dims(), softmax_out->dims(), + platform::errors::InvalidArgument( + "The shape of qk_out_grad and softmax_out is " + "expected to be the same. But recieved qk_out_grad's " + "shape = %s, softmax_out's shape = %s.", + qk_out_grad->dims(), softmax_out->dims())); + + PADDLE_ENFORCE_EQ(src_mask_grad, nullptr, + platform::errors::InvalidArgument( + "src_mask_grad is expected to be nullptr.")); + + phi::SoftmaxBackwardCUDAKernelDriver(dev_ctx_, *softmax_out, + *softmax_out_grad, -1, qk_out_grad); + + // [1, bs, num_head, seq_l, seq_l] -> [bs, num_head, seq_l, seq_l] + if (nonbatched_bias_grad) { + gpuStream_t stream = dev_ctx_.stream(); + TensorReduceImpl>( + dev_ctx_, *qk_out_grad, nonbatched_bias_grad, + kps::IdentityFunctor(), {0, 1}, stream); + } + } + + private: + void ComputeBatchedGEMM(const T* a_ptr, const T* b_ptr, T* c_ptr, + bool trans_a, bool trans_b, int64_t m, int64_t n, + int64_t k, int64_t batch_size, + T alpha = static_cast(1.0), + T beta = static_cast(0.0)) { + CBLAS_TRANSPOSE cblas_trans_a = trans_a ? CblasTrans : CblasNoTrans; + CBLAS_TRANSPOSE cblas_trans_b = trans_b ? CblasTrans : CblasNoTrans; + int64_t stride_a = m * k; + int64_t stride_b = k * n; + + auto blas = phi::funcs::GetBlas(dev_ctx_); + blas.BatchedGEMM(cblas_trans_a, cblas_trans_b, m, n, k, alpha, a_ptr, b_ptr, + beta, c_ptr, batch_size, stride_a, stride_b); + } + + const platform::CUDADeviceContext& dev_ctx_; + bool merge_qkv_; +}; + +} // namespace operators +} // namespace paddle diff --git a/paddle/fluid/operators/fused/fused_gate_attention_op.cc b/paddle/fluid/operators/fused/fused_gate_attention_op.cc new file mode 100644 index 00000000000..ba9dbd82e3d --- /dev/null +++ b/paddle/fluid/operators/fused/fused_gate_attention_op.cc @@ -0,0 +1,317 @@ +/* Copyright (c) 2022 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 +#include +#include "paddle/fluid/framework/op_registry.h" + +namespace paddle { +namespace operators { + +using Tensor = framework::Tensor; +using DDim = framework::DDim; + +class FusedGateAttentionOp : public framework::OperatorWithKernel { + public: + using framework::OperatorWithKernel::OperatorWithKernel; + + void InferShape(framework::InferShapeContext* ctx) const override { + OP_INOUT_CHECK(ctx->HasInput("Query"), "Input", "Query", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasInput("OutLinearWeight"), "Input", "OutLinearWeight", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasInput("OutLinearBias"), "Input", "OutLinearBias", + "fused_gate_attention"); + + OP_INOUT_CHECK(ctx->HasOutput("SoftmaxOut"), "Output", "SoftmaxOut", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasOutput("FMHAOut"), "Output", "FMHAOut", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", + "fused_gate_attention"); + + auto input_q_dims = ctx->GetInputDim("Query"); + int batch_size = input_q_dims[0]; + int seq_len_m = input_q_dims[1]; + int seq_len_r = input_q_dims[2]; + + int num_head, m_size, key_dim; + if (ctx->Attrs().Get("merge_qkv")) { + // QKV's input: [batch_size, seq_len_m, seq_len_r, qkv_dim] + // QKV's weight: [3, num_head, key_dim, qkv_dim] + OP_INOUT_CHECK(ctx->HasInput("QKVWeight"), "Input", "QKVWeight", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasOutput("QKVTransposeOut"), "Output", + "QKVTransposeOut", "fused_gate_attention"); + + auto qkv_w_dims = ctx->GetInputDim("QKVWeight"); + + num_head = qkv_w_dims[1]; + key_dim = qkv_w_dims[2]; + m_size = seq_len_r; + + ctx->SetOutputDim("QKVTransposeOut", {3, batch_size, seq_len_m, num_head, + seq_len_r, key_dim}); + } else { + OP_INOUT_CHECK(ctx->HasInput("QueryWeight"), "Input", "QueryWeight", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasInput("KeyWeight"), "Input", "KeyWeight", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasInput("ValueWeight"), "Input", "ValueWeight", + "fused_gate_attention"); + + auto input_k_dims = ctx->GetInputDim("Key"); + auto q_w_dims = ctx->GetInputDim("QueryWeight"); + + num_head = q_w_dims[1]; + key_dim = q_w_dims[2]; + m_size = input_k_dims[2]; + + ctx->SetOutputDim("QueryTransposeOut", + {batch_size, seq_len_m, num_head, seq_len_r, key_dim}); + ctx->SetOutputDim("KeyTransposeOut", + {batch_size, seq_len_m, num_head, m_size, key_dim}); + ctx->SetOutputDim("ValueTransposeOut", + {batch_size, seq_len_m, num_head, m_size, key_dim}); + } + + ctx->SetOutputDim("SoftmaxOut", + {batch_size, seq_len_m, num_head, seq_len_r, m_size}); + ctx->SetOutputDim("FMHAOut", + {batch_size, seq_len_m, seq_len_r, num_head, key_dim}); + + if (ctx->Attrs().Get("has_gating")) { + OP_INOUT_CHECK(ctx->HasInput("GateWeight"), "Input", "GateWeight", + "fused_gate_attention"); + OP_INOUT_CHECK(ctx->HasInput("GateBias"), "Input", "GateBias", + "fused_gate_attention"); + ctx->SetOutputDim("GateOut", + {batch_size, seq_len_m, seq_len_r, num_head, key_dim}); + } + + ctx->SetOutputDim("Out", ctx->GetInputDim("Query")); + } +}; + +class FusedGateAttentionOpMaker : public framework::OpProtoAndCheckerMaker { + public: + void Make() override { + AddInput("Query", "The query tensor."); + AddInput("Key", "The key tensor.").AsDispensable(); + AddInput("QueryWeight", "(optional) The query weight tensor.") + .AsDispensable(); + AddInput("KeyWeight", "(optional) The key weight tensor.").AsDispensable(); + AddInput("ValueWeight", "(optional) The value weight tensor.") + .AsDispensable(); + AddInput("QKVWeight", "(optional) The qkv weight tensor.").AsDispensable(); + AddInput("NonbatchedBias", "(optional) The nonbatchedBias tensor.") + .AsDispensable(); + AddInput("SrcMask", "The attention mask tensor in fmha."); + AddInput("GateWeight", "(optional) The gate weight tensor.") + .AsDispensable(); + AddInput("GateBias", "(optional) The gate bias tensor.").AsDispensable(); + AddInput("OutLinearWeight", "The out_linear weight tensor."); + AddInput("OutLinearBias", "The out_linear bias tensor."); + AddOutput("QueryTransposeOut", "The transposed result of query matmul.") + .AsIntermediate() + .AsDispensable(); + AddOutput("KeyTransposeOut", "The transposed result of key matmul.") + .AsIntermediate() + .AsDispensable(); + AddOutput("ValueTransposeOut", "The transposed result of value matmul.") + .AsIntermediate() + .AsDispensable(); + AddOutput("QKVTransposeOut", "The transposed result of merged QKV matmul.") + .AsIntermediate() + .AsDispensable(); + AddOutput("SoftmaxOut", "Result in fmha.").AsIntermediate(); + AddOutput("FMHAOut", "Result in fmha.").AsIntermediate(); + AddOutput("GateOut", "Result of the gating module.") + .AsIntermediate() + .AsDispensable(); + AddOutput("Out", "Result after attention."); + AddAttr("has_gating", + "if true, the attention op uses gate architecure, " + "[default true].") + .SetDefault(true); + AddAttr("merge_qkv", + "if true, calculation with merged qkv, " + "[default true].") + .SetDefault(true); + AddComment(R"DOC( + Add fused attention op whose logic is as follows: + { + q = paddle.einsum('nbqa,ahc->nbqhc', q_data, self.query_w) + k = paddle.einsum('nbka,ahc->nbkhc', m_data, self.key_w) + v = paddle.einsum('nbka,ahc->nbkhc', m_data, self.value_w) + + logits = paddle.einsum('nbqhc,nbkhc->nbhqk', q * c , k) + bias + weights = nn.functional.softmax(logits) + weighted_avg = paddle.einsum('nbhqk,nbkhc->nbqhc', weights, v) + if nonbatched_bias is not None: + logits += paddle.unsqueeze(nonbatched_bias, axis=1) + + if self.gating: + gate_values = paddle.einsum('nbqc,chv->nbqhv', q_data, + self.gating_w) + self.gating_b + gate_values_1 = nn.functional.sigmoid(gate_values) + weighted_avg *= gate_values_1 + + output = paddle.einsum('nbqhc,hco->nbqo', weighted_avg, + self.output_w) + self.output_b + + } + )DOC"); + } +}; + +class FusedGateAttentionGradOp : public framework::OperatorWithKernel { + public: + using framework::OperatorWithKernel::OperatorWithKernel; + + void InferShape(framework::InferShapeContext* ctx) const override { + OP_INOUT_CHECK(ctx->HasInput("Query"), "Input", "Query", + "fused_gate_attention_grad"); + if (ctx->HasOutput(framework::GradVarName("Query"))) { + ctx->SetOutputDim(framework::GradVarName("Query"), + ctx->GetInputDim("Query")); + } + if (ctx->HasOutput(framework::GradVarName("Key"))) { + ctx->SetOutputDim(framework::GradVarName("Key"), ctx->GetInputDim("Key")); + } + + if (ctx->Attrs().Get("merge_qkv")) { + OP_INOUT_CHECK(ctx->HasInput("QKVWeight"), "Input", "QKVWeight", + "fused_gate_attention_arad"); + ctx->SetOutputDim(framework::GradVarName("QKVWeight"), + ctx->GetInputDim("QKVWeight")); + } else { + OP_INOUT_CHECK(ctx->HasInput("QueryWeight"), "Input", "QueryWeight", + "fused_aate_attention_arad"); + OP_INOUT_CHECK(ctx->HasInput("KeyWeight"), "Input", "KeyWeight", + "fused_aate_attention_arad"); + OP_INOUT_CHECK(ctx->HasInput("ValueWeight"), "Input", "ValueWeight", + "fused_aate_attention_arad"); + + for (auto& name : {"QueryWeight", "KeyWeight", "ValueWeight"}) { + ctx->SetOutputDim(framework::GradVarName(name), ctx->GetInputDim(name)); + } + } + + OP_INOUT_CHECK(ctx->HasInput("OutLinearWeight"), "Input", "OutLinearWeight", + "fused_aate_attention_arad"); + + if (ctx->Attrs().Get("has_gating")) { + for (auto& name : {"GateWeight", "GateBias", "GateOut"}) { + ctx->SetOutputDim(framework::GradVarName(name), ctx->GetInputDim(name)); + } + } + + if (ctx->HasOutput(framework::GradVarName("NonbatchedBias"))) { + ctx->SetOutputDim(framework::GradVarName("NonbatchedBias"), + ctx->GetInputDim("NonbatchedBias")); + } + + ctx->SetOutputDim(framework::GradVarName("FMHAOut"), + ctx->GetInputDim("FMHAOut")); + + ctx->SetOutputDim(framework::GradVarName("OutLinearWeight"), + ctx->GetInputDim("OutLinearWeight")); + ctx->SetOutputDim(framework::GradVarName("OutLinearBias"), + ctx->GetInputDim("OutLinearBias")); + } +}; + +template +class FusedGateAttentionGradOpMaker : public framework::SingleGradOpMaker { + public: + using framework::SingleGradOpMaker::SingleGradOpMaker; + + protected: + void Apply(GradOpPtr op) const override { + op->SetType("fused_gate_attention_grad"); + op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out")); + + op->SetInput("Query", this->Input("Query")); + op->SetOutput(framework::GradVarName("Query"), this->InputGrad("Query")); + + op->SetAttrMap(this->Attrs()); + bool merge_qkv = BOOST_GET_CONST(bool, op->GetAttr("merge_qkv")); + if (merge_qkv) { + op->SetInput("QKVWeight", this->Input("QKVWeight")); + op->SetOutput(framework::GradVarName("QKVWeight"), + this->InputGrad("QKVWeight")); + op->SetInput("QKVTransposeOut", this->Output("QKVTransposeOut")); + } else { + op->SetInput("Key", this->Input("Key")); + op->SetOutput(framework::GradVarName("Key"), this->InputGrad("Key")); + + for (auto& name : {"QueryWeight", "KeyWeight", "ValueWeight"}) { + op->SetInput(name, this->Input(name)); + op->SetOutput(framework::GradVarName(name), this->InputGrad(name)); + } + + for (auto& name : + {"QueryTransposeOut", "KeyTransposeOut", "ValueTransposeOut"}) { + op->SetInput(name, this->Output(name)); + } + } + + op->SetInput("FMHAOut", this->Output("FMHAOut")); + op->SetOutput(framework::GradVarName("FMHAOut"), + this->OutputGrad("FMHAOut")); + + if (this->HasInput("NonbatchedBias")) { + op->SetInput("NonbatchedBias", this->Input("NonbatchedBias")); + op->SetOutput(framework::GradVarName("NonbatchedBias"), + this->InputGrad("NonbatchedBias")); + } + + op->SetInput("SoftmaxOut", this->Output("SoftmaxOut")); + + bool has_gating = BOOST_GET_CONST(bool, op->GetAttr("has_gating")); + if (has_gating) { + op->SetInput("GateWeight", this->Input("GateWeight")); + op->SetOutput(framework::GradVarName("GateWeight"), + this->InputGrad("GateWeight")); + + op->SetInput("GateBias", this->Input("GateBias")); + op->SetOutput(framework::GradVarName("GateBias"), + this->InputGrad("GateBias")); + + op->SetInput("GateOut", this->Output("GateOut")); + op->SetOutput(framework::GradVarName("GateOut"), + this->OutputGrad("GateOut")); + } + + op->SetInput("OutLinearWeight", this->Input("OutLinearWeight")); + op->SetOutput(framework::GradVarName("OutLinearWeight"), + this->InputGrad("OutLinearWeight")); + + op->SetInput("OutLinearBias", this->Input("OutLinearBias")); + op->SetOutput(framework::GradVarName("OutLinearBias"), + this->InputGrad("OutLinearBias")); + } +}; + +} // namespace operators +} // namespace paddle + +namespace ops = paddle::operators; +REGISTER_OPERATOR( + fused_gate_attention, ops::FusedGateAttentionOp, + ops::FusedGateAttentionOpMaker, + ops::FusedGateAttentionGradOpMaker, + ops::FusedGateAttentionGradOpMaker); +REGISTER_OPERATOR(fused_gate_attention_grad, ops::FusedGateAttentionGradOp); diff --git a/paddle/fluid/operators/fused/fused_gate_attention_op.cu b/paddle/fluid/operators/fused/fused_gate_attention_op.cu new file mode 100644 index 00000000000..b1badf72557 --- /dev/null +++ b/paddle/fluid/operators/fused/fused_gate_attention_op.cu @@ -0,0 +1,488 @@ +/* Copyright (c) 2022 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/fluid/framework/op_registry.h" +#include "paddle/fluid/framework/operator.h" +#include "paddle/fluid/operators/fused/attn_gemm.h" +#include "paddle/fluid/operators/fused/fused_gate_attention.h" +#include "paddle/fluid/platform/device/gpu/gpu_device_function.h" +#include "paddle/phi/kernels/funcs/math_function.h" + +namespace paddle { +namespace operators { + +using Tensor = framework::Tensor; + +template +struct SigmoidMultiplyFunctor { + using MPType = typename phi::dtype::MPTypeTrait::Type; + MPType one = static_cast(1.0f); + + // sigmoid(x) = 1 / (1 + exp(-x)) + // out = sigmoid(x) * y + inline HOSTDEVICE T operator()(T x, T y) const { + MPType x_mp = static_cast(x); + T sigmoid_out = static_cast(one / (one + exp(-x_mp))); + return sigmoid_out * y; + } +}; + +template +struct SigmoidMultiplyGradFunctor { + using MPType = typename phi::dtype::MPTypeTrait::Type; + MPType one = static_cast(1.0f); + + // Gradient of Multiply: + // dx = dout * y + // dy = dout * x + // Gradient of Sigmoid: dx = dout * out * (1 - out) + inline HOSTDEVICE phi::Array operator()(const T dout, const T x, + T y) const { + MPType x_mp = static_cast(x); + T sigmoid_out = static_cast(one / (one + exp(-x_mp))); + T d_sigmoid_out = dout * y; + phi::Array outs; + outs[0] = d_sigmoid_out * sigmoid_out * + (static_cast(1.0f) - sigmoid_out); // dx + outs[1] = dout * sigmoid_out; // dy + return outs; + } +}; + +template +void ComputeMergedQKVMatmulForward(const framework::ExecutionContext &ctx, + const GateAttentionConfig &config, + const Tensor *query, Tensor *qkv_out) { + // query: shape=[batch_size, seq_len_m, seq_len_r, qkv_dim] + // qkv_weight: shape=[3, num_heads, key_dim, qkv_dim] + // qkv_out: shape=[batch_size, seq_len_m, seq_len_r, 3, num_heads, key_dim] + auto *qkv_weight = ctx.Input("QKVWeight"); + + // qkv_out = GEMM(query, qkv_weight^T) + int m = config.batch_size * config.seq_len_m * config.seq_len_r; + int n = 3 * config.num_heads * config.key_dim; + int k = config.q_dim; + auto qkv_compute = + AttnMatMul(ctx.cuda_device_context(), false, true, m, n, k, false); + qkv_compute.ComputeForward(qkv_weight, query, nullptr, qkv_out, nullptr); +} + +template +Tensor *ComputeMergedQKVMatmulBackward(const framework::ExecutionContext &ctx, + const GateAttentionGradConfig &config, + const Tensor *query, + const Tensor *qkv_out_grad, + Tensor *query_grad, bool use_addto) { + auto *qkv_weight = ctx.Input("QKVWeight"); + auto *qkv_weight_grad = + ctx.Output(framework::GradVarName("QKVWeight")); + qkv_weight_grad->mutable_data(ctx.GetPlace()); + + // Gradient of GEMM(query, qkv_weight) + int m = config.batch_size * config.seq_len_m * config.seq_len_r; + int n = 3 * config.num_heads * config.key_dim; + int k = config.q_dim; + auto qkv_compute = + AttnMatMul(ctx.cuda_device_context(), false, true, m, n, k, false); + qkv_compute.ComputeBackward(query, qkv_weight, qkv_out_grad, query_grad, + qkv_weight_grad, nullptr, use_addto); + return query_grad; +} + +template +void ComputeSeparatedQKVMatmulForward(const framework::ExecutionContext &ctx, + const GateAttentionConfig &config, + const Tensor *query, const Tensor *key, + Tensor *query_out, Tensor *key_out, + Tensor *value_out) { + auto *query_weight = ctx.Input("QueryWeight"); + auto *key_weight = ctx.Input("KeyWeight"); + auto *value_weight = ctx.Input("ValueWeight"); + + // query_out = GEMM(query, query_weight) + // query: shape=[batch_size, seq_len_m, seq_len_r, q_dim] + // query_weight: shape=[q_dim, num_heads, key_dim] + // query_out: shape=[batch_size, seq_len_m, seq_len_r, num_heads, key_dim] + int q_m = config.batch_size * config.seq_len_m * config.seq_len_r; + int q_n = config.num_heads * config.key_dim; + int q_k = config.q_dim; + auto q_compute = AttnMatMul(ctx.cuda_device_context(), false, false, q_m, + q_n, q_k, false); + q_compute.ComputeForward(query_weight, query, nullptr, query_out, nullptr); + + // k_out = GEMM(key, key_weight) + // key: shape=[batch_size, seq_len_m, m_size, kv_dim] + // key_weight: shape=[kv_dim, num_heads, key_dim] + // key_out: shape=[batch_size, seq_len_m, m_size, num_heads, key_dim] + int kv_m = config.batch_size * config.seq_len_m * config.m_size; + int kv_n = config.num_heads * config.key_dim; + int kv_k = config.kv_dim; + auto kv_compute = AttnMatMul(ctx.cuda_device_context(), false, false, kv_m, + kv_n, kv_k, false); + kv_compute.ComputeForward(key_weight, key, nullptr, key_out, nullptr); + + // value_out = GEMM(value, value_weight) + kv_compute.ComputeForward(value_weight, key, nullptr, value_out, nullptr); +} + +template +Tensor *ComputeSeparatedQKVMatmulBackward( + const framework::ExecutionContext &ctx, + const GateAttentionGradConfig &config, const Tensor *query, + const Tensor *key, const Tensor *query_out_grad, const Tensor *key_out_grad, + const Tensor *value_out_grad, Tensor *query_grad, Tensor *key_grad, + bool use_addto) { + // Gradient of GEMM(key, k_weight) + const auto *key_weight = ctx.Input("KeyWeight"); + auto *key_weight_grad = + ctx.Output(framework::GradVarName("KeyWeight")); + key_weight_grad->mutable_data(ctx.GetPlace()); + + int kv_m = config.batch_size * config.seq_len_m * config.m_size; + int kv_n = config.num_heads * config.key_dim; + int kv_k = config.kv_dim; + auto kv_compute = AttnMatMul(ctx.cuda_device_context(), false, false, kv_m, + kv_n, kv_k, false); + kv_compute.ComputeBackward(key, key_weight, key_out_grad, key_grad, + key_weight_grad, nullptr, false); + + // Gradient of GEMM(value, v_weight) + auto *value_weight = ctx.Input("ValueWeight"); + auto *value_weight_grad = + ctx.Output(framework::GradVarName("ValueWeight")); + value_weight_grad->mutable_data(ctx.GetPlace()); + + kv_compute.ComputeBackward(key, value_weight, value_out_grad, key_grad, + value_weight_grad, nullptr, true); + + // Gradient of GEMM(query, query_weight) + const auto *query_weight = ctx.Input("QueryWeight"); + auto *query_weight_grad = + ctx.Output(framework::GradVarName("QueryWeight")); + query_weight_grad->mutable_data(ctx.GetPlace()); + + int q_m = config.batch_size * config.seq_len_m * config.seq_len_r; + int q_n = config.num_heads * config.key_dim; + int q_k = config.q_dim; + auto q_compute = AttnMatMul(ctx.cuda_device_context(), false, false, q_m, + q_n, q_k, false); + q_compute.ComputeBackward(query, query_weight, query_out_grad, query_grad, + query_weight_grad, nullptr, use_addto); + return query_grad; +} + +template +Tensor *ComputeGatingLinearForward(const framework::ExecutionContext &ctx, + const GateAttentionConfig &config, + const Tensor *query, + const Tensor *fmha_out) { + auto *gate_weight = ctx.Input("GateWeight"); + auto *gate_bias = ctx.Input("GateBias"); + + auto *gate_out = ctx.Output("GateOut"); + gate_out->mutable_data(ctx.GetPlace()); + VLOG(4) << "[ComputeGatingLinearForward] gate_out: " + << MemoryDebugString(*gate_out); + + // The first gate_bias_out stores the result of the multiplication, + // and the second gate_bias_out stores the result of the multiplication + + // bias. + // gate_out = GEMM(query, gate_weight) + gate_bias + int m = config.batch_size * config.seq_len_m * config.seq_len_r; + int n = config.num_heads * config.key_dim; + int k = config.q_dim; + auto gate_attn_compute = + AttnMatMul(ctx.cuda_device_context(), false, false, m, n, k, true); + gate_attn_compute.ComputeForward(gate_weight, query, gate_bias, gate_out, + gate_out); + + // gate_out = sigmoid(gate_out) * fmha_out + std::vector ins = {gate_out, fmha_out}; + std::vector outs = {gate_out}; + phi::funcs::ElementwiseKernel(ctx.cuda_device_context(), ins, &outs, + SigmoidMultiplyFunctor()); + return gate_out; +} + +template +Tensor *ComputeGatingLinearBackward(const framework::ExecutionContext &ctx, + const GateAttentionGradConfig &config, + const Tensor *fmha_out, + const Tensor *gate_out_grad, + Tensor *query_grad, Tensor *fmha_out_grad) { + const auto *query = ctx.Input("Query"); + const auto *gate_weight = ctx.Input("GateWeight"); + const auto *gate_bias = ctx.Input("GateBias"); + + // Re-compute gate_bias_out + Tensor gate_bias_out; + gate_bias_out.Resize(config.gate_out_dims); + gate_bias_out.mutable_data(ctx.GetPlace()); + + int m = config.batch_size * config.seq_len_m * config.seq_len_r; + int n = config.num_heads * config.key_dim; + int k = config.q_dim; + auto gate_attn_compute = + AttnMatMul(ctx.cuda_device_context(), false, false, m, n, k, true); + gate_attn_compute.ComputeForward(gate_weight, query, gate_bias, + &gate_bias_out, &gate_bias_out); + + // Gradient of sigmoid(gate_bias_out) * fmha_out + // Compute inplace and save gate_bias_out_grad to gate_bias_out. + std::vector ins = {gate_out_grad, &gate_bias_out, fmha_out}; + std::vector outs = {&gate_bias_out, fmha_out_grad}; + phi::funcs::ElementwiseKernel, 2>( + ctx.cuda_device_context(), ins, &outs, SigmoidMultiplyGradFunctor()); + + // Gradient of GEMM(query, gate_weight) + gate_bias + auto *gate_weight_grad = + ctx.Output(framework::GradVarName("GateWeight")); + auto *gate_bias_grad = ctx.Output(framework::GradVarName("GateBias")); + gate_weight_grad->mutable_data(ctx.GetPlace()); + gate_bias_grad->mutable_data(ctx.GetPlace()); + + gate_attn_compute.ComputeBackward(query, gate_weight, &gate_bias_out, + query_grad, gate_weight_grad, + gate_bias_grad); + return fmha_out_grad; +} + +template +Tensor *ComputeOutputLinearForward(const framework::ExecutionContext &ctx, + const GateAttentionConfig &config, + const Tensor *fmha_or_gate_out) { + const auto *out_linear_weight = ctx.Input("OutLinearWeight"); + const auto *out_linear_bias = ctx.Input("OutLinearBias"); + + auto *out = ctx.Output("Out"); + out->mutable_data(ctx.GetPlace()); + VLOG(4) << "[ComputeOutputLinearForward] out: " << MemoryDebugString(*out); + + // out = GEMM(fmha_or_gate_out, out_linear_weight) + out_linear_bias + int m = config.batch_size * config.seq_len_m * config.seq_len_r; + int n = config.q_dim; + int k = config.num_heads * config.key_dim; + auto out_linear_compute = + AttnMatMul(ctx.cuda_device_context(), false, false, m, n, k, true); + out_linear_compute.ComputeForward(out_linear_weight, fmha_or_gate_out, + out_linear_bias, out, out); + return out; +} + +template +Tensor *ComputeOutputLinearBackward(const framework::ExecutionContext &ctx, + const GateAttentionGradConfig &config, + bool has_gating) { + std::string input_name = has_gating ? "GateOut" : "FMHAOut"; + + const auto *out_grad = ctx.Input(framework::GradVarName("Out")); + const auto *out_linear_weight = ctx.Input("OutLinearWeight"); + const auto *input = ctx.Input(input_name); + + auto *out_linear_weight_grad = + ctx.Output(framework::GradVarName("OutLinearWeight")); + auto *out_linear_bias_grad = + ctx.Output(framework::GradVarName("OutLinearBias")); + auto *input_grad = ctx.Output(framework::GradVarName(input_name)); + + out_linear_weight_grad->mutable_data(ctx.GetPlace()); + out_linear_bias_grad->mutable_data(ctx.GetPlace()); + input_grad->mutable_data(ctx.GetPlace()); + + int m = config.batch_size * config.seq_len_m * config.seq_len_r; + int n = config.q_dim; + int k = config.num_heads * config.key_dim; + auto out_linear_compute = + AttnMatMul(ctx.cuda_device_context(), false, false, m, n, k, true); + out_linear_compute.ComputeBackward(input, out_linear_weight, out_grad, + input_grad, out_linear_weight_grad, + out_linear_bias_grad); + return input_grad; +} + +template +class FusedGateAttentionOpKernel : public framework::OpKernel { + public: + void Compute(const framework::ExecutionContext &ctx) const override { + const auto *query = ctx.Input("Query"); + const auto *key = ctx.Input("Key"); + const auto *query_weight = ctx.Input("QueryWeight"); + const auto *qkv_weight = ctx.Input("QKVWeight"); + + const auto *src_mask = ctx.Input("SrcMask"); + const auto *nonbatched_bias = ctx.Input("NonbatchedBias"); + + auto *q_transpose_out = ctx.Output("QueryTransposeOut"); + auto *k_transpose_out = ctx.Output("KeyTransposeOut"); + auto *v_transpose_out = ctx.Output("ValueTransposeOut"); + auto *qkv_transpose_out = ctx.Output("QKVTransposeOut"); + + auto *softmax_out = ctx.Output("SoftmaxOut"); + auto *fmha_out = ctx.Output("FMHAOut"); + + const bool merge_qkv = ctx.Attr("merge_qkv"); + const bool has_gating = ctx.Attr("has_gating"); + + // When seq_len_r = m_size, q_dim = kv_dim, QKV matmul can be merged. + auto &dev_ctx = ctx.template device_context(); + GateAttentionConfig config(query, key, query_weight, qkv_weight, + merge_qkv); + + if (merge_qkv) { + // 1. Merged QKV Matmul: einsum(nbhqk,nbkhc -> nbqhc) + Tensor *qkv_out = config.GetQKVOut(dev_ctx); + ComputeMergedQKVMatmulForward(ctx, config, query, qkv_out); + + qkv_transpose_out->mutable_data(ctx.GetPlace()); + VLOG(4) << "qkv_transpose_out:" << MemoryDebugString(*qkv_transpose_out); + } else { + // 1. Separated QKV Matmul + Tensor *query_out = config.GetQueryOut(dev_ctx); + Tensor *key_out = config.GetKeyOut(dev_ctx); + Tensor *value_out = config.GetValueOut(dev_ctx); + ComputeSeparatedQKVMatmulForward(ctx, config, query, key, query_out, + key_out, value_out); + + q_transpose_out->mutable_data(ctx.GetPlace()); + k_transpose_out->mutable_data(ctx.GetPlace()); + v_transpose_out->mutable_data(ctx.GetPlace()); + VLOG(4) << "q_transpose_out: " << MemoryDebugString(*q_transpose_out); + VLOG(4) << "k_transpose_out: " << MemoryDebugString(*k_transpose_out); + VLOG(4) << "v_transpose_out: " << MemoryDebugString(*v_transpose_out); + } + + softmax_out->mutable_data(ctx.GetPlace()); + fmha_out->mutable_data(ctx.GetPlace()); + VLOG(4) << "softmax_out: " << MemoryDebugString(*softmax_out); + VLOG(4) << "fmha_out: " << MemoryDebugString(*fmha_out); + + // 2. FMHA + auto fmha_compute = FMHAGateRef(dev_ctx, merge_qkv); + fmha_compute.ComputeForward( + nonbatched_bias, src_mask, q_transpose_out, k_transpose_out, + v_transpose_out, qkv_transpose_out, softmax_out, fmha_out, &config); + + // 3. Gating Linear + Tensor *fmha_or_gate_out = + !has_gating ? fmha_out : ComputeGatingLinearForward(ctx, config, + query, fmha_out); + + // 4. Output Linear + ComputeOutputLinearForward(ctx, config, fmha_or_gate_out); + } +}; + +template +class FusedGateAttentionGradKernel : public framework::OpKernel { + public: + void Compute(const framework::ExecutionContext &ctx) const override { + const auto has_gating = ctx.Attr("has_gating"); + const auto merge_qkv = ctx.Attr("merge_qkv"); + + // forward input + const auto *query = ctx.Input("Query"); + const auto *key = ctx.Input("Key"); + const auto *query_weight = ctx.Input("QueryWeight"); + const auto *qkv_weight = ctx.Input("QKVWeight"); + + // forward output, backward input + const auto *q_transpose_out = ctx.Input("QueryTransposeOut"); + const auto *k_transpose_out = ctx.Input("KeyTransposeOut"); + const auto *v_transpose_out = ctx.Input("ValueTransposeOut"); + const auto *qkv_transpose_out = ctx.Input("QKVTransposeOut"); + const auto *softmax_out = ctx.Input("SoftmaxOut"); + const auto *fmha_out = ctx.Input("FMHAOut"); + + // backward output + auto *query_grad = ctx.Output(framework::GradVarName("Query")); + query_grad->mutable_data(ctx.GetPlace()); + auto *nonbatched_bias_grad = + ctx.Output(framework::GradVarName("NonbatchedBias")); + auto *fmha_out_grad = ctx.Output(framework::GradVarName("FMHAOut")); + + auto &dev_ctx = ctx.template device_context(); + GateAttentionGradConfig config(query, key, query_weight, qkv_weight, + merge_qkv); + + // 1. Gradient of Output Linear + Tensor *fhma_or_gate_out_grad = + ComputeOutputLinearBackward(ctx, config, has_gating); + + // 2. Gradient of Gating Linear + if (has_gating) { + // fhma_or_gate_out_grad is actually gate_out_grad. + fmha_out_grad->mutable_data(ctx.GetPlace()); + ComputeGatingLinearBackward(ctx, config, fmha_out, + fhma_or_gate_out_grad, query_grad, + fmha_out_grad); + } + + // 3. Gradient of FMHA + if (nonbatched_bias_grad) { + nonbatched_bias_grad->mutable_data(ctx.GetPlace()); + } + + auto fmha_compute = FMHAGateRef(dev_ctx, merge_qkv); + fmha_compute.ComputeBackward( + q_transpose_out, k_transpose_out, v_transpose_out, qkv_transpose_out, + softmax_out, fmha_out_grad, nullptr, nonbatched_bias_grad, &config); + + bool use_addto = has_gating ? true : false; + if (merge_qkv) { + // 4. Gradient of Merged QKV Matmul + Tensor *qkv_out_grad = config.GetQKVOutGrad(dev_ctx); + ComputeMergedQKVMatmulBackward(ctx, config, query, qkv_out_grad, + query_grad, use_addto); + } else { + // 4. Gradient of Separated QKV Matmul + auto *key_grad = ctx.Output(framework::GradVarName("Key")); + if (key_grad) { + key_grad->mutable_data(ctx.GetPlace()); + } + Tensor *query_out_grad = config.GetQueryOutGrad(dev_ctx); + Tensor *key_out_grad = config.GetKeyOutGrad(dev_ctx); + Tensor *value_out_grad = config.GetValueOutGrad(dev_ctx); + ComputeSeparatedQKVMatmulBackward( + ctx, config, query, key, query_out_grad, key_out_grad, value_out_grad, + query_grad, key_grad, use_addto); + } + } +}; + +} // namespace operators +} // namespace paddle + +namespace ops = paddle::operators; +namespace plat = paddle::platform; +#ifdef PADDLE_WITH_HIP +REGISTER_OP_CUDA_KERNEL(fused_gate_attention, + ops::FusedGateAttentionOpKernel, + ops::FusedGateAttentionOpKernel, + ops::FusedGateAttentionOpKernel); +REGISTER_OP_CUDA_KERNEL(fused_gate_attention_grad, + ops::FusedGateAttentionGradKernel, + ops::FusedGateAttentionGradKernel, + ops::FusedGateAttentionGradKernel); +#else +REGISTER_OP_CUDA_KERNEL(fused_gate_attention, + ops::FusedGateAttentionOpKernel, + ops::FusedGateAttentionOpKernel, + ops::FusedGateAttentionOpKernel, + ops::FusedGateAttentionOpKernel); +REGISTER_OP_CUDA_KERNEL(fused_gate_attention_grad, + ops::FusedGateAttentionGradKernel, + ops::FusedGateAttentionGradKernel, + ops::FusedGateAttentionGradKernel, + ops::FusedGateAttentionGradKernel); +#endif diff --git a/paddle/fluid/platform/device/gpu/gpu_info.cc b/paddle/fluid/platform/device/gpu/gpu_info.cc index 5410638ceb3..8c04e935134 100644 --- a/paddle/fluid/platform/device/gpu/gpu_info.cc +++ b/paddle/fluid/platform/device/gpu/gpu_info.cc @@ -225,9 +225,9 @@ class RecordedGpuMallocHelper { if (UNLIKELY(malloc_managed_memory)) { result = cudaMallocManaged(ptr, size); } else { - VLOG(10) << "[cudaMalloc] size=" << static_cast(size) / (1 << 20) - << " MB"; result = cudaMalloc(ptr, size); + VLOG(10) << "[cudaMalloc] size=" << static_cast(size) / (1 << 20) + << " MB, result=" << result; } #endif if (result == gpuSuccess) { diff --git a/paddle/fluid/pybind/op_function_generator.h b/paddle/fluid/pybind/op_function_generator.h index 2b849968c76..df32f65a794 100644 --- a/paddle/fluid/pybind/op_function_generator.h +++ b/paddle/fluid/pybind/op_function_generator.h @@ -32,6 +32,10 @@ std::map> op_ins_map = { {"fused_attention", {"X", "LnScale", "LnBias", "QKVW", "QKVBias", "CacheKV", "SrcMask", "OutLinearW", "OutLinearBias", "Ln2Scale", "Ln2Bias"}}, + {"fused_gate_attention", + {"Query", "Key", "QueryWeight", "KeyWeight", "ValueWeight", "QKVWeight", + "NonbatchedBias", "SrcMask", "GateWeight", "GateBias", "OutLinearWeight", + "OutLinearBias"}}, {"fused_multi_transformer", {"X", "LnScale", "LnBias", "QKVW", "QKVBias", "CacheKV", "TimeStep", "SrcMask", "OutLinearW", "OutLinearBias", "FFNLnScale", "FFNLnBias", @@ -148,6 +152,9 @@ std::map> op_outs_map = { "DropoutMaskOut", "Ln2Mean", "Ln2Variance", "BiasDropoutResidualOut", "CacheKVOut", "Y"}}, + {"fused_gate_attention", + {"QueryTransposeOut", "KeyTransposeOut", "ValueTransposeOut", + "QKVTransposeOut", "SoftmaxOut", "FMHAOut", "GateOut", "Out"}}, {"sync_batch_norm", {"Y", "MeanOut", "VarianceOut", "SavedMean", "SavedVariance", "ReserveSpace"}}, diff --git a/paddle/phi/kernels/gpudnn/softmax_gpudnn.h b/paddle/phi/kernels/gpudnn/softmax_gpudnn.h index 77159bfc876..58781e8c6e4 100644 --- a/paddle/phi/kernels/gpudnn/softmax_gpudnn.h +++ b/paddle/phi/kernels/gpudnn/softmax_gpudnn.h @@ -888,19 +888,6 @@ void SoftmaxBackwardCudnnKernel(const GPUContext& dev_ctx, #endif } -template -static bool CanUseCudnnSoftmax(const GPUContext& dev_ctx) { - if (dev_ctx.cudnn_handle() != nullptr) { - if (std::is_same::value) { -#if CUDNN_VERSION < 8100 - return false; -#endif - } - return true; - } - return false; -} - #if CUDNN_VERSION < 8100 template <> inline void SoftmaxForwardCudnnKernel( @@ -927,6 +914,25 @@ inline void SoftmaxBackwardCudnnKernel( } #endif +template +bool UseCudnnSoftmax(const GPUContext& ctx, int softmax_dim, bool last_dim) { + bool cudnn_available = ctx.cudnn_handle(); + if (!ctx.cudnn_handle()) { + if (std::is_same::value) { +#if CUDNN_VERSION < 8100 + cudnn_available = false; +#endif + } + } + constexpr int max_dim = 512; + if (!cudnn_available || !last_dim || + (softmax_dim <= max_dim && sizeof(T) <= 4)) { + return false; + } else { + return true; + } +} + template void SoftmaxForwardCUDAKernelDriver(const GPUContext& dev_ctx, const DenseTensor& x, @@ -941,10 +947,7 @@ void SoftmaxForwardCUDAKernelDriver(const GPUContext& dev_ctx, int dim = tensor_dims[1]; int D = tensor_dims[2]; - constexpr int max_dim = 512; - - if (D == 1 && - (!CanUseCudnnSoftmax(dev_ctx) || (dim <= max_dim && sizeof(T) <= 4))) { + if (D == 1 && !UseCudnnSoftmax(dev_ctx, dim, true)) { int dim_log2 = static_cast(Log2Ceil(dim)); int dim_ceil = 1 << dim_log2; int warp_size = (dim_ceil < 32) ? dim_ceil : 32; @@ -1016,10 +1019,7 @@ void SoftmaxBackwardCUDAKernelDriver(const GPUContext& dev_ctx, int dim = tensor_dims[1]; int D = tensor_dims[2]; - constexpr int max_dim = 512; - - if (D == 1 && - (!CanUseCudnnSoftmax(dev_ctx) || (dim <= max_dim && sizeof(T) <= 4))) { + if (D == 1 && !UseCudnnSoftmax(dev_ctx, dim, true)) { int dim_log2 = Log2Ceil(dim); int dim_ceil = 1 << dim_log2; int warp_size = (dim_ceil < 32) ? dim_ceil : 32; diff --git a/python/paddle/fluid/tests/unittests/CMakeLists.txt b/python/paddle/fluid/tests/unittests/CMakeLists.txt index a78c820e1e6..e0cd0c4bf4d 100755 --- a/python/paddle/fluid/tests/unittests/CMakeLists.txt +++ b/python/paddle/fluid/tests/unittests/CMakeLists.txt @@ -327,6 +327,7 @@ if ((NOT WITH_NCCL) AND (NOT WITH_RCCL)) endif() if(((NOT WITH_ROCM) AND (NOT WITH_GPU)) OR WIN32) + LIST(REMOVE_ITEM TEST_OPS test_fused_gate_attention_op) LIST(REMOVE_ITEM TEST_OPS test_boxps) endif() list(REMOVE_ITEM TEST_OPS test_seq_concat_op) # FIXME(helin): https://github.com/PaddlePaddle/Paddle/issues/8290 diff --git a/python/paddle/fluid/tests/unittests/test_fused_gate_attention_op.py b/python/paddle/fluid/tests/unittests/test_fused_gate_attention_op.py new file mode 100644 index 00000000000..6f9ba5f5e4e --- /dev/null +++ b/python/paddle/fluid/tests/unittests/test_fused_gate_attention_op.py @@ -0,0 +1,252 @@ +# Copyright (c) 2022 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. + +import numpy as np + +import paddle +import paddle.nn as nn +from paddle import tensor +import unittest +from op_test import OpTest, convert_float_to_uint16 +from test_sparse_attention_op import get_cuda_version +from paddle import _C_ops +from paddle.fluid.framework import default_main_program +from paddle.fluid import core + + +@unittest.skipIf(not core.is_compiled_with_cuda(), + "Paddle is not compiled with CUDA") +class TestFusedGateAttentionOp(OpTest): + def setUp(self): + self.__class__.op_type = "fused_gate_attention" + # use autograd to check grad in this unittest. + self.__class__.no_need_check_grad = True + self.config() + self.merge_qkv = self.q_dim == self.kv_dim + self.generate_input_data() + + def config(self): + self.dtype = "float32" + self.has_gating = True + self.batch_size = 1 + self.msa_len = 3 + self.res_len = 5 + self.q_dim = 6 + self.num_heads = 2 + self.key_dim = 4 + self.m_size = self.res_len + self.kv_dim = self.q_dim + self.out_dim = self.q_dim + self.bias_attr = True + + def generate_input_data(self): + def _random(shape): + if self.dtype == "bfloat16": + data = np.random.random(shape).astype("float32") + return convert_float_to_uint16(data) + else: + return np.random.random(shape).astype(self.dtype) + + np.random.seed(123) + self.query = _random( + (self.batch_size, self.msa_len, self.res_len, self.q_dim)) + self.q_weight = _random((self.q_dim, self.num_heads, self.key_dim)) + self.k_weight = _random((self.kv_dim, self.num_heads, self.key_dim)) + self.v_weight = _random((self.kv_dim, self.num_heads, self.key_dim)) + if self.merge_qkv: + self.key = None + # (3, self.num_heads, self.key_dim, self.q_dim) + q_weight_t = np.transpose(self.q_weight, axes=[1, 2, 0]) + k_weight_t = np.transpose(self.k_weight, axes=[1, 2, 0]) + v_weight_t = np.transpose(self.v_weight, axes=[1, 2, 0]) + self.qkv_weight = np.stack([q_weight_t, k_weight_t, v_weight_t]) + else: + self.key = _random( + (self.batch_size, self.msa_len, self.m_size, self.kv_dim)) + self.qkv_weight = None + + self.attn_mask = _random( + (self.batch_size, self.msa_len, 1, 1, self.m_size)) + + if self.bias_attr: + self.nonbatched_bias = _random( + (self.batch_size, 1, self.num_heads, self.res_len, self.m_size)) + + if self.has_gating: + self.gating_w = _random((self.q_dim, self.num_heads, self.key_dim)) + self.gating_b = _random((self.num_heads, self.key_dim)) + + self.output_w = _random((self.num_heads, self.key_dim, self.out_dim)) + self.output_b = _random((self.out_dim)) + + self.dout = _random( + (self.batch_size, self.msa_len, self.res_len, self.q_dim)) + + def get_reference_out(self): + paddle.disable_static(place=paddle.CUDAPlace(0)) + + query = paddle.to_tensor(self.query, stop_gradient=False) + key = query if self.merge_qkv else paddle.to_tensor( + self.key, stop_gradient=False) + q_weight = paddle.to_tensor(self.q_weight, stop_gradient=False) + k_weight = paddle.to_tensor(self.k_weight, stop_gradient=False) + v_weight = paddle.to_tensor(self.v_weight, stop_gradient=False) + src_mask = paddle.to_tensor(self.attn_mask, stop_gradient=True) + + c = self.key_dim**(-0.5) + # [batch_size, msa_len, num_heads, res_len, key_dim] + q = paddle.einsum('nbqa,ahc->nbqhc', query, q_weight) * c + # [batch_size, msa_len, num_heads, m_size, key_dim] + k = paddle.einsum('nbka,ahc->nbkhc', key, k_weight) + # [batch_size, msa_len, num_heads, m_size, key_dim] + v = paddle.einsum('nbka,ahc->nbkhc', key, v_weight) + + # [batch_size, msa_len, num_heads, res_len, m_size] + logits = paddle.einsum('nbqhc,nbkhc->nbhqk', q, k) # qk_out + logits = logits + src_mask + if self.bias_attr: + nonbatched_bias = paddle.to_tensor( + self.nonbatched_bias, stop_gradient=False) + logits = logits + nonbatched_bias + + weights = nn.functional.softmax(logits) # softmax_out + weighted_avg = paddle.einsum('nbhqk,nbkhc->nbqhc', weights, v) + + if self.has_gating: + gating_w = paddle.to_tensor(self.gating_w, stop_gradient=False) + gating_b = paddle.to_tensor(self.gating_b, stop_gradient=False) + gate_values = paddle.einsum('nbqc,chv->nbqhv', query, + gating_w) + gating_b + gate_values = nn.functional.sigmoid(gate_values) + weighted_avg = weighted_avg * gate_values + + output_b = paddle.to_tensor(self.output_b, stop_gradient=False) + output_w = paddle.to_tensor(self.output_w, stop_gradient=False) + + out = paddle.einsum('nbqhc,hco->nbqo', weighted_avg, + output_w) + output_b + paddle.autograd.backward( + [out], [paddle.to_tensor(self.dout)], retain_graph=True) + if self.merge_qkv: + return out, query.grad, None + else: + return out, query.grad, key.grad + + def get_fused_gate_attention_out(self): + paddle.disable_static(place=paddle.CUDAPlace(0)) + + query = paddle.to_tensor(self.query, stop_gradient=False) + if self.merge_qkv: + key = None + q_weight = None + k_weight = None + v_weight = None + qkv_weight = paddle.to_tensor(self.qkv_weight, stop_gradient=False) + else: + key = paddle.to_tensor(self.key, stop_gradient=False) + q_weight = paddle.to_tensor(self.q_weight, stop_gradient=False) + k_weight = paddle.to_tensor(self.k_weight, stop_gradient=False) + v_weight = paddle.to_tensor(self.v_weight, stop_gradient=False) + qkv_weight = None + + src_mask = paddle.to_tensor(self.attn_mask, stop_gradient=True) + + if self.bias_attr: + nonbatched_bias = paddle.to_tensor( + self.nonbatched_bias, stop_gradient=False) + else: + nonbatched_bias = None + if self.has_gating: + gating_w = paddle.to_tensor(self.gating_w, stop_gradient=False) + gating_b = paddle.to_tensor(self.gating_b, stop_gradient=False) + else: + gating_w = None + gating_b = None + + output_w = paddle.to_tensor(self.output_w, stop_gradient=False) + output_b = paddle.to_tensor(self.output_b, stop_gradient=False) + + _, _, _, _, _, _, _, out = _C_ops.fused_gate_attention( + query, key, q_weight, k_weight, v_weight, qkv_weight, + nonbatched_bias, src_mask, gating_w, gating_b, output_w, output_b, + 'has_gating', self.has_gating, 'merge_qkv', self.merge_qkv) + + paddle.autograd.backward( + [out], [paddle.to_tensor(self.dout)], retain_graph=True) + if key is not None: + return out, query.grad, key.grad + else: + return out, query.grad, None + + def check_output_and_grad(self, atol, rtol): + out_ref, query_grad_ref, key_grad_ref = self.get_reference_out() + out, query_grad, key_grad = self.get_fused_gate_attention_out() + np.testing.assert_allclose(out_ref, out.numpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose( + query_grad_ref, query_grad.numpy(), atol=atol, rtol=rtol) + if key_grad_ref is not None and key_grad is not None: + np.testing.assert_allclose( + key_grad_ref, key_grad.numpy(), atol=atol, rtol=rtol) + + def test_output_and_grad(self): + self.check_output_and_grad(atol=1e-5, rtol=1e-5) + + +class TestSeparatedQKVCase(TestFusedGateAttentionOp): + def config(self): + self.dtype = "float32" + self.has_gating = False + self.batch_size = 1 + self.msa_len = 3 + self.res_len = 5 + self.q_dim = 6 + self.num_heads = 2 + self.key_dim = 4 + self.m_size = 4 + self.kv_dim = 2 + self.out_dim = self.q_dim + self.bias_attr = False + + +class TestMergeQKVNoBiasGatingCase(TestFusedGateAttentionOp): + def config(self): + super().config() + self.has_gating = False + self.bias_attr = False + + +class TestMergeQKVFp16Case(TestFusedGateAttentionOp): + def config(self): + super().config() + self.dtype = "float16" + + def test_output_and_grad(self): + self.check_output_and_grad(atol=1e-1, rtol=1e-5) + + +@unittest.skipIf( + not core.is_compiled_with_cuda() or get_cuda_version() < 11000, + "core is not compiled with CUDA and cuda version need larger than or equal to 11.3" +) +class TestMergeQKVBF16Case(TestFusedGateAttentionOp): + def config(self): + super().config() + self.dtype = "bfloat16" + + def test_output_and_grad(self): + self.check_output_and_grad(atol=1e-1, rtol=1e-3) + + +if __name__ == "__main__": + unittest.main() -- GitLab