[LLM] Fix error when config.bs is more than input bs.

llm/predictor.py

@@ -385,19 +385,25 @@ 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]
+        inputs = dybatch_preprocess(
+            self.tokenizer,
+            source,
+            self.config.src_length,
+            self.config.max_length,
+            self.architectures,
+            top_p=self.config.top_p,
+            temperature=self.config.temperature,
+            benchmark=self.config.benchmark,
+            pre_caches_length=pre_caches_length,
+        )
+
+        bsz = inputs["input_ids"].shape[0]
+        self.attention_mask = self.attention_mask[:bsz]
+        self.tgt_generation_mask = self.tgt_generation_mask[:bsz]
 
         if "chatglm" in self.architectures:
-            inputs = dybatch_preprocess(
-                self.tokenizer,
-                source,
-                self.config.src_length,
-                self.config.max_length,
-                self.architectures,
-                top_p=self.config.top_p,
-                temperature=self.config.temperature,
-                benchmark=self.config.benchmark,
-                pre_caches_length=pre_caches_length,
-            )
+            if bsz < self.config.batch_size:
+                self.tgt_pos = self.tgt_pos[:bsz]
             for i in range(inputs["input_ids"].shape[0]):
                 length = inputs["seq_len_encoder"][i][0]
                 self.attention_mask[i, 0, :length, :length] = 1
@@ -427,16 +433,8 @@ def _preprocess(self, source):
 
             inputs["tgt_pos"] = self.tgt_pos
         elif "bloom" in self.architectures:
-            inputs = dybatch_preprocess(
-                self.tokenizer,
-                source,
-                self.config.src_length,
-                self.config.max_length,
-                self.architectures,
-                top_p=self.config.top_p,
-                temperature=self.config.temperature,
-                benchmark=self.config.benchmark,
-            )
+            if bsz < self.config.batch_size:
+                self.arange_tensor_encoder = self.arange_tensor_encoder[:bsz]
             for i in range(inputs["input_ids"].shape[0]):
                 length = inputs["seq_len_encoder"][i][0]
                 self.attention_mask[i, :, :length, :length] = paddle.tril(
@@ -491,18 +489,6 @@ def _preprocess(self, source):
             )
 
         else:
-            inputs = dybatch_preprocess(
-                self.tokenizer,
-                source,
-                self.config.src_length,
-                self.config.max_length,
-                self.architectures,
-                top_p=self.config.top_p,
-                temperature=self.config.temperature,
-                pre_caches_length=pre_caches_length,
-                benchmark=self.config.benchmark,
-            )
-
             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(
@@ -618,7 +604,6 @@ def _infer(self, inputs):
         for i in range(len(self.cache_kvs_shape)):
             input_tensor = self.predictor.get_input_handle("cache_kvs_" + str(i))
             input_tensor.share_external_data(self.cache_kvs[i])
-
         input_tensor = self.predictor.get_input_handle("pre_ids")
         input_tensor.share_external_data(self.pre_ids)
 

paddlenlp/experimental/transformers/chatglm/modeling.py

@@ -270,8 +270,8 @@ def forward(
         coses = []
         sines = []
         if self.position_encoding_2d:
-            block_position_ids = position_ids[:, 1, :].transpose([1, 0])
-            position_ids = position_ids[:, 0, :].transpose([1, 0])
+            block_position_ids = position_ids[:batch_size, 1, :].transpose([1, 0])
+            position_ids = position_ids[:batch_size, 0, :].transpose([1, 0])
             coses.append(cos.squeeze(1)[position_ids].unsqueeze(2))
             sines.append(sin.squeeze(1)[position_ids].unsqueeze(2))
 

paddlenlp/experimental/transformers/fused_transformer_layers.py

@@ -531,6 +531,7 @@ def forward(
         """
         if caches is not None:
             assert len(caches) == len(self.qkv_weights)
+        bsz = input_ids.shape[0]
         bias_residual_input = src
         ln_out = src
         for i in range(len(caches)):
@@ -576,8 +577,8 @@ def forward(
                     )
 
                 if pre_caches is not None:
-                    k_out = paddle.concat([pre_caches[i][0], k_out], axis=2)
-                    v_out = paddle.concat([pre_caches[i][1], v_out], axis=2)
+                    k_out = paddle.concat([pre_caches[i][0, :bsz], k_out], axis=2)
+                    v_out = paddle.concat([pre_caches[i][1, :bsz], v_out], axis=2)
 
                 # write cache kv (inplace)
                 write_cache_kv(k_out, v_out, caches[i], seq_lens + pre_caches_length)