[cherry-pick]add op: fused_feedforward(forward) (#36729)

This is a fusion operator to compute feed forward layer in transformer model architecture.

cmake/operators.cmake

@@ -217,7 +217,7 @@ function(op_library TARGET)
 "fusion_transpose_flatten_concat_op" "fusion_conv_inception_op"
 "sync_batch_norm_op" "sparse_attention_op"  "dgc_op" "fused_fc_elementwise_layernorm_op"
 "skip_layernorm_op" "multihead_matmul_op" "fusion_group_op" "fused_bn_activation_op" "fused_embedding_eltwise_layernorm_op" "fusion_gru_op" "fusion_lstm_op"
-"fused_bn_add_activation_op" "fused_attention_op")
+"fused_bn_add_activation_op" "fused_attention_op" "fused_feedforward_op")
         if ("${TARGET}" STREQUAL "${manual_pybind_op}")
             set(pybind_flag 1)
         endif()

paddle/fluid/operators/fused/CMakeLists.txt

@@ -17,6 +17,7 @@ register_operators(EXCLUDES
     fusion_lstm_op
     fused_bn_add_activation_op
     fused_attention_op
+    fused_feedforward_op
     fused_transformer_op)
 
 # fusion_gru_op does not have CUDA kernel
@@ -78,6 +79,11 @@ if (WITH_GPU OR WITH_ROCM)
         nv_test(test_fused_residual_dropout_bias SRCS fused_residual_dropout_bias_test.cu DEPS tensor op_registry dropout_op layer_norm_op device_context generator memory)
         nv_test(test_fused_dropout_act_bias SRCS fused_dropout_act_bias_test.cu DEPS tensor op_registry dropout_op layer_norm_op device_context generator memory)
         nv_test(test_fused_layernorm_residual_dropout_bias SRCS fused_layernorm_residual_dropout_bias_test.cu DEPS tensor op_registry dropout_op layer_norm_op device_context generator memory)
+
+
+        op_library(fused_feedforward_op)
+        file(APPEND ${pybind_file} "USE_CUDA_ONLY_OP(fused_feedforward);\n")
+
         # fused_attention_op
         op_library(fused_attention_op)
         file(APPEND ${pybind_file} "USE_CUDA_ONLY_OP(fused_attention);\n")

paddle/fluid/operators/fused/fused_dropout_common.h

@@ -110,27 +110,34 @@ inline __device__ void CalculateDBias(const T *tmp_sum, T *dbias,
   }
   __syncthreads();
   // reduce sum
-  T sum = static_cast<T>(0);
+  T sum[2] = {static_cast<T>(0)};
   int tid = threadIdx.y * blockDim.x + threadIdx.x;
   int x = tid >> 5;  // warp id
   int y = tid & 31;  // thread id on warp 0~31
 
   // need BlockSizeX * VecSize warps
-  if (x < BlockSizeX * VecSize) {
+  for (int j = x; j < BlockSizeX * VecSize; j += 32) {
 // reduce 128 to 32
 #pragma unroll
     for (int i = 0; i < (BlockSizeY >> 5); i++) {
-      sum += cache[x][y + i * 32];
+      sum[(j >> 5)] += cache[j][y + i * 32];
     }
   }
 
+  int reduce_num_pre_thread = (BlockSizeX * VecSize + 31) / 32;
   // reduce 32 to 1
-  sum = WarpReduceSum(sum);
+  for (int i = 0; i < reduce_num_pre_thread; i++) {
+    sum[i] = WarpReduceSum(sum[i]);
+  }
 
   // save sum to dbias
-  int bias_id = blockIdx.x * blockDim.x * VecSize + x;
-  if (y == 0 && x < VecSize * BlockSizeX && bias_id < cols) {
-    dbias[bias_id] = sum;
+  if (y == 0 && x < BlockSizeX * VecSize) {
+    for (int i = 0; i < reduce_num_pre_thread; i++) {
+      int bias_id = blockIdx.x * BlockSizeX * VecSize + x + i * 32;
+      if (bias_id < cols) {
+        dbias[bias_id] = sum[i];
+      }
+    }
   }
 }
 

paddle/fluid/operators/fused/fused_feedforward_op.cc

+/* 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.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <algorithm>
+#include <utility>
+#include <vector>
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/op_version_registry.h"
+#include "paddle/fluid/operators/math/blas.h"
+#include "paddle/fluid/operators/matmul_v2_op.h"
+
+namespace paddle {
+namespace operators {
+using Tensor = framework::Tensor;
+
+class FusedFeedForwardOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+ protected:
+  void InferShape(framework::InferShapeContext *context) const override {
+    OP_INOUT_CHECK(context->HasInput("X"), "Input", "X", "fused_feedforward");
+    OP_INOUT_CHECK(context->HasInput("Linear1Weight"), "Input", "Linear1Weight",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasInput("Linear2Weight"), "Input", "Linear2Weight",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Out"), "Output", "Out",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Dropout1Mask"), "Output", "Dropout1Mask",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Dropout2Mask"), "Output", "Dropout2Mask",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Ln1Mean"), "Output", "Ln1Mean",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Ln1Variance"), "Output", "Ln1Variance",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Ln2Mean"), "Output", "Ln2Mean",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Ln2Variance"), "Output", "Ln2Variance",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Linear1Out"), "Output", "Linear1Out",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Ln1Out"), "Output", "Ln1Out",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Dropout1Out"), "Output", "Dropout1Out",
+                   "fused_feedforward");
+    OP_INOUT_CHECK(context->HasOutput("Dropout2Out"), "Output", "Dropout2Out",
+                   "fused_feedforward");
+
+    auto dim_x = context->GetInputDim("X");
+    auto mat_dim_x =
+        math::CreateMatrixDescriptor(RowMatrixFromVector(dim_x), 0, false);
+    // verify for the pre layer_norm, the feature size must be larger than 1
+    PADDLE_ENFORCE_GT(
+        mat_dim_x.width_, static_cast<size_t>(1),
+        platform::errors::InvalidArgument("Product from the X shape[1] to "
+                                          "shape[n-1] must be larger than 1!"));
+    auto dim_Linear1Weight = context->GetInputDim("Linear1Weight");
+    auto tmp_dim_x = dim_x;
+    tmp_dim_x[dim_x.size() - 1] =
+        dim_Linear1Weight[dim_Linear1Weight.size() - 1];
+    context->SetOutputDim("Out", dim_x);
+    if (context->Attrs().Get<bool>("dropout1_is_test") == false) {
+      context->SetOutputDim("Dropout1Mask", tmp_dim_x);
+    }
+    context->SetOutputDim("Dropout1Out", tmp_dim_x);
+    context->SetOutputDim("Linear1Out", tmp_dim_x);
+    context->SetOutputDim("Ln1Out", dim_x);
+    context->SetOutputDim("Dropout2Out", dim_x);
+
+    if (context->Attrs().Get<bool>("dropout2_is_test") == false) {
+      context->SetOutputDim("Dropout2Mask", dim_x);
+    }
+    framework::DDim mean_dim =
+        framework::make_ddim({mat_dim_x.batch_size_ * mat_dim_x.height_});
+    context->SetOutputDim("Ln1Mean", mean_dim);
+    context->SetOutputDim("Ln1Variance", mean_dim);
+    context->SetOutputDim("Ln2Mean", mean_dim);
+    context->SetOutputDim("Ln2Variance", mean_dim);
+    context->ShareLoD("X", "Out");
+  }
+
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext &ctx) const override {
+    return framework::OpKernelType(
+        OperatorWithKernel::IndicateVarDataType(ctx, "X"), ctx.GetPlace());
+  }
+};
+
+class FusedFeedForwardOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X", "The input of FusedFeedForward op");
+    AddInput(
+        "Dropout1Seed",
+        "The seed of first dropout op, it has higher priority than the attr "
+        "fix_seed and seed")
+        .AsDispensable();
+    AddInput(
+        "Dropout2Seed",
+        "The seed of second dropout op, it has higher priority than the attr "
+        "fix_seed and seed")
+        .AsDispensable();
+
+    AddInput("Linear1Weight", "The linear1 weight of FusedFeedForward op");
+    AddInput("Linear1Bias", "The linear1 bias of FusedFeedForward op")
+        .AsDispensable();
+    AddInput("Linear2Weight", "The linear2 weight of FusedFeedForward op");
+    AddInput("Linear2Bias", "The linear2 bias input of FusedFeedForward op")
+        .AsDispensable();
+    AddInput("Ln1Scale", "The layer_norm1 scale of FusedFeedForward op")
+        .AsDispensable();
+    AddInput("Ln1Bias", "The layer_norm1 bias of FusedFeedForward op")
+        .AsDispensable();
+    AddInput("Ln2Scale", "The layer_norm2 scale of FusedFeedForward op")
+        .AsDispensable();
+    AddInput("Ln2Bias", "The layer_norm2 bias of FusedFeedForward op")
+        .AsDispensable();
+    AddOutput("Out", "The output of FusedFeedForward op");
+    AddOutput("Dropout1Mask", "The mask of dropout1").AsIntermediate();
+    AddOutput("Dropout2Mask", "The mask of dropout2").AsIntermediate();
+    AddOutput("Ln1Mean", "The mean of layer_norm1").AsIntermediate();
+    AddOutput("Ln1Variance", "The variance of layer_norm1").AsIntermediate();
+    AddOutput("Ln2Mean", "The mean of layer_nomr2").AsIntermediate();
+    AddOutput("Ln2Variance", "The variance of layer_norm2").AsIntermediate();
+    AddOutput("Linear1Out", "The output of linear1").AsIntermediate();
+    AddOutput("Ln1Out", "The output of layer_norm1").AsIntermediate();
+    AddOutput("Dropout1Out", "The output of dropout1").AsIntermediate();
+    AddOutput("Dropout2Out", "The output of dropout2").AsIntermediate();
+
+    AddAttr<bool>("pre_layer_norm", "true is pre layernorm").SetDefault(false);
+    AddAttr<float>("ln1_epsilon", "epsilon of pre layer_norm")
+        .SetDefault(1e-5f);
+    AddAttr<float>("ln2_epsilon", "epsilon of post layer_norm")
+        .SetDefault(1e-5f);
+    AddAttr<std::string>("act_method", "act_method").SetDefault("gelu");
+    AddAttr<float>("dropout1_rate", "the dropout rate of first dropout")
+        .SetDefault(.5f)
+        .AddCustomChecker([](const float &drop_p) {
+          PADDLE_ENFORCE_EQ(
+              drop_p >= 0.0f && drop_p <= 1.0f, true,
+              platform::errors::InvalidArgument(
+                  "'dropout1_rate' must be between 0.0 and 1.0."));
+        });
+    AddAttr<float>("dropout2_rate", "the dropout rate of second dropout")
+        .SetDefault(.5f)
+        .AddCustomChecker([](const float &drop_p) {
+          PADDLE_ENFORCE_EQ(
+              drop_p >= 0.0f && drop_p <= 1.0f, true,
+              platform::errors::InvalidArgument(
+                  "'dropout2_rate' must be between 0.0 and 1.0."));
+        });
+    AddAttr<std::string>("dropout1_implementation",
+                         "the dropout implementation of first dropout")
+        .SetDefault("downgrade_in_infer")
+        .AddCustomChecker([](const std::string &type) {
+          PADDLE_ENFORCE_EQ(
+              type == "downgrade_in_infer" || type == "upscale_in_train", true,
+              platform::errors::InvalidArgument(
+                  "dropout1_implementation can only be downgrade_in_infer or "
+                  "upscale_in_train"));
+        });
+    AddAttr<std::string>("dropout2_implementation",
+                         "the dropout implementation of second dropout")
+        .SetDefault("downgrade_in_infer")
+        .AddCustomChecker([](const std::string &type) {
+          PADDLE_ENFORCE_EQ(
+              type == "downgrade_in_infer" || type == "upscale_in_train", true,
+              platform::errors::InvalidArgument(
+                  "dropout2_implementation can only be downgrade_in_infer or "
+                  "upscale_in_train"));
+        });
+    AddAttr<bool>("dropout1_is_test", "the is_test of first dropout")
+        .SetDefault(false);
+    AddAttr<bool>("dropout2_is_test", "the is_test of second dropout")
+        .SetDefault(false);
+    AddAttr<bool>("dropout1_fix_seed", "the is_test of first dropout")
+        .SetDefault(false);
+    AddAttr<bool>("dropout2_fix_seed", "the is_test of second dropout")
+        .SetDefault(false);
+    AddAttr<int>("dropout1_seed", "Dropout1 random seed.").SetDefault(0);
+    AddAttr<int>("dropout2_seed", "Dropout2 random seed.").SetDefault(0);
+    AddComment(R"DOC(
+        the function of fused_feedforward operator is the same as the following pseudo code:
+        residual = src;
+        ln1_out = src;
+        if(pre_layer_norm){
+            ln1_out = layer_norm(src);
+        }
+        out = linear(dropout(activation(dropout(linear(ln1_out)))));
+        if(!pre_layer_norm) {
+            out = layer_norm(out);
+        }
+        )DOC");
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OPERATOR(fused_feedforward, ops::FusedFeedForwardOp,
+                  ops::FusedFeedForwardOpMaker);

paddle/fluid/operators/fused/fused_feedforward_op.cu

+/* 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.
+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/op_version_registry.h"
+#include "paddle/fluid/operators/math/blas.h"
+#include "paddle/fluid/operators/matmul_v2_op.h"
+
+#include "paddle/fluid/operators/fused/fused_dropout_helper.h"
+#include "paddle/fluid/operators/layer_norm_kernel.cu.h"
+
+namespace paddle {
+namespace operators {
+
+using Tensor = framework::Tensor;
+
+template <typename DeviceContext, typename T>
+class FusedFeedForwardKernel : public framework::OpKernel<T> {
+ public:
+  void MatMul(const platform::CUDADeviceContext& ctx,
+              const framework::Tensor& a, const framework::Tensor& b,
+              framework::Tensor* c) const {
+    auto blas = math::GetBlas<DeviceContext, T>(ctx);
+    auto a_2d = FoldInitDims(a);
+    auto b_2d = FoldInitDims(b);
+    auto mat_dim_a = math::CreateMatrixDescriptor(a_2d.dims(), 0, false);
+    auto mat_dim_b = math::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));
+  }
+
+  void FFN(const framework::Tensor& x, const framework::Tensor& linear1_weight,
+           const framework::Tensor* linear1_bias,
+           const framework::Tensor& linear2_weight,
+           const framework::Tensor* linear2_bias,
+           const framework::Tensor* ln1_scale,
+           const framework::Tensor* ln1_bias,
+           const framework::Tensor* ln2_scale,
+           const framework::Tensor* ln2_bias, framework::Tensor* out,
+           framework::Tensor* dropout1_mask, framework::Tensor* dropout2_mask,
+           framework::Tensor* ln1_mean, framework::Tensor* ln1_variance,
+           framework::Tensor* ln2_mean, framework::Tensor* ln2_variance,
+           framework::Tensor* linear1_out, framework::Tensor* ln1_out,
+           framework::Tensor* dropout1_out, framework::Tensor* 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 DropoutParam& dropout_param1,
+           const DropoutParam& dropout_param2,
+           const platform::CUDADeviceContext& ctx) const {
+    FusedDropoutLayerNormHelper<T, uint8_t> pre_layernorm_helper(
+        bsz_seq, d_model, epsilon1);
+    FusedDropoutHelper<T, uint8_t> fused_act_dropout_helper(
+        ctx, bsz_seq, dim_feedforward, dropout_param1);
+    FusedDropoutLayerNormHelper<T, uint8_t> fused_dropout_layernorm_helper(
+        ctx, bsz_seq, d_model, dropout_param2, epsilon2);
+
+    auto place = ctx.GetPlace();
+    using U = LayerNormParamType<T>;
+    const framework::Tensor* 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(
+          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(ctx, *in, linear1_weight, linear1_out);
+    fused_act_dropout_helper.DropoutActBias(
+        ctx, linear1_out->data<T>(), linear1_bias_ptr, act_method,
+        dropout1_out->data<T>(), dropout1_mask->data<uint8_t>());
+    framework::Tensor linear2_out;
+    linear2_out.mutable_data<T>({bsz_seq, d_model}, place);
+    MatMul(ctx, *dropout1_out, linear2_weight, &linear2_out);
+    if (!pre_layer_norm) {
+      fused_dropout_layernorm_helper.LayernormResidualDropoutBias(
+          ctx, linear2_out.data<T>(), x.data<T>(), 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(
+          ctx, linear2_out.data<T>(), x.data<T>(), linear2_bias_ptr,
+          out->data<T>(), dropout2_mask->data<uint8_t>());
+    }
+  }
+
+  void Compute(const framework::ExecutionContext& context) const override {
+    auto* x = context.Input<framework::Tensor>("X");
+    auto* linear1_weight = context.Input<framework::Tensor>("Linear1Weight");
+    auto* linear1_bias = context.Input<framework::Tensor>("Linear1Bias");
+    auto* linear2_weight = context.Input<framework::Tensor>("Linear2Weight");
+    auto* linear2_bias = context.Input<framework::Tensor>("Linear2Bias");
+    auto* ln1_scale = context.Input<framework::Tensor>("Ln1Scale");
+    auto* ln1_bias = context.Input<framework::Tensor>("Ln1Bias");
+    auto* ln2_scale = context.Input<framework::Tensor>("Ln2Scale");
+    auto* ln2_bias = context.Input<framework::Tensor>("Ln2Bias");
+
+    auto* ln1_mean = context.Output<framework::Tensor>("Ln1Mean");
+    auto* ln1_variance = context.Output<framework::Tensor>("Ln1Variance");
+    auto* ln2_mean = context.Output<framework::Tensor>("Ln2Mean");
+    auto* ln2_variance = context.Output<framework::Tensor>("Ln2Variance");
+    auto* out = context.Output<framework::Tensor>("Out");
+    auto* dropout1_mask = context.Output<framework::Tensor>("Dropout1Mask");
+    auto* dropout2_mask = context.Output<framework::Tensor>("Dropout2Mask");
+    auto* linear1_out = context.Output<framework::Tensor>("Linear1Out");
+    auto* ln1_out = context.Output<framework::Tensor>("Ln1Out");
+    auto* dropout1_out = context.Output<framework::Tensor>("Dropout1Out");
+    auto* dropout2_out = context.Output<framework::Tensor>("Dropout2Out");
+
+    const std::string act_method = context.Attr<std::string>("act_method");
+
+    const bool pre_layer_norm = context.Attr<bool>("pre_layer_norm");
+    const float epsilon1 = context.Attr<float>("ln1_epsilon");
+    const float epsilon2 = context.Attr<float>("ln2_epsilon");
+
+    DropoutParam dropout_param1(context, 1);
+    DropoutParam dropout_param2(context, 2);
+
+    using U = LayerNormParamType<T>;
+    auto place = context.GetPlace();
+    out->mutable_data<T>(place);
+    dropout1_mask->mutable_data<uint8_t>(place);
+    dropout2_mask->mutable_data<uint8_t>(place);
+    ln1_mean->mutable_data<U>(place);
+    ln1_variance->mutable_data<U>(place);
+    ln2_mean->mutable_data<U>(place);
+    ln2_variance->mutable_data<U>(place);
+    linear1_out->mutable_data<T>(place);
+    ln1_out->mutable_data<T>(place);
+    dropout1_out->mutable_data<T>(place);
+    dropout2_out->mutable_data<T>(place);
+
+    auto x_dim = x->dims();
+    auto mat_dim_x =
+        math::CreateMatrixDescriptor(RowMatrixFromVector(x_dim), 0, false);
+
+    auto dim = linear1_weight->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_;
+
+    FFN(*x, *linear1_weight, linear1_bias, *linear2_weight, linear2_bias,
+        ln1_scale, ln1_bias, ln2_scale, ln2_bias, 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, epsilon1, epsilon2,
+        dropout_param1, dropout_param2, context.cuda_device_context());
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OP_CUDA_KERNEL(
+    fused_feedforward,
+    ops::FusedFeedForwardKernel<paddle::platform::CUDADeviceContext, float>,
+    ops::FusedFeedForwardKernel<paddle::platform::CUDADeviceContext, double>,
+    ops::FusedFeedForwardKernel<paddle::platform::CUDADeviceContext,
+                                paddle::platform::float16>);

paddle/fluid/pybind/op_function_generator.cc

@@ -72,6 +72,9 @@ std::map<std::string, std::set<std::string>> op_ins_map = {
     {"sparse_momentum", {"Param", "Grad", "Velocity", "Index", "LearningRate"}},
     {"rnn", {"Input", "PreState", "WeightList", "SequenceLength"}},
     {"run_program", {"X", "Params"}},
+    {"fused_feedforward",
+     {"Dropout1Seed", "Dropout2Seed", "Linear1Bias", "Linear2Bias", "Ln1Scale",
+      "Ln1Bias", "Ln2Scale", "Ln2Bias"}},
     {"matrix_rank", {"X", "TolTensor"}},
     {"adam",
      {"Param", "Grad", "LearningRate", "Moment1", "Moment2", "Beta1Pow",

python/paddle/fluid/tests/unittests/CMakeLists.txt

@@ -91,6 +91,8 @@ foreach(TEST_OP ${MIXED_DIST_TEST_OPS})
 endforeach()
 
 if(NOT WITH_GPU)
+
+    LIST(REMOVE_ITEM TEST_OPS test_fused_feedforward_op)
     LIST(REMOVE_ITEM TEST_OPS test_fused_attention_op)
 endif()
 

python/paddle/fluid/tests/unittests/test_fused_feedforward_op.py

+#   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.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+
+import paddle
+import paddle.fluid as fluid
+import paddle.fluid.core as core
+from paddle.nn.layer import transformer
+import paddle.nn.functional as F
+from paddle.nn.layer.norm import LayerNorm
+from paddle.nn.layer.common import Linear, Dropout
+import unittest
+from op_test import OpTest
+
+
+class TestFusedFFNOp(OpTest):
+    def getDtype(self):
+        self.dtype = "float32"
+        self.layer_norm_dtype = "float32"
+
+    def getShape(self):
+        self.batch_size = np.random.randint(1, 64)
+        self.query_length = np.random.randint(32, 256)
+        self.d_model = np.random.randint(32, 1024)
+        self.dim_feedforward = np.random.randint(32, 1024)
+
+    def getDiff(self):
+        self.rtol = 1e-3
+        self.atol = 1e-4
+
+    def getActivation(self):
+        self.act_method = "gelu"
+
+    def getNormalizeBefore(self):
+        self.pre_layer_norm = False
+
+    def setUp(self):
+        paddle.disable_static()
+        self.__class__.op_type = "fused_feedforward"
+        self.getDtype()
+        self.getShape()
+        self.getDiff()
+        self.getActivation()
+        self.getNormalizeBefore()
+        paddle.set_default_dtype(self.dtype)
+        self.weight_attr = None
+        self.bias_attr = None
+
+        self.weight_attrs = transformer._convert_param_attr_to_list(
+            self.weight_attr, 2)
+        self.bias_attrs = transformer._convert_param_attr_to_list(
+            self.bias_attr, 2)
+        self.linear1 = Linear(
+            self.d_model,
+            self.dim_feedforward,
+            self.weight_attrs[1],
+            bias_attr=self.bias_attrs[1])
+        self.linear2 = Linear(
+            self.dim_feedforward,
+            self.d_model,
+            self.weight_attrs[1],
+            bias_attr=self.bias_attrs[1])
+
+        paddle.set_default_dtype(self.layer_norm_dtype)
+        self.norm1 = LayerNorm(self.d_model)
+        self.norm2 = LayerNorm(self.d_model)
+        self.dropout = Dropout(0.0, mode="upscale_in_train")
+        self.dropout1 = Dropout(0.0, mode="upscale_in_train")
+        self.dropout2 = Dropout(0.0, mode="upscale_in_train")
+        self.activation = getattr(F, self.act_method)
+
+        self.src = np.random.random((self.batch_size, self.query_length,
+                                     self.d_model)).astype(self.dtype)
+
+    def Base(self):
+        paddle.disable_static()
+        tensor_src = paddle.to_tensor(self.src, stop_gradient=False)
+        residual = paddle.to_tensor(self.src)
+        if self.pre_layer_norm:
+            ln1_out = self.norm1(tensor_src)
+            linear2_out = self.linear2(
+                self.dropout(self.activation(self.linear1(ln1_out))))
+            dropout2_out = residual + self.dropout2(linear2_out)
+        else:
+            linear2_out = self.linear2(
+                self.dropout(self.activation(self.linear1(tensor_src))))
+            dropout2_out = residual + self.dropout2(linear2_out)
+            dropout2_out = self.norm2(dropout2_out)
+        return dropout2_out
+
+    def FusedFFN(self):
+        paddle.disable_static()
+        linear1_weight = paddle.to_tensor(
+            self.linear1.weight, stop_gradient=False)
+        linear1_bias = paddle.to_tensor(self.linear1.bias, stop_gradient=False)
+        linear2_weight = paddle.to_tensor(
+            self.linear2.weight, stop_gradient=False)
+        linear2_bias = paddle.to_tensor(self.linear2.bias, stop_gradient=False)
+        ln1_scale = paddle.to_tensor(self.norm1.weight, stop_gradient=False)
+        ln1_bias = paddle.to_tensor(self.norm1.bias, stop_gradient=False)
+        ln2_scale = paddle.to_tensor(self.norm2.weight, stop_gradient=False)
+        ln2_bias = paddle.to_tensor(self.norm2.bias, stop_gradient=False)
+        x = paddle.to_tensor(self.src, stop_gradient=False)
+        out = F.fused_feedforward(
+            x,
+            linear1_weight,
+            linear2_weight,
+            linear1_bias,
+            linear2_bias,
+            ln1_scale,
+            ln1_bias,
+            ln2_scale,
+            ln2_bias,
+            0.0,
+            0.0,
+            activation=self.act_method,
+            pre_layer_norm=self.pre_layer_norm)
+        return out
+
+    def test_fused_ffn(self):
+        base_out = self.Base()
+        fused_out = self.FusedFFN()
+        np.testing.assert_allclose(
+            base_out.numpy(), fused_out.numpy(), rtol=self.rtol, atol=self.atol)
+
+
+class TestFusedFFNOpFp16(TestFusedFFNOp):
+    def getDtype(self):
+        self.dtype = "float16"
+        self.layer_norm_dtype = "float32"
+
+    def getDiff(self):
+        self.rtol = 1e-1
+        self.atol = 1e-2
+
+    def getShape(self):
+        self.batch_size = 8
+        self.query_length = 128
+        self.d_model = 512
+        self.dim_feedforward = 512
+
+
+class TestFusedFFNOpFp64(TestFusedFFNOp):
+    def getDtype(self):
+        self.dtype = "float64"
+        self.layer_norm_dtype = "float64"
+
+
+class TestFusedFFNOpActivation(TestFusedFFNOp):
+    def getActivation(self):
+        self.act_method = "relu"
+
+
+class TestFusedFFNOpNormalizeBefore(TestFusedFFNOp):
+    def getNormalizeBefore(self):
+        self.pre_layer_norm = True
+
+    def getShape(self):
+        self.batch_size = 1
+        self.query_length = 1
+        self.d_model = 8
+        self.dim_feedforward = 8
+
+
+class APITestStaticFusedFFN(unittest.TestCase):
+    def test_static(self):
+        paddle.enable_static()
+        dtype = "float32"
+        layer_norm_dtype = "float32"
+        batch_size = 1
+        d_model = 8
+        dim_feedforward = 8
+
+        x = paddle.static.data(
+            name='x', shape=[batch_size, d_model, dim_feedforward], dtype=dtype)
+        linear1_weight = paddle.static.data(
+            name='linear1_weight',
+            shape=[d_model, dim_feedforward],
+            dtype=dtype)
+        linear1_bias = paddle.static.data(
+            name='linear1_bias', shape=[dim_feedforward])
+        linear2_weight = paddle.static.data(
+            name='linear2_weight',
+            shape=[dim_feedforward, d_model],
+            dtype=dtype)
+        linear2_bias = paddle.static.data(name='linear2_bias', shape=[d_model])
+        ln1_scale = paddle.static.data(name='ln1_scale', shape=[d_model])
+        ln1_bias = paddle.static.data(name='ln1_scale', shape=[d_model])
+        ln2_scale = paddle.static.data(name='ln2_scale', shape=[d_model])
+        ln2_bias = paddle.static.data(name='ln2_scale', shape=[d_model])
+
+        fused_out = F.fused_feedforward(
+            x,
+            linear1_weight,
+            linear2_weight,
+            linear1_bias,
+            linear2_bias,
+            ln1_scale,
+            ln1_bias,
+            ln2_scale,
+            ln2_bias,
+            0.0,
+            0.0,
+            activation="relu",
+            pre_layer_norm=False)
+
+        ######base ffn######
+        linear1_out = F.linear(x, linear1_weight, linear1_bias)
+        act_out = F.relu(linear1_out)
+        dropout1_out = F.dropout(x=act_out, p=0.0, training=False)
+        linear2_out = F.linear(dropout1_out, linear2_weight, linear2_bias)
+        dropout2_out = x + F.dropout(x=linear2_out, p=0.0, training=False)
+        ln_out = F.layer_norm(
+            dropout2_out,
+            normalized_shape=list([d_model]),
+            weight=ln2_scale,
+            bias=ln2_bias)
+        ######base ffn######
+
+        exe = paddle.static.Executor(paddle.CUDAPlace(0))
+
+        x_data = np.random.random(
+            (batch_size, d_model, dim_feedforward)).astype(dtype)
+        linear1_weight_data = np.random.random(
+            (d_model, dim_feedforward)).astype(dtype)
+        linear1_bias_data = np.zeros((dim_feedforward)).astype(dtype)
+        linear2_weight_data = np.random.random(
+            (dim_feedforward, d_model)).astype(dtype)
+        linear2_bias_data = np.zeros((d_model)).astype(dtype)
+
+        ln1_scale_data = np.ones((d_model)).astype(layer_norm_dtype)
+        ln1_bias_data = np.zeros((d_model)).astype(layer_norm_dtype)
+        ln2_scale_data = np.ones((d_model)).astype(layer_norm_dtype)
+        ln2_bias_data = np.zeros((d_model)).astype(layer_norm_dtype)
+
+        res_list = [fused_out, ln_out]
+        real_res = []
+
+        for res in res_list:
+            fetch = exe.run(feed={
+                'x': x_data,
+                'linear1_weight': linear1_weight_data,
+                'linear1_bias': linear1_bias_data,
+                'linear2_weight': linear2_weight_data,
+                'linear2_bias': linear2_bias_data,
+                'ln1_scale': ln1_scale_data,
+                'ln1_bias': ln1_bias_data,
+                'ln2_scale': ln2_scale_data,
+                'ln2_bias': ln2_bias_data
+            },
+                            fetch_list=[res])
+            real_res.append(fetch)
+        self.assertTrue(
+            np.allclose(
+                real_res[0], real_res[1], atol=1e-5),
+            "two value is check diff")
+
+
+class TestFusedFFNOpError(unittest.TestCase):
+    def test_errors(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(paddle.static.Program(),
+                                         paddle.static.Program()):
+
+            def test_dtype():
+                x = paddle.static.data(
+                    name='x', shape=[1, 10, 10], dtype="int32")
+                linear1_weight = paddle.static.data(
+                    name='linear1_weight', shape=[1, 10, 10], dtype="float32")
+                linear2_weight = paddle.static.data(
+                    name='linear2_weight', shape=[1, 10, 10], dtype="float32")
+                paddle.nn.functional.fused_feedforward(x, linear1_weight,
+                                                       linear2_weight)
+
+            self.assertRaises(TypeError, test_dtype)
+
+            def test_dropout_rate_type():
+                x = paddle.static.data(
+                    name='x1', shape=[1, 10, 10], dtype="float32")
+                linear1_weight = paddle.static.data(
+                    name='linear1_weight1', shape=[10, 10], dtype="float32")
+                linear2_weight = paddle.static.data(
+                    name='linear2_weight1', shape=[10, 10], dtype="float32")
+                paddle.nn.functional.fused_feedforward(
+                    x, linear1_weight, linear2_weight, dropout1_rate="a")
+
+            self.assertRaises(TypeError, test_dropout_rate_type)
+
+            def test_dropout_rate_value():
+                x = paddle.static.data(
+                    name='x2', shape=[1, 10, 10], dtype="float32")
+                linear1_weight = paddle.static.data(
+                    name='linear1_weight2', shape=[10, 10], dtype="float32")
+                linear2_weight = paddle.static.data(
+                    name='linear2_weight2', shape=[10, 10], dtype="float32")
+                paddle.nn.functional.fused_feedforward(
+                    x, linear1_weight, linear2_weight, dropout2_rate=-1)
+
+            self.assertRaises(ValueError, test_dropout_rate_value)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/paddle/nn/functional/__init__.py

@@ -110,6 +110,7 @@ from .vision import grid_sample  # noqa: F401
 from .vision import pixel_shuffle  # noqa: F401
 from .input import one_hot  # noqa: F401
 from .input import embedding  # noqa: F401
+from .fused_transformer import fused_feedforward  # noqa: F401
 from ...fluid.layers import gather_tree  # noqa: F401
 from ...fluid.layers import temporal_shift  # noqa: F401
 
@@ -210,6 +211,7 @@ __all__ = [     #noqa
            'layer_norm',
            'instance_norm',
            'class_center_sample',
+            'fused_feedforward',
            'fused_multi_head_attention',
            'sparse_attention',
 ]

python/paddle/nn/functional/fused_transformer.py

@@ -12,13 +12,166 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import paddle
+from ...fluid.layer_helper import LayerHelper
 from ...fluid.framework import in_dygraph_mode
+from ...fluid.data_feeder import check_variable_and_dtype, check_dtype
 from paddle import _C_ops
 
 __all__ = []
 
 
+def _verify_dropout_rate(dropout_rate):
+    if not isinstance(dropout_rate, (float, int)):
+        raise TypeError("dropout_rate argument should be a number")
+    if dropout_rate < 0 or dropout_rate > 1:
+        raise ValueError("dropout_rate argument should between 0 and 1")
+
+
+def fused_feedforward(x,
+                      linear1_weight,
+                      linear2_weight,
+                      linear1_bias=None,
+                      linear2_bias=None,
+                      ln1_scale=None,
+                      ln1_bias=None,
+                      ln2_scale=None,
+                      ln2_bias=None,
+                      dropout1_rate=0.5,
+                      dropout2_rate=0.5,
+                      activation="relu",
+                      ln1_epsilon=1e-5,
+                      ln2_epsilon=1e-5,
+                      pre_layer_norm=False,
+                      name=None):
+    """
+    This is a fusion operator to compute feed forward layer in transformer model architecture.
+    This operator only supports running on GPU. The function of the operator is consistent with
+    the following pseudo code:
+
+    .. code-block:: python
+
+        residual = src;
+        if pre_layer_norm:
+            src = layer_norm(src)
+        src = linear(dropout(activation(dropout(linear(src)))))
+        if not pre_layer_norm:
+            src = layer_norm(out)
+
+    Args:
+        x (Tensor): the input tensor could be 3-D tensor, the input data type could be float16, float32 or float64, the shape is`[batch\_size, sequence\_length, d_model]`.
+        linear1_weight (Tensor): The weight of first linear, the data type is same as `x`, the shape is `[d\_model, dim\_feedforward]`.
+        linear2_weight (Tensor): The weight of second linear, the data type is same as `x`, the shape is `[dim\_feedforward, d\_model]`.
+        linear1_bias (Tensor, optional): The bias of first linear, the data type is same as `x`, the shape is `[dim_feedforward]`. Default None.
+        linear2_bias (Tensor, optional): The bias of second linear, the data type is same as `x`, the shape is `[d_model]`. Default None.
+        ln1_scale (Tensor, optional): the weight of first layer_norm, the data type is float32 or float64, the shape is same as `x`. Default None.
+        ln1_bias (Tensor, optional): The bias of first layer_norm, the data type is float32 or float64, the shape is `[d\_model]`. Default None.
+        ln2_scale (Tensor, optional): The weight of second layer_norm, the data type is float32 or float64, the shape is same as `x`. Default None.
+        ln2_bias (Tensor, optional): The bias of second layer_norm, the data type is float32 or float64, the shape is `[d\_model]`. Default None.
+        dropout1_rate (float, optional): The first dropout probability of setting units to zero. Default 0.5.
+        dropout2_rate (float, optional): The second dropout probability of setting units to zero. Default 0.5.
+        activation (str, optional): The activation. Default "relu".
+        ln1_epsilon (float, optional): Small float of first layer_norm added to denominator to avoid dividing by zero. Default is 1e-5.
+        ln2_epsilon (float, optional): Small float of second layer_norm added to denominator to avoid dividing by zero. Default is 1e-5.
+        pre_layer_norm (bool, optional): add layer_norm in the pre-processing stage or post-processing state.
+        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        Tensor: The output Tensor, the data type and shape is same as `x`.
+
+    Examples:
+        .. code-block:: python
+
+            # required: gpu
+            import paddle
+            import numpy as np
+            x_data = np.random.random((1, 8, 8)).astype("float32")
+            linear1_weight_data = np.random.random((8, 8)).astype("float32")
+            linear2_weight_data = np.random.random((8, 8)).astype("float32")
+            x = paddle.to_tensor(x_data)
+            linear1_weight = paddle.to_tensor(linear1_weight_data)
+            linear2_weight = paddle.to_tensor(linear2_weight_data)
+            out = paddle.nn.functional.fused_feedforward(x, linear1_weight, linear2_weight)
+            print(out.numpy().shape)
+            # (1, 8, 8)
+    """
+    _verify_dropout_rate(dropout1_rate)
+    _verify_dropout_rate(dropout2_rate)
+
+    if in_dygraph_mode():
+        out, _, _, _, _, _, _, _, _, _, _ = _C_ops.fused_feedforward(
+            x, None, None, linear1_weight, linear1_bias, linear2_weight,
+            linear2_bias, ln1_scale, ln1_bias, ln2_scale, ln2_bias,
+            'pre_layer_norm', pre_layer_norm, 'ln1_epsilon', ln1_epsilon,
+            'ln2_epsilon', ln2_epsilon, 'act_method', activation,
+            'dropout1_rate', dropout1_rate, 'dropout2_rate', dropout2_rate)
+        return out
+
+    helper = LayerHelper("fused_feedforward")
+    dtype = x.dtype
+    check_variable_and_dtype(x, 'x', ['float16', 'float32', 'float64'],
+                             'fused_feedforward')
+    check_dtype(dtype, 'dtype', ['float16', 'float32', 'float64'],
+                'fused_feedforward')
+
+    out = helper.create_variable_for_type_inference(x.dtype)
+    dropout1_mask = helper.create_variable_for_type_inference(
+        'uint8', stop_gradient=True)
+    dropout2_mask = helper.create_variable_for_type_inference(
+        'uint8', stop_gradient=True)
+    ln1_mean = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+    ln1_variance = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+    ln2_mean = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+    ln2_variance = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+    linear1_out = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+    ln1_out = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+    dropout1_out = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+    dropout2_out = helper.create_variable_for_type_inference(
+        x.dtype, stop_gradient=True)
+
+    helper.append_op(
+        type='fused_feedforward',
+        inputs={
+            'X': x,
+            'Linear1Weight': linear1_weight,
+            'Linear1Bias': linear1_bias,
+            'Linear2Weight': linear2_weight,
+            'Linear2Bias': linear2_bias,
+            'Ln1Scale': ln1_scale,
+            'Ln1Bias': ln1_bias,
+            'Ln2Scale': ln2_scale,
+            'Ln2Bias': ln2_bias,
+        },
+        outputs={
+            'Out': out,
+            'Dropout1Mask': dropout1_mask,
+            'Dropout2Mask': dropout2_mask,
+            'Ln1Mean': ln1_mean,
+            'Ln1Variance': ln1_variance,
+            'Ln2Mean': ln2_mean,
+            'Ln2Variance': ln2_variance,
+            'Linear1Out': linear1_out,
+            'Ln1Out': ln1_out,
+            'Dropout1Out': dropout1_out,
+            'Dropout2Out': dropout2_out,
+        },
+        attrs={
+            'dropout1_rate': dropout1_rate,
+            'dropout2_rate': dropout2_rate,
+            'act_method': activation,
+            'pre_layer_norm': pre_layer_norm,
+            'ln1_epsilon': ln1_epsilon,
+            'ln2_epsilon': ln2_epsilon,
+        })
+    return out
+
+
 def fused_multi_head_attention(x,
                                qkv_weight,
                                linear_weight,