add data copy (#475)

[FasterTransformer] Add data copy to implement ChooseKernel and DataTransform in Paddle

Author: FrostML

examples/machine_translation/transformer/faster_transformer/encoder_decoding_predict.py

@@ -87,6 +87,10 @@ def do_predict(args):
     place = paddle.set_device(place)
 
     # Define data loader
+    # NOTE: Data yielded by DataLoader may be on CUDAPinnedPlace,
+    # but custom op doesn't support CUDAPinnedPlace. Hence,
+    # disable using CUDAPinnedPlace in DataLoader.
+    paddle.fluid.reader.use_pinned_memory(False)
     test_loader, to_tokens = reader.create_infer_loader(args)
 
     # Define model

paddlenlp/ops/src/fusion_decoding_op.cc

@@ -80,13 +80,23 @@ std::vector<paddle::Tensor> DecodingForward(
   } else {
     PD_THROW("Not supported decoding strategy. ");
   }
-  auto output_ids = paddle::Tensor(input.place(), output_dims);
-  auto parent_ids = paddle::Tensor(input.place(), parent_ids_dims);
-  auto sequence_length = paddle::Tensor(input.place(), sequence_length_dims);
 
   if (input.place() == paddle::PlaceType::kGPU) {
+    auto output_ids = paddle::Tensor(paddle::PlaceType::kGPU, output_dims);
+    auto parent_ids = paddle::Tensor(paddle::PlaceType::kGPU, parent_ids_dims);
+    auto sequence_length =
+        paddle::Tensor(paddle::PlaceType::kGPU, sequence_length_dims);
+
+    paddle::Tensor seq_len = paddle::Tensor(paddle::PlaceType::kGPU);
+
+    if (mem_seq_len.place() != paddle::PlaceType::kGPU) {
+      seq_len = mem_seq_len.copy_to<int>(paddle::PlaceType::kGPU);
+    } else {
+      seq_len = mem_seq_len;
+    }
+
     return DecodingCUDAForward(input,
-                               mem_seq_len,
+                               seq_len,
                                word_embedding,
                                self_ln_weight,
                                self_ln_bias,

paddlenlp/ops/src/fusion_gpt_op.cc

@@ -44,8 +44,16 @@ std::vector<paddle::Tensor> GPT2Forward(
   std::vector<int64_t> output_dims({total_len, batch_size});
   auto output_ids = paddle::Tensor(input.place(), output_dims);
 
-  if (input.place() == paddle::PlaceType::kGPU) {
-    return GPT2CUDAForward(input,
+  if (word_embedding.place() == paddle::PlaceType::kGPU) {
+    paddle::Tensor input_ids = paddle::Tensor(paddle::PlaceType::kCPU);
+
+    if (input.place() != paddle::PlaceType::kCPU) {
+      input_ids = input.copy_to<int>(paddle::PlaceType::kCPU);
+    } else {
+      input_ids = input;
+    }
+
+    return GPT2CUDAForward(input_ids,
                            word_embedding,
                            self_ln_weight,
                            self_ln_bias,

paddlenlp/ops/src/fusion_gpt_op.cu

@@ -61,20 +61,16 @@ std::vector<paddle::Tensor> gpt2_kernel(
   DecodingInitParam<DataType_> decoding_params;
   decoding_params.cublas_handle = cublas_handle_;
 
-  decoding_params.output_ids = output_ids.mutable_data<int>(input.place());
+  decoding_params.output_ids = output_ids.mutable_data<int>(word_emb.place());
 
   typedef DecoderTransformerTraits<traits_::OpType> DecodingTraits_;
   decoding_params.stream = stream;
   fastertransformer::Allocator<AllocatorType::PD> allocator_(stream);
 
   DecodingGpt2<DecodingTraits_::OpType>* gpt2_decoding;
 
-  // input data is on gpu.
-  int* h_input_data = new int[batch_size_ * start_len];
-  cudaMemcpy(h_input_data,
-             input.data<int>(),
-             sizeof(int) * batch_size_ * start_len,
-             cudaMemcpyDeviceToHost);
+  // input data should be on CPU.
+  int* h_input_data = input.data<int>();
   gpt2_decoding = new DecodingGpt2<DecodingTraits_::OpType>(allocator_,
                                                             batch_size_,
                                                             max_len,
@@ -189,7 +185,7 @@ std::vector<paddle::Tensor> GPT2CUDAForward(
     const int& eos_id,
     const float& temperature,
     const bool& use_fp16 = false) {
-  auto stream = input.stream();
+  auto stream = word_embedding.stream();
   cublasHandle_t cublas_handle_;
   cublasCreate(&cublas_handle_);
   cublasSetStream(cublas_handle_, stream);