[LLM] support bloom dybatch (#7003)

* support_llama

* support_llama

* delete_loss

Author: xiaoxiaohehe001

csrc/generation/set_alibi_mask_value.cu

@@ -0,0 +1,136 @@
+// Copyright (c) 2023 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 "helper.h"
+
+template <typename T>
+__global__ void set_value_by_id(const int *seq_lens, 
+                               const bool *stop_flags, 
+                              const float *alibi_slopes, 
+                              const int64_t *tgt_pos, 
+                              T *output_data, 
+                              int *sequence_lengths, 
+                              int bs, 
+                              int length,
+                              int num_head) {
+    int bs_id = blockIdx.x;                          
+    int hid = threadIdx.x;
+    if (bs_id < bs) {
+        T *output_data_now = output_data + bs_id * num_head * length + hid * length;
+        float tgt_pos_now = static_cast<float>(tgt_pos[bs_id]);
+        output_data_now[seq_lens[bs_id]] = static_cast<T>(tgt_pos_now * alibi_slopes[hid]);
+        if (stop_flags[bs_id]) {
+            sequence_lengths[bs_id] = 0;
+        }
+    }
+}
+
+template <paddle::DataType D>
+std::vector<paddle::Tensor> set_mask_value(const paddle::Tensor& input_data, 
+                                           const paddle::Tensor& stop_flags, 
+                                          const paddle::Tensor& seq_lens,
+                                          const paddle::Tensor& alibi_slopes,
+                                          const paddle::Tensor& tgt_pos
+                                          ) {
+    typedef PDTraits<D> traits_;
+    typedef typename traits_::DataType DataType_;
+    typedef typename traits_::data_t data_t;
+
+    PD_CHECK(seq_lens.dtype() == paddle::DataType::INT32);
+    PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
+    auto cu_stream = input_data.stream();
+    std::vector<int64_t> input_data_shape = input_data.shape();
+    std::vector<int64_t> seq_lens_shape = seq_lens.shape();
+    auto sequence_lengths = seq_lens.copy_to(seq_lens.place(), false);
+
+    int input_bs = input_data_shape[0];
+    int length = input_data_shape[3];
+    int seq_bs = seq_lens_shape[0];
+    int num_head = alibi_slopes.shape()[0];
+
+    int grid_size = input_bs;
+    int block_size = num_head;
+    set_value_by_id<<<grid_size, block_size, 0, cu_stream>>>(seq_lens.data<int>(), 
+                                                     stop_flags.data<bool>(), 
+                                                     alibi_slopes.data<float>(),
+                                                     tgt_pos.data<int64_t>(),
+                                                     reinterpret_cast<DataType_*>(const_cast<data_t*>(input_data.data<data_t>())), 
+                                                     sequence_lengths.data<int>(), seq_bs, length, num_head);
+    return {sequence_lengths};
+}
+
+std::vector<paddle::Tensor> SetMaskValue(const paddle::Tensor& input_data, 
+                                          const paddle::Tensor& stop_flags, 
+                                          const paddle::Tensor& seq_lens,
+                                          const paddle::Tensor& alibi_slopes,
+                                          const paddle::Tensor& tgt_pos) {
+    switch (input_data.type()) {
+        case paddle::DataType::BFLOAT16: {
+            return set_mask_value<paddle::DataType::BFLOAT16>(
+                input_data,
+                stop_flags,
+                seq_lens,
+                alibi_slopes,
+                tgt_pos
+            );
+        }
+        case paddle::DataType::FLOAT16: {
+            return set_mask_value<paddle::DataType::FLOAT16>(
+                input_data,
+                stop_flags,
+                seq_lens,
+                alibi_slopes,
+                tgt_pos
+            );
+        }
+        case paddle::DataType::FLOAT32: {
+            return set_mask_value<paddle::DataType::FLOAT32>(
+                input_data,
+                stop_flags,
+                seq_lens,
+                alibi_slopes,
+                tgt_pos
+            );
+        }
+        default: {
+            PD_THROW(
+                "NOT supported data type. "
+                "Only float16, bfloat16 and float32 are supported. ");
+            break;
+        }
+    }
+}
+
+std::vector<std::vector<int64_t>> SetMaskValueInferShape(const std::vector<int64_t>& input_data_shape, 
+                                                         const std::vector<int64_t>& stop_flags_shape, 
+                                                         const std::vector<int64_t>& seq_lens_shape,
+                                                         const std::vector<int64_t>& alibi_slopes_shape,
+                                                         const std::vector<int64_t>& tgt_pos) {
+    return {seq_lens_shape};
+}
+
+std::vector<paddle::DataType> SetMaskValueInferDtype(const paddle::DataType& input_data_dtype, 
+                                                      const paddle::DataType& stop_flags_dtype, 
+                                                      const paddle::DataType& seq_lens_dtype,
+                                                      const paddle::DataType& alibi_slopes_dtype,
+                                                      const paddle::DataType& tgt_pos_dtype) {
+    return {seq_lens_dtype};
+}
+
+PD_BUILD_OP(set_alibi_mask_value)
+    .Inputs({"input_data", "stop_flags", "seq_lens", "alibi_slopes", "tgt_pos"})
+    .Outputs({"sequence_lengths"})
+    .SetKernelFn(PD_KERNEL(SetMaskValue))
+    .SetInferShapeFn(PD_INFER_SHAPE(SetMaskValueInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(SetMaskValueInferDtype));

csrc/setup_cuda.py

@@ -31,6 +31,7 @@
             "./generation/write_cache_kv.cu",
             "./generation/encode_rotary_qk.cu",
             "./generation/top_p_sampling.cu",
+            "./generation/set_alibi_mask_value.cu",
         ]
     ),
 )

llm/predictor.py

@@ -19,7 +19,6 @@
 import time
 from abc import abstractmethod
 from dataclasses import dataclass, field
-from distutils.command.config import config
 
 import numpy as np
 import paddle
@@ -139,6 +138,7 @@ def __init__(self, config: PredictorArgument, tokenizer: PretrainedTokenizer = N
         self.tokenizer = tokenizer
         self.return_tensors = "pd"
         self.tensor_parallel_rank, self.tensor_parallel_degree = init_dist_env()
+        self.model_config.tensor_parallel_rank, self.model_config.tensor_parallel_degree = init_dist_env()
 
     def _preprocess(self, source):
         tokenized_source = self.tokenizer(
@@ -284,11 +284,11 @@ def __init__(self, config: PredictorArgument, tokenizer: PretrainedTokenizer):
             self.cache_kvs[0].shape[-3],
             self.cache_kvs[0].shape[-1],
         )
-        total_max_length = config.src_length + config.max_length
-        self.pre_ids = paddle.full([config.batch_size, total_max_length], -1, dtype="int64")
+        self.total_max_length = config.src_length + config.max_length
+        self.pre_ids = paddle.full([config.batch_size, self.total_max_length], -1, dtype="int64")
         if "chatglm" in self.architectures:
             self.attention_mask = paddle.ones(
-                shape=(config.batch_size, 1, total_max_length, total_max_length),
+                shape=(config.batch_size, 1, self.total_max_length, self.total_max_length),
                 dtype=self.dtype,
             )
             self.tgt_pos = paddle.ones(
@@ -297,15 +297,17 @@ def __init__(self, config: PredictorArgument, tokenizer: PretrainedTokenizer):
             )
         else:
             self.attention_mask = paddle.zeros(
-                shape=(config.batch_size, 1, total_max_length, total_max_length),
+                shape=(config.batch_size, 1, self.total_max_length, self.total_max_length),
                 dtype=self.dtype,
             )
 
         self.tgt_generation_mask = paddle.zeros(
-            shape=[config.batch_size, 1, 1, total_max_length],
+            shape=[config.batch_size, 1, 1, self.total_max_length],
             dtype=self.dtype,
         )
-        self.arange_tensor_encoder = paddle.zeros(shape=(config.batch_size, 1, total_max_length), dtype=self.dtype)
+        self.arange_tensor_encoder = paddle.zeros(
+            shape=(config.batch_size, 1, self.total_max_length), dtype=self.dtype
+        )
 
         if config.export_precache:
             if config.prefix_path:
@@ -342,6 +344,10 @@ def _postprocess(self, predictions):
             return None
 
     def _preprocess(self, source):
+        self.attention_mask[:] = 0
+        self.tgt_generation_mask[:] = 0
+        pre_caches_length = 0 if not self.config.export_precache else self.pre_caches[0].shape[-2]
+
         if "chatglm" in self.architectures:
             inputs = dybatch_preprocess(
                 self.tokenizer,
@@ -370,12 +376,12 @@ def _preprocess(self, source):
             )
             for i in range(inputs["input_ids"].shape[0]):
                 length = inputs["seq_len_encoder"][i][0]
-                self.attention_mask[i, 0, :length, :length] = paddle.tril(
+                self.attention_mask[i, :, :length, :length] = paddle.tril(
                     paddle.ones(shape=(length, length), dtype=self.config.dtype)
                 )
-                self.arange_tensor_encoder[i, 0, :length] = paddle.arange(length).astype(self.config.dtype)
+                self.arange_tensor_encoder[i, :, :length] = paddle.arange(length).astype(self.config.dtype)
 
-                self.tgt_generation_mask[i, 0, 0, :length] = paddle.ones(shape=[1, length], dtype=self.config.dtype)
+                self.tgt_generation_mask[i, :, 0, :length] = paddle.ones(shape=[1, length], dtype=self.config.dtype)
             # alibi encoder
             alibi_slopes = get_alibi_slopes(self.model_config.n_head)
             inputs["position_ids"] = paddle.to_tensor(alibi_slopes, dtype="float32")
@@ -402,16 +408,16 @@ def _preprocess(self, source):
                 [
                     inputs["input_ids"].shape[0],
                     self.model_config.n_head // self.model_config.tensor_parallel_degree,
-                    self.config.max_length,
-                    self.config.max_length,
+                    self.total_max_length,
+                    self.total_max_length,
                 ]
             )
             alibi_decoder = alibi.expand(
                 [
                     inputs["input_ids"].shape[0],
                     self.model_config.n_head // self.model_config.tensor_parallel_degree,
                     1,
-                    self.config.max_length,
+                    self.total_max_length,
                 ]
             )
             self.attention_mask = (
@@ -422,7 +428,6 @@ def _preprocess(self, source):
             )
 
         else:
-            pre_caches_length = 0 if not self.config.export_precache else self.pre_caches[0].shape[-2]
             inputs = dybatch_preprocess(
                 self.tokenizer,
                 source,
@@ -655,7 +660,7 @@ def create_predictor(
                     from paddlenlp.experimental.transformers import (
                         LlamaForCausalLMInferenceModel as LlamaInferenceModel,
                     )
-                  
+
                     config.tensor_parallel_degree = tensor_parallel_degree
                     config.tensor_parallel_rank = tensor_parallel_rank
                     config.quant_bits = -1
@@ -679,6 +684,20 @@ def create_predictor(
                     dtype=predictor_args.dtype,
                 )
                 model.eval()
+            elif "bloom" in config.architectures[0].lower():
+                from paddlenlp.experimental.transformers import (
+                    BloomForCausalLMInferenceModel,
+                )
+
+                config.tensor_parallel_degree = tensor_parallel_degree
+                config.tensor_parallel_rank = tensor_parallel_rank
+                model = BloomForCausalLMInferenceModel.from_pretrained(
+                    predictor_args.model_name_or_path,
+                    config=config,
+                    dtype=predictor_args.dtype,
+                )
+                cache_kvs_shape = BloomForCausalLMInferenceModel.get_cache_kvs_shape(config, predictor_args.batch_size)
+                model.eval()
             predictor = DygraphInferencePredictor(predictor_args, model=model, tokenizer=tokenizer)
         elif predictor_args.mode == "static":
             config = AutoConfig.from_pretrained(predictor_args.model_name_or_path)
@@ -698,6 +717,15 @@ def create_predictor(
                     config, predictor_args.batch_size
                 )
                 predictor = StaticInferencePredictor(predictor_args, cache_kvs_shape, tokenizer=tokenizer)
+            elif "bloom" in config.architectures[0].lower():
+                from paddlenlp.experimental.transformers import (
+                    BloomForCausalLMInferenceModel,
+                )
+
+                cache_kvs_shape = BloomForCausalLMInferenceModel.get_cache_kvs_shape(config, predictor_args.batch_size)
+                predictor = StaticInferencePredictor(
+                    predictor_args, cache_kvs_shape=cache_kvs_shape, tokenizer=tokenizer
+                )
         else:
             raise ValueError("the `mode` should be one of [dynamic, static]")
     return predictor

paddlenlp/experimental/transformers/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .bloom import *
 from .chatglm import *
 from .fused_transformer_layers import *
 from .llama import *

paddlenlp/experimental/transformers/bloom/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023 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.
+
+from .modeling import *