Support pre-normalization for FT Encoder fp32 (PaddlePaddle#974)

* add encoder op

* support fp16

* support pre-normalization for FT Encoder fp32

* update CMakelists.txt

* move code to cuda_kernel_h and cuda_kernel_cu

* uncomment FT sample

* add jiaqi code

* self-attn output diff

* finish pre-normalization v1

* implement generalized kernel

* add v1 encoder

* delete assert for generalize version kernel

* support post-normalization fp32 for FT Encoder

* delete unused CMake commands

Co-authored-by: LiuChiaChi <[email protected]>
Co-authored-by: LiuChiachi <[email protected]>
Co-authored-by: Jiaqi Liu <[email protected]>

Author: 3 people

paddlenlp/ops/CMakeLists.txt

@@ -178,7 +178,6 @@ file(TO_NATIVE_PATH ${OPS_SOURCE_DIR}/patches/FasterTransformer/bert_encoder_tra
 file(TO_NATIVE_PATH ${CMAKE_BINARY_DIR}/${THIRD_PATH}/source/${THIRD_PARTY_NAME}/fastertransformer/bert_encoder_transformer.h bert_encoder_transformer_h_dst)
 
 set(OPT_OPEN_ATTN_COMMAND sed -i -e "370,392d" -e "410,454d" -e "229d" ${open_attention_h_dst})
-#set(OPT_BERT_ENCODER_COMMAND sed -i -e "552,592d" -e "118a   bool is_gelu_=true;" ${bert_encoder_transformer_h_dst})
 
 # TODO(guosheng): `find` seems meeting errors missing argument to `-exec', fix it
 set(MUTE_COMMAND grep -rl "printf(\"\\[WARNING\\]" ${CMAKE_BINARY_DIR}/${THIRD_PATH}/source/${THIRD_PARTY_NAME}/ | xargs -i{} sed -i "s/printf(\"\\WWARNING\\W decoding[^)]\\{1,\\})/ /" {})
@@ -305,5 +304,3 @@ if(ON_INFER AND WITH_GPT AND WITH_SP)
 endif()
 
 add_subdirectory(faster_transformer)
-
-

paddlenlp/ops/faster_transformer/src/fusion_encoder_op.cc

@@ -46,7 +46,8 @@ std::vector<paddle::Tensor> EncoderForward(
     const int64_t& num_layer,
     const int64_t& layer_idx,
     const bool& allow_gemm_test,
-    const bool& use_trt_kernel) {
+    const bool& use_trt_kernel,
+    const bool& normalize_before) {
   if (input.place() == paddle::PlaceType::kGPU) {
     auto shape = input.shape();
     auto encoder_out = paddle::Tensor(paddle::PlaceType::kGPU, shape);
@@ -80,7 +81,8 @@ std::vector<paddle::Tensor> EncoderForward(
                               num_layer,
                               layer_idx,
                               allow_gemm_test,
-                              use_trt_kernel);
+                              use_trt_kernel,
+                              normalize_before);
   } else {
     PD_THROW("Not implemented place. Only GPU is supported. ");
   }
@@ -116,7 +118,8 @@ std::vector<std::vector<int64_t>> EncoderInferShape(
     const int64_t& num_layer,
     const int64_t& layer_idx,
     const bool& allow_gemm_test,
-    const bool& use_trt_kernel) {
+    const bool& use_trt_kernel,
+    const bool& normalize_before) {
   return {input_shape};
 }
 
@@ -179,7 +182,8 @@ PD_BUILD_OP(fusion_encoder)
             "num_layer: int64_t",
             "layer_idx: int64_t",
             "allow_gemm_test: bool",
-            "use_trt_kernel: bool"})
+            "use_trt_kernel: bool",
+            "normalize_before: bool"})
     .SetKernelFn(PD_KERNEL(EncoderForward))
     .SetInferShapeFn(PD_INFER_SHAPE(EncoderInferShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(EncoderInferDtype));

paddlenlp/ops/faster_transformer/src/fusion_encoder_op.cu

@@ -61,6 +61,7 @@ std::vector<paddle::Tensor> encoder_kernel(
     int64_t layer_idx_,
     bool allow_gemm_test,
     bool use_trt_kernel_,
+    bool normalize_before,
     cublasHandle_t cublas_handle_,
     cudaStream_t stream) {
   int batch_size_ = input.shape()[0];
@@ -148,7 +149,8 @@ std::vector<paddle::Tensor> encoder_kernel(
   //   }
 
   BertEncoderTransformer<EncoderTraits_>* encoder =
-      new BertEncoderTransformer<EncoderTraits_>(int8_mode, allow_gemm_test);
+      new BertEncoderTransformer<EncoderTraits_>(
+          int8_mode, allow_gemm_test, normalize_before);
 
   encoder->allocateBuffer(allocator_,
                           batch_size_,
@@ -199,7 +201,8 @@ std::vector<paddle::Tensor> EncoderCUDAForward(
     int64_t num_layer,
     int64_t layer_idx,
     bool allow_gemm_test,
-    bool use_trt_kernel) {
+    bool use_trt_kernel,
+    bool normalize_before) {
   auto stream = input.stream();
   cublasHandle_t cublas_handle_;
   cublasCreate(&cublas_handle_);
@@ -241,6 +244,7 @@ std::vector<paddle::Tensor> EncoderCUDAForward(
           layer_idx,
           allow_gemm_test,
           use_trt_kernel,
+          normalize_before,
           cublas_handle_,
           stream);
 
@@ -279,6 +283,7 @@ std::vector<paddle::Tensor> EncoderCUDAForward(
           layer_idx,
           allow_gemm_test,
           use_trt_kernel,
+          normalize_before,
           cublas_handle_,
           stream);
       break;

paddlenlp/ops/faster_transformer/src/fusion_encoder_op.h

@@ -57,4 +57,5 @@ std::vector<paddle::Tensor> EncoderCUDAForward(
     int64_t num_layer_,
     int64_t layer_idx_,
     bool allow_gemm_test,
-    bool use_trt_kernel_);
+    bool use_trt_kernel_,
+    bool normalize_before);

paddlenlp/ops/faster_transformer/transformer/encoder.py

@@ -52,7 +52,8 @@ def infer_transformer_encoder(
         int8_mode=0,
         layer_idx=0,
         allow_gemm_test=False,
-        use_trt_kernel=False):
+        use_trt_kernel=False,
+        normalize_before=False):
     """
     Fusion Encoder API intergrating Encoder inference in FasterTransformer. It
     accepts the weight and bias of TransformerEncoder and some other parameters
@@ -92,6 +93,7 @@ def infer_transformer_encoder(
         'layer_idx': layer_idx,
         'allow_gemm_test': allow_gemm_test,
         'use_trt_kernel': use_trt_kernel,
+        'normalize_before': normalize_before
     }
     encoder_out = helper.create_variable(dtype=input.dtype)
     outputs = {"EncoderOut": encoder_out}
@@ -173,7 +175,8 @@ def encoder_layer_forward(self,
         # amax_list=paddle.to_tensor([]),  # int8 mode is not supported.
         n_head=self._config['nhead'],
         size_per_head=self._config['d_model'] // self._config['nhead'],
-        is_gelu=self._config['activation'] == 'gelu')
+        is_gelu=self._config['activation'] == 'gelu',
+        normalize_before=self._config['normalize_before'] == True)
     return src
 
 

paddlenlp/ops/patches/FasterTransformer/CMakeLists.txt

@@ -223,4 +223,4 @@ link_directories(
 
 add_subdirectory(fastertransformer)
 add_subdirectory(tools)
-add_subdirectory(sample)
+#add_subdirectory(sample)