support ModuleToSave original module offloading (#282)

Author: tastelikefeet

swift/tuners/base.py

@@ -49,13 +49,21 @@ def __init__(self,
             model = model.base_model
 
         if isinstance(config, SwiftConfig):
-            self.adapters[DEFAULT_ADAPTER] = self._prepare_model(
-                model, config, DEFAULT_ADAPTER)
+            if DEFAULT_ADAPTER not in self.adapters:
+                self.adapters[DEFAULT_ADAPTER] = self._prepare_model(
+                    model, config, DEFAULT_ADAPTER)
+            else:
+                logger.warn(
+                    f'Adater {DEFAULT_ADAPTER} has been patched, skip.')
         elif isinstance(config, dict):
             assert (all(isinstance(c, SwiftConfig) for c in config.values()))
             for adapter_name, _config in config.items():
-                self.adapters[adapter_name] = self._prepare_model(
-                    model, _config, adapter_name)
+                if adapter_name not in self.adapters:
+                    self.adapters[adapter_name] = self._prepare_model(
+                        model, _config, adapter_name)
+                else:
+                    logger.warn(
+                        f'Adater {adapter_name} has been patched, skip.')
         self.model = model
 
         self.extra_state_keys = extra_state_keys or []
@@ -195,7 +203,8 @@ def load_state_file(path):
     def from_pretrained(cls,
                         model: Union[nn.Module, 'SwiftModel'],
                         model_id: str = None,
-                        adapter_name: Union[str, List[str]] = None,
+                        adapter_name: Union[str, List[str], Dict[str,
+                                                                 str]] = None,
                         inference_mode: bool = False,
                         revision: str = None,
                         **kwargs):
@@ -205,7 +214,7 @@ def from_pretrained(cls,
             model (`Union[torch.nn.Module, 'SwiftModel']`): The model to be tuned,
                 if the model is already a `SwiftModel` it will be un-wrapped and re-wrapped..
             model_id (`str`): The model_id or a local model dir of tuners to use to tune the model.
-            adapter_name (`Union[str, List[str]]`): The adapter_names saved in the model repo to load.
+            adapter_name (`Union[str, List[str], Dict[str, str]]`): The adapter_names saved in the model repo to load.
                 Default `None`, means load all tuners saved in the model_id
             inference_mode (`bool`): Use in the inference mode or not.
             revision (`str`): The model revision to use.
@@ -236,7 +245,8 @@ def from_pretrained(cls,
                 os.path.isfile(os.path.join(model_dir, sub_dir, CONFIG_NAME))
             ]
         for _name in adapter_name if isinstance(adapter_name,
-                                                list) else [adapter_name]:
+                                                list) else [adapter_name] \
+                if isinstance(adapter_name, str) else adapter_name.keys():
             sub_folder = os.path.join(model_dir, _name)
             config_file = os.path.join(sub_folder, CONFIG_NAME)
 
@@ -250,26 +260,31 @@ def from_pretrained(cls,
             if SWIFT_TYPE_KEY not in json_object:
                 raise ValueError('Mixed using with peft is not allowed now.')
             else:
-                adapters[_name] = SwiftConfig.from_pretrained(sub_folder)
+                key = _name if not isinstance(adapter_name,
+                                              dict) else adapter_name[_name]
+                adapters[key] = SwiftConfig.from_pretrained(sub_folder)
 
         self = SwiftModel(model, adapters, extra_state_keys, inference_mode,
                           **kwargs)
         for _name in adapter_name if isinstance(adapter_name,
-                                                list) else [adapter_name]:
+                                                list) else [adapter_name] \
+                if isinstance(adapter_name, str) else adapter_name.keys():
             sub_folder = os.path.join(model_dir, _name)
             state_dict = cls.load_state_file(sub_folder)
+            _adapter = _name if not isinstance(adapter_name,
+                                               dict) else adapter_name[_name]
             if state_dict is not None:
                 model_is_qlora = len([
                     k for k in self.state_dict().keys()
-                    if k.endswith('.lora_A.default.weight')
-                    or k.endswith('.lora_B.default.weight')
+                    if k.endswith(f'.lora_A.{_adapter}.weight')
+                    or k.endswith(f'.lora_B.{_adapter}.weight')
                 ])
                 if not model_is_qlora:
                     # model is lora, state_dict: qlora->lora
                     state_dict = {
-                        k[:-len('.default.weight') if k.
-                          endswith('.lora_A.default.weight') or k.
-                          endswith('.lora_B.default.weight') else None]: v
+                        k[:-len(f'.{_name}.weight') if k.
+                          endswith(f'.lora_A.{_name}.weight') or k.
+                          endswith(f'.lora_B.{_name}.weight') else None]: v
                         for k, v in state_dict.items()
                     }
                 if any(['loramodule' in key for key in state_dict]):
@@ -288,7 +303,13 @@ def from_pretrained(cls,
                                     f'lora_B.{_name}.weight'): value
                         for key, value in state_dict.items()
                     }
-                self.load_state_dict(state_dict, adapter_name=_name)
+                if isinstance(adapter_name, dict):
+                    # TODO this logic is fragile! replace `_name` may cause other parts replaced
+                    state_dict = {
+                        key.replace(_name, adapter_name[_name]): value
+                        for key, value in state_dict.items()
+                    }
+                self.load_state_dict(state_dict, adapter_name=_adapter)
         state_dict = cls.load_state_file(model_dir)
         if state_dict is not None:
             self.load_state_dict(state_dict)
@@ -569,7 +590,8 @@ def unmerge(model: Union[PeftModel, SwiftModel], **kwargs):
     @staticmethod
     def from_pretrained(model: Union[nn.Module, SwiftModel],
                         model_id: str = None,
-                        adapter_name: Union[str, List[str]] = None,
+                        adapter_name: Union[str, List[str], Dict[str,
+                                                                 str]] = None,
                         revision: str = None,
                         **kwargs):
         """Prepare a model by a model_id in the ModelScope hub or a local dir.
@@ -593,7 +615,8 @@ def from_pretrained(model: Union[nn.Module, SwiftModel],
             is_peft_model = SWIFT_TYPE_KEY not in _json
 
         _name = adapter_name if isinstance(
-            adapter_name, str) or adapter_name is None else adapter_name[0]
+            adapter_name, str) or adapter_name is None else adapter_name[0] \
+            if isinstance(adapter_name, list) else list(adapter_name.keys())[0]
         _name = _name or ''
         if os.path.exists(os.path.join(model_id, _name, CONFIG_NAME)):
             with open(os.path.join(model_id, _name, CONFIG_NAME), 'r') as f:

swift/tuners/lora.py

@@ -78,7 +78,8 @@ def activate_adapter(module: torch.nn.Module,
         for sub_module in module.modules():
             if isinstance(sub_module, (LoraLayer, LoRALayer)):
                 sub_module.set_activation(adapter_name, activate)
-                sub_module.save_memory(adapter_name, activate, offload)
+                if hasattr(sub_module, 'save_memory'):
+                    sub_module.save_memory(adapter_name, activate, offload)
 
     @staticmethod
     def unpatch_lora(model, config: LoRAConfig, adapter_name: str):

swift/tuners/lora_layers.py

@@ -272,7 +272,7 @@ def inject_adapter(self, model: nn.Module, adapter_name: str):
 
                 if not isinstance(target, ModulesToSaveWrapper):
                     new_module = ModulesToSaveWrapper(
-                        target, adapter_name, module_key=key)
+                        target, adapter_name=adapter_name, module_key=key)
                     setattr(parent, target_name, new_module)
                 else:
                     target.update(adapter_name)
@@ -489,6 +489,7 @@ def _create_new_module(lora_config, adapter_name, target, **kwargs):
         elif lora_config.use_merged_linear:
             new_module = MergedLinear(
                 adapter_name,
+                current_key,
                 target,
                 bias=bias,
                 enable_lora=lora_config.enable_lora,

swift/tuners/neftune.py

@@ -64,8 +64,10 @@ def mark_trainable_callback(model):
                            mark_trainable_callback)
 
     @staticmethod
-    def activate_adapter(module: torch.nn.Module, adapter_name: str,
-                         activate: bool):
+    def activate_adapter(module: torch.nn.Module,
+                         adapter_name: str,
+                         activate: bool,
+                         offload: str = None):
         for sub_module in module.modules():
             if isinstance(sub_module, torch.nn.Embedding):
                 sub_module.nef_activated = activate

swift/tuners/scetuning/scetuning.py

@@ -154,7 +154,8 @@ def _get_module(module):
 
         # refactor forward function
         def _forward_encoder_mode(self, *args, **kwargs):
-            args = self.forward_origin(*args, **kwargs)
+            args = getattr(self, f'forward_origin_{adapter_name}')(*args,
+                                                                   **kwargs)
             args_type = type(args)
             if args_type is tuple:
                 args = args[0]
@@ -185,12 +186,15 @@ def _forward_decoder_mode(self, *args, **kwargs):
             if args_type is tuple:
                 args_main = (args_sub_tuner_new, *args_sub_extra)
 
-            args_main = self.forward_origin(*args_main, **kwargs)
+            args_main = getattr(self,
+                                f'forward_origin_{adapter_name}')(*args_main,
+                                                                  **kwargs)
             return args_main
 
         # 3. inject the tuners
         for tuner_id, t_module in enumerate(target_module_ins_list):
-            t_module.forward_origin = getattr(t_module, 'forward')
+            setattr(t_module, f'forward_origin_{adapter_name}',
+                    getattr(t_module, 'forward'))
             if config.tuner_mode in ('encoder', 'identity'):
                 _forward = _forward_encoder_mode
             elif config.tuner_mode == 'decoder':

swift/tuners/utils.py

@@ -12,6 +12,7 @@
 import json
 import numpy as np
 import torch
+from packaging import version
 from peft.utils import CONFIG_NAME
 from peft.utils import ModulesToSaveWrapper as _ModulesToSaveWrapper
 from peft.utils import _get_submodules
@@ -252,7 +253,27 @@ def load_disk(module: torch.nn.Module, adapter_name, module_key):
             file = os.path.join(sub_folder, f'{key}.dat')
             state_dict[key] = OffloadHelper.load_offloaded_weight(
                 file, OffloadHelper.index[md5][key])
-        module.load_state_dict(state_dict, assign=True)
+        if version.parse(torch.__version__) >= version.parse('2.1.0'):
+            module.load_state_dict(state_dict, assign=True)
+        else:
+            for name, _module in module.named_modules():
+                if len(list(_module.modules())) > 1:
+                    continue
+
+                buffers = {}
+                prefix = name if not name else name + '.'
+                for sub_name, buffer in _module.named_buffers():
+                    buffer_cls = type(buffer)
+                    buffers[sub_name] = buffer_cls(state_dict[prefix
+                                                              + sub_name])
+                _module._buffers.update(buffers)
+                params = {}
+                for sub_name, param in _module.named_parameters():
+                    param_cls = type(param)
+                    params[sub_name] = param_cls(
+                        state_dict[prefix + sub_name],
+                        requires_grad=param.requires_grad)
+                _module._parameters.update(params)
         shutil.rmtree(sub_folder, ignore_errors=True)
 
 
@@ -295,7 +316,7 @@ def offload(module: torch.nn.Module, adapter_name, module_key,
         if offload == 'cpu':
             if str(device) != 'cpu':
                 module.to('cpu')
-        if offload == 'meta':
+        elif offload == 'meta':
             if str(device) != 'meta':
                 OffloadHelper.offload_disk(
                     module, adapter_name=adapter_name, module_key=module_key)
@@ -331,6 +352,12 @@ class ModulesToSaveWrapper(ActivationMixin, _ModulesToSaveWrapper):
     def __init__(self, *args, module_key, **kwargs):
         super(ModulesToSaveWrapper, self).__init__(module_key)
         super(ActivationMixin, self).__init__(*args, **kwargs)
+        SwiftAdapter.save_memory(
+            self.original_module,
+            'original_module',
+            self.module_key,
+            False,
+            offload='cpu')
 
     @property
     def active_adapter(self):
@@ -343,7 +370,7 @@ def active_adapter(self):
             )
         return active_adapters[0]
 
-    def set_adapter(self, adapter_name: str, offload: str):
+    def set_adapter(self, adapter_name: str, offload: str = None):
         if adapter_name not in self.modules_to_save:
             raise ValueError(
                 f'Adapter {adapter_name} not found in {self.modules_to_save.keys()}'
@@ -352,8 +379,14 @@ def set_adapter(self, adapter_name: str, offload: str):
         self.set_activation(adapter_name, True)
         SwiftAdapter.save_memory(self.modules_to_save[adapter_name],
                                  adapter_name, self.module_key, True)
+        SwiftAdapter.save_memory(
+            self.original_module,
+            'original_module',
+            self.module_key,
+            False,
+            offload=offload)
 
-    def deactivate_adapter(self, adapter_name: str, offload: str):
+    def deactivate_adapter(self, adapter_name: str, offload: str = None):
         if adapter_name in self.modules_to_save and self.unique_thread:
             self.modules_to_save[adapter_name].requires_grad_(False)
         self.set_activation(adapter_name, False)
@@ -363,6 +396,22 @@ def deactivate_adapter(self, adapter_name: str, offload: str):
             self.module_key,
             False,
             offload=offload)
+        if not self.get_activated_adapters():
+            SwiftAdapter.save_memory(self.original_module, 'original_module',
+                                     self.module_key, True)
+
+    def enable_adapters(self, enabled: bool):
+        super().enable_adapters(enabled)
+        if not enabled:
+            SwiftAdapter.save_memory(
+                self.original_module,
+                'original_module',
+                self.module_key,
+                False,
+                offload='meta')
+        else:
+            SwiftAdapter.save_memory(self.original_module, 'original_module',
+                                     self.module_key, True)
 
 
 def set_adapter(model, adapter_name, activate, offload):
@@ -385,6 +434,7 @@ def set_trainable(model, adapter_name):
                 target.update(adapter_name)
                 target.set_adapter(target.active_adapter)
             else:
-                new_module = ModulesToSaveWrapper(target, adapter_name)
+                new_module = ModulesToSaveWrapper(
+                    target, module_key=key, adapter_name=adapter_name)
                 new_module.set_adapter(adapter_name)
                 setattr(parent, target_name, new_module)

tests/tuners/test_swift_base.py

@@ -202,9 +202,12 @@ def reset_lora_parameters(self, adapter_name, init_lora_weights):
             os.path.exists(
                 os.path.join(self.tmp_dir, 'default', WEIGHTS_NAME)))
 
-        model2 = Swift.from_pretrained(model2, self.tmp_dir)
+        model2 = Swift.from_pretrained(
+            model2, self.tmp_dir, adapter_name={'default': 'test'})
+        self.assertTrue('test' in model2.adapters)
         output2 = model2(**input)
         self.assertTrue(torch.allclose(output1.logits, output2.logits))
+        model2 = Swift.from_pretrained(model2, self.tmp_dir)
         state_dict = model.state_dict()
         state_dict2 = model2.state_dict()
         for key in state_dict: