Fix data_collator (#674)

Author: tastelikefeet

swift/llm/sft.py

@@ -182,9 +182,8 @@ def llm_sft(args: SftArguments) -> Dict[str, Union[str, Any]]:
         if val_dataset is not None:
             val_dataset = LazyLLMDataset(val_dataset, template)
 
-    pad_to_multiple_of = 8 if args.sft_type == 'longlora' else None
-    data_collator = partial(
-        template.data_collator, pad_to_multiple_of=pad_to_multiple_of)
+    padding_to = args.max_length if args.sft_type == 'longlora' else None
+    data_collator = partial(template.data_collator, padding_to=padding_to)
 
     # Trainer
     logger.info(f'training_args: {training_args}')

swift/llm/utils/template.py

@@ -8,8 +8,7 @@
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.utils.rnn import pad_sequence
-from transformers import (DataCollatorForSeq2Seq, PreTrainedTokenizerBase,
-                          StoppingCriteria)
+from transformers import PreTrainedTokenizerBase, StoppingCriteria
 
 from swift.llm.agent.utils import calculate_loss_scale
 
@@ -187,10 +186,6 @@ def _init_template(self,
         self.truncation_strategy = truncation_strategy
         self.model = kwargs.get('model', None)
         self.use_loss_scale = kwargs.get('use_loss_scale', False)
-        self._data_collator = DataCollatorForSeq2Seq(
-            tokenizer=self.tokenizer,
-            label_pad_token_id=self.tokenizer.pad_token_id,
-        )
         for key in [
                 'prefix', 'prompt', 'chat_sep', 'suffix', 'prefix_has_system'
         ]:
@@ -391,28 +386,55 @@ def concat_tokenizer_kwargs(
         assert len(old_tokenizer_kwargs) == 0
         return curr_tokenizer_kwargs
 
-    def data_collator(
-            self,
-            batch: List[Dict[str, Any]],
-            pad_to_multiple_of: Optional[int] = None) -> Dict[str, Any]:
+    def data_collator(self,
+                      batch: List[Dict[str, Any]],
+                      padding_to: Optional[int] = None) -> Dict[str, Any]:
         """
         Args:
             batch(`List[Dict[str, Any]]`): The input data in batch
-            pad_to_multiple_of(`int`, optional): Whether padding to the multiple of an integer value.
+            padding_to(`int`, optional): Whether padding the batch to a fixed length, if none, the batch
+                will be padded to the `longest`
         """
-        self._data_collator.pad_to_multiple_of = pad_to_multiple_of
-        if pad_to_multiple_of:
-            self.tokenizer.padding_side = 'right'
-        loss_scale = [torch.tensor(b.pop('loss_scale'))
+        tokenizer = self.tokenizer
+        assert tokenizer.pad_token_id is not None
+        input_ids = [torch.tensor(b['input_ids']) for b in batch]
+        labels = [torch.tensor(b['labels']) for b in batch]
+        loss_scale = [torch.tensor(b['loss_scale'])
                       for b in batch] if 'loss_scale' in batch[0] else None
-        res = self._data_collator(batch, return_tensors='pt')
-        padding_to = res['input_ids'].shape[1]
+        attention_mask = [
+            torch.ones(len(input_ids[i]), dtype=torch.int64)
+            for i in range(len(input_ids))
+        ]
+
+        if padding_to is not None:
+            padding_len = padding_to - input_ids[0].shape[-1]
+            if padding_len > 0:
+                input_ids[0] = F.pad(input_ids[0], (0, padding_len),
+                                     'constant', tokenizer.pad_token_id)
+                attention_mask[0] = F.pad(attention_mask[0], (0, padding_len),
+                                          'constant', 0)
+                labels[0] = F.pad(labels[0], (0, padding_len), 'constant',
+                                  -100)
         if loss_scale:
             loss_scale[0] = F.pad(loss_scale[0],
-                                  (0, padding_to - loss_scale[0].shape[-1]),
+                                  (0, padding_to - labels[0].shape[-1]),
                                   'constant', 0.)
+
+        input_ids = pad_sequence(
+            input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
+        attention_mask = pad_sequence(
+            attention_mask, batch_first=True, padding_value=0)
+        if loss_scale:
             loss_scale = pad_sequence(
                 loss_scale, batch_first=True, padding_value=0.)
+        labels = pad_sequence(labels, batch_first=True, padding_value=-100)
+
+        res = {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask,
+            'labels': labels,
+        }
+        if loss_scale is not None:
             res['loss_scale'] = loss_scale
         return res
 
@@ -579,11 +601,10 @@ def encode(
         inputs['images'] = image_tensor.to(model.dtype)
         return inputs, {}
 
-    def data_collator(
-            self,
-            batch: List[Dict[str, Any]],
-            pad_to_multiple_of: Optional[int] = None) -> Dict[str, Any]:
-        res = super().data_collator(batch, pad_to_multiple_of)
+    def data_collator(self,
+                      batch: List[Dict[str, Any]],
+                      padding_to: Optional[int] = None) -> Dict[str, Any]:
+        res = super().data_collator(batch, padding_to)
         res['images'] = torch.concat([b['images'] for b in batch])
         return res
 
@@ -887,11 +908,10 @@ def encode(
         inputs['image_sizes'] = image_sizes
         return inputs, {}
 
-    def data_collator(
-            self,
-            batch: List[Dict[str, Any]],
-            pad_to_multiple_of: Optional[int] = None) -> Dict[str, Any]:
-        res = super().data_collator(batch, pad_to_multiple_of)
+    def data_collator(self,
+                      batch: List[Dict[str, Any]],
+                      padding_to: Optional[int] = None) -> Dict[str, Any]:
+        res = super().data_collator(batch, padding_to)
         res['images'] = torch.concat([b['images'] for b in batch])
         res['image_sizes'] = sum([b['image_sizes'] for b in batch], start=[])
         return res
@@ -1073,11 +1093,10 @@ def encode(
             len(inputs['input_ids']) - len(token_type_ids))
         return inputs, {}
 
-    def data_collator(
-            self,
-            batch: List[Dict[str, Any]],
-            pad_to_multiple_of: Optional[int] = None) -> Dict[str, Any]:
-        res = super().data_collator(batch, pad_to_multiple_of)
+    def data_collator(self,
+                      batch: List[Dict[str, Any]],
+                      padding_to: Optional[int] = None) -> Dict[str, Any]:
+        res = super().data_collator(batch, padding_to)
         is_cogagent = 'cross_images' in batch[0]
         keys = ['images', 'cross_images'] if is_cogagent else ['images']
         for key in keys: