Support longlora for transformers 4.38 (#456)

Author: tastelikefeet

docs/source/LLM/支持的模型和数据集.md

@@ -275,6 +275,7 @@
 |generated-chat-zh|[AI-ModelScope/generated_chat_0.4M](https://modelscope.cn/datasets/AI-ModelScope/generated_chat_0.4M/summary)|396004|0|273.3±52.0, min=32, max=873|chat, character-dialogue|
 |cls-fudan-news-zh|[damo/zh_cls_fudan-news](https://modelscope.cn/datasets/damo/zh_cls_fudan-news/summary)|4959|0|3234.4±2547.5, min=91, max=19548|chat, classification|
 |ner-jave-zh|[damo/zh_ner-JAVE](https://modelscope.cn/datasets/damo/zh_ner-JAVE/summary)|1266|0|118.3±45.5, min=44, max=223|chat, ner|
+|long-alpaca-12k|[AI-ModelScope/LongAlpaca-12k](https://modelscope.cn/datasets/AI-ModelScope/LongAlpaca-12k/summary)|11998|0|9619.0±8295.8, min=36, max=78925|longlora, QA|
 |coco-en|[modelscope/coco_2014_caption](https://modelscope.cn/datasets/modelscope/coco_2014_caption/summary)|414113|40504|298.8±2.8, min=294, max=351|chat, multi-modal, vision|
 |🔥coco-mini-en|[modelscope/coco_2014_caption](https://modelscope.cn/datasets/modelscope/coco_2014_caption/summary)|20000|200|298.8±2.8, min=294, max=339|chat, multi-modal, vision|
 |🔥coco-mini-en-2|[modelscope/coco_2014_caption](https://modelscope.cn/datasets/modelscope/coco_2014_caption/summary)|20000|200|36.8±2.8, min=32, max=77|chat, multi-modal, vision|

swift/llm/sft.py

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

swift/llm/tuner.py

@@ -96,7 +96,6 @@ def prepare_model(model, args: SftArguments):
                 longlora_config = LongLoRAConfig(
                     lora_dtype=args.lora_dtype,
                     model_type=LongLoRAModelType.LLAMA,
-                    use_flash_attn=args.use_flash_attn,
                     **lora_kwargs)
                 model = Swift.prepare_model(model, longlora_config)
                 logger.info(f'longlora_config: {longlora_config}')

swift/llm/utils/dataset.py

@@ -110,6 +110,7 @@ class DatasetName:
     # example dataset for specific model
     cls_fudan_news_zh = 'cls-fudan-news-zh'  # seqgpt-560m
     ner_java_zh = 'ner-jave-zh'  # seqgpt-560m
+    long_alpaca_12k = 'long-alpaca-12k'
 
     # multi-modal
     # for qwen-vl
@@ -457,6 +458,24 @@ def _repair_ms_bench(conversations: str) -> Dict[str, str]:
     return conversations
 
 
+def long_alpaca_preprocessor(dataset: HfDataset):
+
+    def map_row(row):
+        response = row['response']
+        if response and response.startswith('Answer:'):
+            response = response[len('Answer:') + 1:].strip()
+        return {'query': row['query'], 'response': response}
+    return dataset.rename_columns({'instruction': 'query', 'output': 'response'})\
+        .remove_columns(['input', 'file']).map(map_row).filter(lambda row: row['response'] is not None)
+
+
+register_dataset(
+    DatasetName.long_alpaca_12k,
+    'AI-ModelScope/LongAlpaca-12k', ['train'], [],
+    long_alpaca_preprocessor,
+    get_dataset_from_repo,
+    tags=['longlora', 'QA'])
+
 register_dataset(
     DatasetName.ms_bench,
     'iic/ms_bench', ['train'], [],

swift/llm/utils/template.py

@@ -8,7 +8,8 @@
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.utils.rnn import pad_sequence
-from transformers import PreTrainedTokenizerBase, StoppingCriteria
+from transformers import (DataCollatorForSeq2Seq, PreTrainedTokenizerBase,
+                          StoppingCriteria)
 
 from swift.llm.agent.utils import calculate_loss_scale
 
@@ -186,6 +187,10 @@ 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'
         ]:
@@ -386,55 +391,28 @@ def concat_tokenizer_kwargs(
         assert len(old_tokenizer_kwargs) == 0
         return curr_tokenizer_kwargs
 
-    def data_collator(self,
-                      batch: List[Dict[str, Any]],
-                      padding_to: Optional[int] = None) -> Dict[str, Any]:
+    def data_collator(
+            self,
+            batch: List[Dict[str, Any]],
+            pad_to_multiple_of: Optional[int] = None) -> Dict[str, Any]:
         """
         Args:
             batch(`List[Dict[str, Any]]`): The input data in batch
-            padding_to(`int`, optional): Whether padding the batch to a fixed length, if none, the batch
-                will be padded to the `longest`
+            pad_to_multiple_of(`int`, optional): Whether padding to the multiple of an integer value.
         """
-        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'])
+        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'))
                       for b in batch] if 'loss_scale' in batch[0] else None
-        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 - 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)
+        res = self._data_collator(batch, return_tensors='pt')
+        padding_to = res['input_ids'].shape[1]
         if loss_scale:
+            loss_scale[0] = F.pad(loss_scale[0],
+                                  (0, padding_to - loss_scale[0].shape[-1]),
+                                  'constant', 0.)
             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
 
@@ -601,10 +579,11 @@ def encode(
         inputs['images'] = image_tensor.to(model.dtype)
         return inputs, {}
 
-    def data_collator(self,
-                      batch: List[Dict[str, Any]],
-                      padding_to: Optional[int] = None) -> Dict[str, Any]:
-        res = super().data_collator(batch, padding_to)
+    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)
         res['images'] = torch.concat([b['images'] for b in batch])
         return res
 
@@ -908,10 +887,11 @@ def encode(
         inputs['image_sizes'] = image_sizes
         return inputs, {}
 
-    def data_collator(self,
-                      batch: List[Dict[str, Any]],
-                      padding_to: Optional[int] = None) -> Dict[str, Any]:
-        res = super().data_collator(batch, padding_to)
+    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)
         res['images'] = torch.concat([b['images'] for b in batch])
         res['image_sizes'] = sum([b['image_sizes'] for b in batch], start=[])
         return res
@@ -1093,10 +1073,11 @@ def encode(
             len(inputs['input_ids']) - len(token_type_ids))
         return inputs, {}
 
-    def data_collator(self,
-                      batch: List[Dict[str, Any]],
-                      padding_to: Optional[int] = None) -> Dict[str, Any]:
-        res = super().data_collator(batch, padding_to)
+    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)
         is_cogagent = 'cross_images' in batch[0]
         keys = ['images', 'cross_images'] if is_cogagent else ['images']
         for key in keys: