Fix offload (#288)

Author: tastelikefeet

swift/aigc/animatediff.py

@@ -525,7 +525,7 @@ def state_dict(self,
                                **kwargs):
                     state_dict = self.state_dict_origin()
                     return {
-                        key: value
+                        key.replace('base_layer.', ''): value
                         for key, value in state_dict.items()
                         if 'lora' not in key
                     }

swift/llm/sft.py

@@ -9,22 +9,19 @@
 from modelscope import BitsAndBytesConfig, GenerationConfig
 
 from swift.trainers import (IntervalStrategy, Seq2SeqTrainer,
-                            Seq2SeqTrainingArguments)
-from swift.tuners import (LongLoRAConfig, LongLoRAModelType, LoraConfig,
-                          LoRAConfig, NEFTuneConfig, Swift)
+                            Seq2SeqTrainingArguments, TrainerCallback)
 from swift.utils import (check_json_format, compute_acc_metrics,
-                         compute_nlg_metrics, freeze_model_parameters,
-                         get_dist_setting, get_logger, get_main,
-                         get_model_info, is_ddp_plus_mp, is_dist, is_master,
-                         plot_images, preprocess_logits_for_metrics,
+                         compute_nlg_metrics, get_dist_setting, get_logger,
+                         get_main, get_model_info, is_ddp_plus_mp, is_dist,
+                         is_master, plot_images, preprocess_logits_for_metrics,
                          seed_everything, show_layers)
 from .tuner import prepare_model
 from .utils import (LazyLLMDataset, SftArguments, Template,
                     add_self_cognition_dataset, data_collate_fn, dataset_map,
-                    find_all_linear_for_lora, get_additional_saved_files,
-                    get_dataset, get_model_tokenizer, get_template,
-                    get_time_info, print_example, set_generation_config,
-                    sort_by_max_length, stat_dataset)
+                    get_additional_saved_files, get_dataset,
+                    get_model_tokenizer, get_template, get_time_info,
+                    print_example, set_generation_config, sort_by_max_length,
+                    stat_dataset)
 
 logger = get_logger()
 
@@ -234,13 +231,19 @@ def llm_sft(args: SftArguments) -> Dict[str, Union[str, Any]]:
     if args.check_model_is_latest is False:
         trainer_kwargs['check_model'] = False
 
+    class TrainerAdapterCallback(TrainerCallback):
+
+        def on_train_begin(*args, **kwargs):
+            model.set_active_adapters(model.adapters.keys(), offload='meta')
+
     trainer = Seq2SeqTrainer(
         model=model,
         args=training_args,
         data_collator=data_collator,
         train_dataset=train_dataset,
         eval_dataset=val_dataset,
         tokenizer=tokenizer,
+        callbacks=[TrainerAdapterCallback()],
         **trainer_kwargs)
     trainer.sft_args = args
     if is_master():

swift/trainers/__init__.py

@@ -8,7 +8,7 @@
     from .dpo_trainers import DPOTrainer
     from .trainers import Seq2SeqTrainer, Trainer
     from .utils import EvaluationStrategy, FSDPOption, HPSearchBackend, HubStrategy, \
-        IntervalStrategy, SchedulerType, ShardedDDPOption
+        IntervalStrategy, SchedulerType, ShardedDDPOption, TrainerCallback
 else:
     _import_structure = {
         'arguments': ['Seq2SeqTrainingArguments', 'TrainingArguments'],
@@ -17,7 +17,7 @@
         'utils': [
             'EvaluationStrategy', 'FSDPOption', 'HPSearchBackend',
             'HubStrategy', 'IntervalStrategy', 'SchedulerType',
-            'ShardedDDPOption'
+            'ShardedDDPOption', 'TrainerCallback'
         ]
     }
 

swift/trainers/utils.py

@@ -6,6 +6,7 @@
 from typing import List, Union
 
 from torch.nn import Module
+from transformers.trainer_callback import TrainerCallback
 from transformers.trainer_utils import (EvaluationStrategy, FSDPOption,
                                         HPSearchBackend, HubStrategy,
                                         IntervalStrategy, SchedulerType)

swift/tuners/base.py

@@ -464,20 +464,20 @@ def set_active_adapters(self,
 
         adapter_names = set(adapter_names)
         for adapter_name in (adapter_names & set(self.adapters.keys())):
-            self.activate_adapter(adapter_name)
+            self.activate_adapter(adapter_name, offload)
 
         for adapter_name in (set(self.adapters.keys()) - adapter_names):
             self.deactivate_adapter(adapter_name, offload)
 
-    def activate_adapter(self, adapter_name):
+    def activate_adapter(self, adapter_name, offload=None):
         if adapter_name not in self.adapters:
             logger.warning(
                 f'{adapter_name} not in adapters: {self.adapters.keys()}')
             return
 
         from .mapping import SWIFT_MAPPING
         SWIFT_MAPPING[self.adapters[adapter_name].config.swift_type][1]\
-            .activate_adapter(self.base_model, adapter_name, True)
+            .activate_adapter(self.base_model, adapter_name, True, offload)
 
     def deactivate_adapter(self, adapter_name, offload=None):
         if adapter_name not in self.adapters:

swift/tuners/lora_layers.py

@@ -16,6 +16,7 @@
 from peft.tuners.lora import Conv2d as _Conv2d
 from peft.tuners.lora import Embedding as _Embedding
 from peft.tuners.lora import Linear as _Linear
+from peft.tuners.lora import LoraLayer
 from peft.tuners.lora import LoraModel as _LoraModel
 from peft.tuners.lora.tp_layer import LoraParallelLinear as _LoraParallelLinear
 from peft.tuners.tuners_utils import BaseTunerLayer