Add ability for force bos id for mbart

examples/cpp/bart/config.ini

@@ -17,7 +17,7 @@ repetition_penalty=1.0 ; Use for sampling
 presence_penalty=0.0  ; Only one of repetition_penalty and presence_penalty are allowed.
 len_penalty=0.0
 beam_search_diversity_rate=0.0
-request_batch_size=8 # determine by the request
+request_batch_size=1 # determine by the request
 request_output_len=32 # determine by the request
 
 [encoder]

src/fastertransformer/kernels/decoding_kernels.cu

@@ -64,6 +64,34 @@ void invokeDecodingInitialize(bool*        finished,
         finished, sequence_length, word_ids, cum_log_probs, sentence_ids, batch_size, beam_width, max_input_length);
 }
 
+__global__ void forceId(int*       word_ids,
+                        const int* force_bos_ids,
+                        const int  batch_size,
+                        const int  beam_width,
+                        const int  step)
+{
+    for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < batch_size * beam_width;
+         index += blockDim.x * gridDim.x) {
+        if (word_ids != nullptr) {
+            word_ids[index+step*batch_size*beam_width] = force_bos_ids[index / beam_width];
+        }
+    }
+}
+
+void invokeForceId(int*         word_ids,
+                   const int*   force_bos_ids,
+                   const int    batch_size,
+                   const int    beam_width,
+                   const int    step,
+                   cudaStream_t stream)
+{
+    dim3 grid((int)ceil(batch_size * beam_width * 1.0 / 256));
+    dim3 block(256);
+
+    forceId<<<grid, block, 0, stream>>>(
+        word_ids, force_bos_ids, batch_size, beam_width, step);
+}
+
 template void invokeDecodingInitialize(bool*        finished,
                                        int*         sequence_length,
                                        int*         word_ids,

src/fastertransformer/kernels/decoding_kernels.h

@@ -33,6 +33,13 @@ void invokeDecodingInitialize(bool*        finished,
                               const int    max_input_length,
                               cudaStream_t stream);
 
+void invokeForceId(int*         word_ids,
+                   const int*   force_bos_ids,
+                   const int    batch_size,
+                   const int    beam_width,
+                   const int    step,
+                   cudaStream_t stream);
+
 // get token from all_ids at step, then lookup from the embedding table
 // by the token
 template<typename T>

src/fastertransformer/models/bart/BartDecoder.cc

@@ -546,35 +546,33 @@ void BartDecoder<T>::forward(std::vector<Tensor>*                           outp
                                                        stream_);
         }
         sync_check_cuda_error();
-// {
-//     {
-//                     T* buf;
-//                     int st = local_batch_size * d_model_;
-//                     buf = new T[st];
-//                     cudaMemcpy(buf, decoder_output, sizeof(T) * st, cudaMemcpyDeviceToHost);
-//                     auto step_ptr = input_tensors->at(4).data;
-//                     int step = ((int*)step_ptr)[0];
-//                     if (step == 1) {
-//                         printf("decoder_output at layer %d step %d\n", l, step);
-//                         for (int i=0; i<50; i++) {
-//                             printf("%f ", double(buf[i]));
-//                         }
-//                         printf("buf last: %f\n", double(buf[st-1]));
-//                         printf("\n");
-//                     }
-//     }
-// }
+
         if (isLastLayerParallelId(l) == true && pipeline_para_.rank_ != pipeline_para_.world_size_ - 1
             && pipeline_para_.world_size_ > 1) {
             // ftNcclSend(decoder_output, local_batch_size * d_model_, pipeline_para_.rank_ + 1,
             // pipeline_para_, stream_);
-
             ftNcclSend(decoder_output + local_batch_size * d_model_ / tensor_para_.world_size_ * tensor_para_.rank_,
                        local_batch_size * d_model_ / tensor_para_.world_size_,
                        pipeline_para_.rank_ + 1,
                        pipeline_para_,
                        stream_);
         }
+        // {
+        //             T* buf;
+        //             int st = local_batch_size * d_model_;
+        //             buf = new T[st];
+        //             cudaMemcpy(buf, decoder_output, sizeof(T) * st, cudaMemcpyDeviceToHost);
+        //             auto step_ptr = input_tensors->at(4).data;
+        //             int step = ((int*)step_ptr)[0];
+        //             if (step == 1) {
+        //                 printf("decoder_output at layer %d step %d\n", l, step);
+        //                 for (int i=0; i<50; i++) {
+        //                     printf("%f ", double(buf[i]));
+        //                 }
+        //                 printf("buf last: %f\n", double(buf[st-1]));
+        //                 printf("\n");
+        //             }
+        // }
     }
 
     if (is_free_buffer_after_forward_ == true) {

src/fastertransformer/models/bart/BartDecoding.cc

@@ -116,6 +116,7 @@ void BartDecoding<T>::allocateBuffer(
 
     start_ids_buf_ = (int*)(allocator_->reMalloc(start_ids_buf_, sizeof(int) * batch_size, false));
     end_ids_buf_   = (int*)(allocator_->reMalloc(end_ids_buf_, sizeof(int) * batch_size, false));
+    forced_bos_ids_buf_ = (int*)(allocator_->reMalloc(forced_bos_ids_buf_, sizeof(int) * batch_size, false));
 
     output_ids_buf_ =
         (int*)(allocator_->reMalloc(output_ids_buf_, sizeof(int) * batchxbeam * (max_seq_len + 1), false));
@@ -182,6 +183,7 @@ void BartDecoding<T>::freeBuffer()
         allocator_->free((void**)(&tiled_encoder_sequence_length_));
 
         allocator_->free((void**)(&start_ids_buf_));
+        allocator_->free((void**)(&forced_bos_ids_buf_));
         allocator_->free((void**)(&end_ids_buf_));
 
         allocator_->free((void**)(&output_ids_buf_));
@@ -343,6 +345,7 @@ void BartDecoding<T>::forward(TensorMap*                   output_tensors,
     //      stop_words_list [batch_size, 2, stop_words_length], optional
     //      bad_words_list [batch_size, 2, stop_words_length], optional
     //      start_id [batch_size] on cpu, optional
+    //      forced_bos_id [batch_size] on cpu, optional
     //      end_id [batch_size] on cpu, optional
     //      runtime_top_k [1] or [batch_size] on cpu, optional, uint.
     //      runtime_top_p [1] or [batch_size] on cpu, optional, float.
@@ -382,6 +385,7 @@ void BartDecoding<T>::forward(TensorMap*                   output_tensors,
         dynamic_decode_layer_->setup(batch_size, beam_width, &input_map);
         handleOptArg(&input_map, "start_id", start_ids_buf_, start_id_, batch_size);
         handleOptArg(&input_map, "end_id", end_ids_buf_, end_id_, batch_size);
+        handleOptArg(&input_map, "forced_bos_id", forced_bos_ids_buf_, -1, batch_size);
     }
 
     FT_CHECK_WITH_INFO(input_tensors->at("encoder_output").shape[2] == d_model_,
@@ -734,18 +738,6 @@ void BartDecoding<T>::forward(TensorMap*                   output_tensors,
                      {"local_batch_size", Tensor{MEMORY_CPU, TYPE_INT32, {1}, &tmp_local_batch_size}},
                      {"is_initialize_random_table", Tensor{MEMORY_CPU, TYPE_BOOL, {1}, &is_initialize_random_table}}});
 
-// {
-//                     T* buf;
-//                     int st = batch_size * beam_width * vocab_size_padded_;
-//                     buf = new T[st];
-//                     cudaMemcpy(buf, logits_buf_, sizeof(T) * st, cudaMemcpyDeviceToHost);
-//                     printf("logits_buf_\n");
-//                     for (int i=0; i<50; i++) {
-//                         printf("%f ", double(buf[i]));
-//                     }
-//                     printf("buf last: %f\n", double(buf[st-1]));
-//                     printf("\n");
-// }
                 if (cache_indirections_[src_indir_idx] != nullptr) {
                     dynamic_decode_input_tensors.insert(
                         "src_cache_indirection",
@@ -803,19 +795,33 @@ void BartDecoding<T>::forward(TensorMap*                   output_tensors,
                     }
                     dynamic_decode_output_tensors.insert(*t);
                 }
-    // {
-    //     int* buf;
-    //     int st = batch_size * (max_seq_len+1);
-    //     buf = new int[st];
-    //     cudaMemcpy(buf, output_ids_buf_, sizeof(int) * st, cudaMemcpyDeviceToHost);
-    //     printf("start_ids_buf_ before forward: %d\n", batch_size);
-    //     for (int i=0; i<st; i++) {
-    //         printf("%d ", buf[i]);
-    //     }
-    //     printf("\n");
-    // }
-
                 dynamic_decode_layer_->forward(&dynamic_decode_output_tensors, &dynamic_decode_input_tensors);
+                if (step == 1 && input_tensors->isExist("forced_bos_id")) {
+                    invokeForceId(output_ids_buf_,
+                                  forced_bos_ids_buf_,
+                                  batch_size,
+                                  beam_width,
+                                  step,
+                                  stream_);
+                    sync_check_cuda_error();
+                }
+                // {
+                // for (auto t = dynamic_decode_output_tensors.begin(); t != dynamic_decode_output_tensors.end(); ++t) {
+                //         printf("step: %d, t->first: %s\n", step, t->first.c_str());
+                //         // printf("%s\n", t->second.toString().c_str());
+                //         {
+                //                     int* buf;
+                //                     int st = t->second.size();
+                //                     buf = new int[st];
+                //                     cudaMemcpy(buf, t->second.data, sizeof(int) * t->second.size(), cudaMemcpyDeviceToHost);
+                //                     for (int i=0; i<st; i++) {
+                //                         printf("%d ", buf[i]);
+                //                     }
+                //                     printf("\n");
+                //         }
+                //     }
+                //     printf("\n\n");
+                // }
             }
         }
 
@@ -944,19 +950,6 @@ void BartDecoding<T>::forward(TensorMap*                   output_tensors,
         }
     }
 
-    // {
-    //     int* buf;
-    //     int st = batch_size * (max_seq_len+1);
-    //     buf = new int[st];
-    //     cudaMemcpy(buf, output_ids_buf_, sizeof(int) * st, cudaMemcpyDeviceToHost);
-    //     printf("output_ids_buf_ after finalize: %d\n", batch_size);
-    //     for (int i=0; i<st; i++) {
-    //         printf("%d ", buf[i]);
-    //     }
-    //     printf("\n");
-
-    // }
-
     if (pipeline_para_.world_size_ > 1) {
         ftNcclGroupStart();
         if (pipeline_para_.rank_ == pipeline_para_.world_size_ - 1) {
@@ -1023,17 +1016,6 @@ void BartDecoding<T>::forward(TensorMap*                   output_tensors,
     // throw errors when detected
     ftNcclStreamSynchronize(tensor_para_, pipeline_para_, stream_);
 
-    // {
-    //     int* buf;
-    //     int st = 32;
-    //     buf = new int[st];
-    //     cudaMemcpy(buf, output_tensors->at("output_ids").data, sizeof(int) * st, cudaMemcpyDeviceToHost);
-    //     printf("output_ids after finalize: %s %d\n", output_tensors->at("output_ids").toString().c_str(), batch_size);
-    //     for (int i=0; i<st; i++) {
-    //         printf("%d ", buf[i]);
-    //     }
-    //     printf("\n");
-    // }
     if (is_free_buffer_after_forward_) {
         freeBuffer();
     }

src/fastertransformer/models/bart/BartDecoding.h

@@ -94,8 +94,9 @@ class BartDecoding: public BaseLayer {
     bool*              finished_buf_              = nullptr;
     bool*              h_finished_buf_            = nullptr;
 
-    int* start_ids_buf_ = nullptr;
-    int* end_ids_buf_   = nullptr;
+    int* start_ids_buf_         = nullptr;
+    int* forced_bos_ids_buf_    = nullptr;
+    int* end_ids_buf_           = nullptr;
 
     T*   key_cache_             = nullptr;
     T*   value_cache_           = nullptr;

src/fastertransformer/models/bart/BartDecodingWeight.cc

@@ -265,11 +265,28 @@ void BartDecodingWeight<T>::loadModel(std::string dir_path)
     loadWeightFromBin<T>(
         weights_ptr[3], {(size_t)weights_size[3]}, dir_path + "/decoder.final_layer_norm.weight.bin", model_file_type);
     if (bart_with_bias) {
-        loadWeightFromBin<T>(weights_ptr[4],
-                             {(size_t)weights_size[4]},
-                             dir_path + "/decoder.final_layer_norm.bias.bin",
-                             model_file_type);
-        loadWeightFromBin<T>(weights_ptr[5], {(size_t)weights_size[5]}, dir_path + "/decoder.final_logits_bias.bin", model_file_type);
+        if (mbart) {
+            loadWeightFromBin<T>(weights_ptr[4],
+                                {(size_t)weights_size[4]},
+                                dir_path + "/decoder.layer_norm.weight.bin",
+                                model_file_type);
+            loadWeightFromBin<T>(weights_ptr[5],
+                                {(size_t)weights_size[5]},
+                                dir_path + "/decoder.final_layer_norm.bias.bin",
+                                model_file_type);
+            loadWeightFromBin<T>(weights_ptr[6],
+                                {(size_t)weights_size[6]},
+                                dir_path + "/decoder.layer_norm.bias.bin",
+                                model_file_type);
+            loadWeightFromBin<T>(weights_ptr[7], {(size_t)weights_size[7]}, dir_path + "/decoder.final_logits_bias.bin", model_file_type);
+
+        } else {
+            loadWeightFromBin<T>(weights_ptr[4],
+                                {(size_t)weights_size[4]},
+                                dir_path + "/decoder.final_layer_norm.bias.bin",
+                                model_file_type);
+            loadWeightFromBin<T>(weights_ptr[5], {(size_t)weights_size[5]}, dir_path + "/decoder.final_logits_bias.bin", model_file_type);
+        }
     }
 
     for (int l = 0; l < num_layer_; l++) {

src/fastertransformer/models/bart/BartEncoder.cc

@@ -455,15 +455,15 @@ void BartEncoder<T>::forward(TensorMap*                  output_tensors,
 // {
 //         T* buf;
 //         int batch_size = 1;
-//         int seq_len = 11;
+//         int seq_len = 13;
 //         int st = batch_size * seq_len * d_model_;
 //         printf("st: %d %d %d %d\n",batch_size, seq_len, d_model_, st);
 //         buf = new T[st];
 //         cudaMemcpy(buf, bart_encoder_emb_buf_, sizeof(T) * st, cudaMemcpyDeviceToHost);
 //         printf("bart_encoder_emb_buf_\n");
 //         for (int i=0; i < seq_len; i++) {
 //             for (int j=0; j<d_model_; j++) {
-//                 printf("%f ", double(buf[i+j*seq_len]));
+//                 printf("%f ", double(buf[i * d_model_+ j]));
 //                 if (j > 10) {
 //                     break;
 //                 }

src/fastertransformer/models/bart/BartEncoderWeight.cc

@@ -257,10 +257,25 @@ void BartEncoderWeight<T>::loadModel(std::string dir_path)
     loadWeightFromBin<T>(
         weights_ptr[2], {(size_t)weights_size[2]}, dir_path + "/encoder.final_layer_norm.weight.bin", model_file_type);
     if (bart_with_bias) {
-        loadWeightFromBin<T>(weights_ptr[3],
-                             {(size_t)weights_size[3]},
-                             dir_path + "/encoder.final_layer_norm.bias.bin",
-                             model_file_type);
+        if (mbart) {
+            loadWeightFromBin<T>(weights_ptr[3],
+                                {(size_t)weights_size[3]},
+                                dir_path + "/encoder.final_layer_norm.bias.bin",
+                                model_file_type);
+            loadWeightFromBin<T>(weights_ptr[4],
+                                {(size_t)weights_size[4]},
+                                dir_path + "/encoder.layer_norm.weight.bin",
+                                model_file_type);
+            loadWeightFromBin<T>(weights_ptr[5],
+                                {(size_t)weights_size[5]},
+                                dir_path + "/encoder.layer_norm.bias.bin",
+                                model_file_type);    
+        } else {
+            loadWeightFromBin<T>(weights_ptr[3],
+                                {(size_t)weights_size[3]},
+                                dir_path + "/encoder.final_layer_norm.bias.bin",
+                                model_file_type);
+        }
     }
 
     for (int l = 0; l < num_layer_; l++) {

src/fastertransformer/triton_backend/bart/BartTritonModel.cc

@@ -23,14 +23,15 @@ namespace ft = fastertransformer;
 
 std::shared_ptr<AbstractTransformerModel> AbstractTransformerModel::createBartModel(std::string model_dir)
 {
+    printf("createBartModel\n");
     INIReader reader = INIReader(model_dir + "/config.ini");
     if (reader.ParseError() < 0) {
         std::cout << "[ERROR] Can't load '" << model_dir << "/config.ini"
                   << "'\n";
         return nullptr;
     }
 
-    const std::string data_type = reader.Get("ft_instance_hyperparameter", "data_type");
+    const std::string data_type = "fp32"; // reader.Get("ft_instance_hyperparameter", "data_type");
     if (data_type == "fp16") {
         // return std::make_shared<BartTritonModel<half>>(reader, model_dir);
         return std::make_shared<BartTritonModel<half>>(1, 1, 0, model_dir, 0);
@@ -60,6 +61,12 @@ BartTritonModel<T>::BartTritonModel(INIReader reader, std::string model_dir): mo
     encoder_num_layer_     = reader.GetInteger("encoder", "num_layers");
     encoder_vocab_size_    = reader.GetInteger("encoder", "vocab_size");
     encoder_max_pos_seq_len_ = reader.GetInteger("encoder", "max_pos_seq_len");
+    mbart_para_              = reader.GetBoolean("encoder", "mbart", false);
+    if (mbart_para_) {
+        layernorm_type_ = ft::LayerNormType::pre_layernorm;
+    } else {
+        layernorm_type_ = ft::LayerNormType::post_layernorm;
+    }
 
     // decoding
     decoding_head_num_      = reader.GetInteger("decoder", "num_heads");
@@ -111,7 +118,13 @@ BartTritonModel<T>::BartTritonModel(size_t      tensor_para_size,
     encoder_vocab_size_    = reader.GetInteger("encoder", "vocab_size");
     encoder_max_pos_seq_len_ =
         reader.GetInteger("encoder", "max_pos_seq_len");
-
+    mbart_para_              = reader.GetBoolean("encoder", "mbart", false);
+    if (mbart_para_) {
+        layernorm_type_ = ft::LayerNormType::pre_layernorm;
+    } else {
+        layernorm_type_ = ft::LayerNormType::post_layernorm;
+    }
+
     // decoding
     decoding_head_num_      = reader.GetInteger("decoder", "num_heads");
     decoding_size_per_head_ = reader.GetInteger("decoder", "d_kv");

src/fastertransformer/triton_backend/bart/BartTritonModel.h

@@ -91,7 +91,6 @@ struct BartTritonModel: public AbstractTransformerModel {
 
     // bart structure difference
     bool                      bart_with_bias_ = true;
-    // TODO(zhwang): support mbart.
     bool                      mbart_para_ = false;
     bool                      use_gated_activation_ = false;
     ft::PositionEmbeddingType position_embedding_type_ = ft::PositionEmbeddingType::absolute;