Add longlora for llama (#115)

Author: tastelikefeet

examples/pytorch/llm/src/llm_sft.py

@@ -11,8 +11,8 @@
 from utils import (SftArguments, dataset_map, get_dataset, get_model_tokenizer,
                    get_preprocess)
 
-from swift import (LoraConfig, LoRAConfig, Seq2SeqTrainer,
-                   Seq2SeqTrainingArguments, Swift, get_logger)
+from swift import (LongLoRAConfig, LongLoRAModelType, LoraConfig, LoRAConfig,
+                   Seq2SeqTrainer, Seq2SeqTrainingArguments, Swift, get_logger)
 from swift.utils import (add_version_to_work_dir, broadcast_string,
                          check_json_format, compute_nlg_metrics,
                          data_collate_fn, find_all_linear_for_lora,
@@ -54,27 +54,40 @@ def llm_sft(args: SftArguments) -> None:
         args.model_type, torch_dtype=args.torch_dtype, **kwargs)
 
     # ### Preparing LoRA
-    if args.sft_type == 'lora':
+    if args.sft_type == 'lora' or args.sft_type == 'longlora':
         if args.resume_from_checkpoint is None:
             if 'ALL' in args.lora_target_modules:
                 assert len(args.lora_target_modules) == 1
                 args.lora_target_modules = find_all_linear_for_lora(
                     model, args.quantization_bit, args.model_type)
                 logger.info(
                     f'Setting lora_target_modules: {args.lora_target_modules}')
-            lora_kwargs = {}
-            if args.tuner_bankend == 'peft':
-                global LoRAConfig
-                LoRAConfig = LoraConfig
-                lora_kwargs['task_type'] = 'CAUSAL_LM'
-            lora_config = LoRAConfig(
-                r=args.lora_rank,
-                target_modules=args.lora_target_modules,
-                lora_alpha=args.lora_alpha,
-                lora_dropout=args.lora_dropout_p,
-                **lora_kwargs)
-            model = Swift.prepare_model(model, lora_config)
-            logger.info(f'lora_config: {lora_config}')
+            if args.sft_type == 'lora':
+                lora_kwargs = {}
+                if args.tuner_bankend == 'peft':
+                    global LoRAConfig
+                    LoRAConfig = LoraConfig
+                    lora_kwargs['task_type'] = 'CAUSAL_LM'
+                lora_config = LoRAConfig(
+                    r=args.lora_rank,
+                    target_modules=args.lora_target_modules,
+                    lora_alpha=args.lora_alpha,
+                    lora_dropout=args.lora_dropout_p,
+                    **lora_kwargs)
+                model = Swift.prepare_model(model, lora_config)
+                logger.info(f'lora_config: {lora_config}')
+            elif args.sft_type == 'longlora':
+                assert args.tuner_bankend != 'peft'
+                assert LongLoRAModelType.LLAMA in args.model_type
+                longlora_config = LongLoRAConfig(
+                    r=args.lora_rank,
+                    target_modules=args.lora_target_modules,
+                    lora_alpha=args.lora_alpha,
+                    lora_dropout=args.lora_dropout_p,
+                    model_type=LongLoRAModelType.LLAMA,
+                    use_flash_attn=args.use_flash_attn)
+                model = Swift.prepare_model(model, longlora_config)
+                logger.info(f'longlora_config: {longlora_config}')
         else:
             model = Swift.from_pretrained(
                 model, args.resume_from_checkpoint, is_trainable=True)
@@ -109,7 +122,10 @@ def llm_sft(args: SftArguments) -> None:
     if args.test_oom_error:
         train_dataset = sort_by_max_length(train_dataset, 20000)
     # Data analysis
-    data_collator = partial(data_collate_fn, tokenizer=tokenizer)
+    data_collator = partial(
+        data_collate_fn,
+        tokenizer=tokenizer,
+        padding_to=args.max_length if args.sft_type == 'longlora' else None)
     print_example(train_dataset[0], tokenizer)
     stat_dataset(train_dataset)
     stat_dataset(val_dataset)

examples/pytorch/llm/src/utils/argument.py

@@ -24,7 +24,7 @@ class SftArguments:
         default=ModelType.qwen_7b_chat,
         metadata={'choices': list(MODEL_MAPPING.keys())})
     sft_type: str = field(
-        default='lora', metadata={'choices': ['lora', 'full']})
+        default='lora', metadata={'choices': ['longlora', 'lora', 'full']})
     tuner_bankend: str = field(
         default='swift', metadata={'choices': ['swift', 'peft']})
     template_type: Optional[str] = field(
@@ -147,7 +147,7 @@ def init_argument(self):
             # Initialize in advance
             dist.init_process_group(backend=self.ddp_backend)
 
-        if self.sft_type == 'lora':
+        if self.sft_type == 'lora' or self.sft_type == 'longlora':
             if self.learning_rate is None:
                 self.learning_rate = 1e-4
             if self.only_save_model is None:
@@ -223,7 +223,7 @@ class InferArguments:
         default=ModelType.qwen_7b_chat,
         metadata={'choices': list(MODEL_MAPPING.keys())})
     sft_type: str = field(
-        default='lora', metadata={'choices': ['lora', 'full']})
+        default='lora', metadata={'choices': ['longlora', 'lora', 'full']})
     template_type: Optional[str] = field(
         default=None, metadata={'choices': list(TEMPLATE_MAPPING.keys())})
     ckpt_dir: str = '/path/to/your/vx_xxx/checkpoint-xxx'

swift/__init__.py

@@ -13,7 +13,8 @@
         PrefixTuningConfig, PromptEncoderConfig, PromptLearningConfig,
         PromptTuningConfig, get_peft_config, get_peft_model,
         get_peft_model_state_dict, Prompt, PromptConfig, PromptModule,
-        SwiftConfig, SwiftOutput, Swift, SwiftTuners)
+        SwiftConfig, SwiftOutput, Swift, SwiftTuners, LongLoRAConfig, LongLoRA,
+        LongLoRAModelType)
     from .hub import snapshot_download, push_to_hub, push_to_hub_async, push_to_hub_in_queue
     from .trainers import (EvaluationStrategy, FSDPOption, HPSearchBackend,
                            HubStrategy, IntervalStrategy, SchedulerType,
@@ -38,7 +39,7 @@
             'PromptTuningConfig', 'get_peft_config', 'get_peft_model',
             'get_peft_model_state_dict', 'Prompt', 'PromptConfig',
             'PromptModule', 'SwiftConfig', 'SwiftOutput', 'Swift',
-            'SwiftTuners'
+            'SwiftTuners', 'LongLoRAConfig', 'LongLoRA', 'LongLoRAModelType'
         ],
         'trainers': [
             'EvaluationStrategy', 'FSDPOption', 'HPSearchBackend',

swift/tuners/__init__.py

@@ -9,6 +9,7 @@
     from .lora import LoRA, LoRAConfig
     from .mapping import SWIFT_MAPPING, SwiftTuners
     from .side import Side, SideConfig, SideModule
+    from .longlora.longlora import LongLoRAModelType, LongLoRAConfig, LongLoRA
     from .restuning import ResTuning, ResTuningConfig, ResTuningBypassModule
     from .peft import (LoraConfig, PeftConfig, PeftModel, PeftModelForCausalLM,
                        PeftModelForSeq2SeqLM,
@@ -24,6 +25,8 @@
         'adapter': ['Adapter', 'AdapterConfig', 'AdapterModule'],
         'base': ['SwiftModel', 'Swift'],
         'lora': ['LoRA', 'LoRAConfig'],
+        'longlora.longlora':
+        ['LongLoRAModelType', 'LongLoRAConfig', 'LongLoRA'],
         'mapping': ['SWIFT_MAPPING', 'SwiftTuners'],
         'side': ['Side', 'SideConfig', 'SideModule'],
         'restuning': ['ResTuning', 'ResTuningConfig', 'ResTuningBypassModule'],