fix(internlm): Prevent errors by padding the dimensions of wrap tokens.

modeling_internlm_xcomposer2.py

@@ -258,15 +258,15 @@ class InternLMXComposer2ForCausalLM(InternLM2PreTrainedModel):
             wrap_target = wrap_target[:, :self.max_length].to(self.device)
             wrap_im_mask = wrap_im_mask[:, :self.max_length].to(self.device)
 
-            wrap_embeds_list.append(wrap_embeds)
-            wrap_atts_list.append(wrap_atts)
-            wrap_target_list.append(wrap_target)
-            wrap_im_mask_list.append(wrap_im_mask)
-
-        wrap_embeds = torch.cat(wrap_embeds_list)
-        wrap_atts = torch.cat(wrap_atts_list)
-        wrap_target = torch.cat(wrap_target_list)
-        wrap_im_mask = torch.cat(wrap_im_mask_list)
+            wrap_embeds_list.append(wrap_embeds.squeeze(0))
+            wrap_atts_list.append(wrap_atts.squeeze(0))
+            wrap_target_list.append(wrap_target.squeeze(0))
+            wrap_im_mask_list.append(wrap_im_mask.squeeze(0))
+
+        wrap_embeds = torch.nn.utils.rnn.pad_sequence(wrap_embeds_list, batch_first=True, padding_value=self.tokenizer._pad_token_type_id)
+        wrap_atts = torch.nn.utils.rnn.pad_sequence(wrap_atts_list, batch_first=True, padding_value=self.tokenizer._pad_token_type_id)
+        wrap_target = torch.nn.utils.rnn.pad_sequence(wrap_target_list, batch_first=True, padding_value=self.tokenizer._pad_token_type_id)
+        wrap_im_mask = torch.nn.utils.rnn.pad_sequence(wrap_im_mask_list, batch_first=True, padding_value=self.tokenizer._pad_token_type_id)
         return wrap_embeds, wrap_atts, wrap_target, wrap_im_mask
 
     def mask_human_targets(self, input_ids, pure=False):