[PEFT] Peft integration alternative design

src/transformers/__init__.py

@@ -48,6 +48,7 @@
 
 # Base objects, independent of any specific backend
 _import_structure = {
+    "adapters": [],
     "audio_utils": [],
     "benchmark": [],
     "commands": [],

src/transformers/adapters/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .peft_mixin import PeftAdapterMixin

src/transformers/adapters/peft_mixin.py

@@ -0,0 +1,217 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import inspect
+from typing import Optional
+
+from ..utils import find_adapter_config_file, is_accelerate_available, is_peft_available, logging, requires_backends
+
+
+if is_accelerate_available():
+    from accelerate import dispatch_model
+    from accelerate.utils import get_balanced_memory, infer_auto_device_map
+
+
+logger = logging.get_logger(__name__)
+
+
+class PeftAdapterMixin:
+    """
+    A class containing all functions for loading and using adapters weights that are supported in PEFT library.
+    Currently supported PEFT methods are all non-prefix tuning methods
+    """
+
+    _hf_peft_config_loaded = False
+
+    def load_adapter(
+        self,
+        peft_model_id: str,
+        adapter_name: Optional[str] = "default",
+        revision: Optional[str] = None,
+        use_auth_token: Optional[str] = None,
+        commit_hash: Optional[str] = None,
+        device_map: Optional[str] = "auto",
+        max_memory: Optional[int] = None,
+        offload_dir: Optional[str] = None,
+        offload_index: Optional[int] = None,
+    ) -> None:
+        """
+        Load adapter weights from file. Requires peft as a backend to load the adapter weights
+        """
+        requires_backends(self.load_adapter, "peft")
+
+        from peft import PeftConfig, create_and_replace, load_peft_weights
+        from peft.utils import set_peft_model_state_dict
+        from peft.utils.other import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING
+
+        if not self._hf_peft_config_loaded:
+            self.peft_config = {}
+            self._hf_peft_config_loaded = True
+
+        adapter_config_file = find_adapter_config_file(
+            peft_model_id,
+            revision=revision,
+            use_auth_token=use_auth_token,
+            commit_hash=commit_hash,
+        )
+
+        if adapter_config_file is None:
+            raise ValueError(
+                f"adapter model file not found in {peft_model_id}. Make sure you are passing the correct path to the "
+                "adapter model."
+            )
+
+        loaded_peft_config = PeftConfig.from_pretrained(
+            peft_model_id,
+            revision=revision,
+            use_auth_token=use_auth_token,
+            commit_hash=commit_hash,
+        )
+
+        if not hasattr(loaded_peft_config, "target_modules"):
+            target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[self.config.model_type]
+            loaded_peft_config.target_modules = target_modules
+
+        if adapter_name not in self.peft_config:
+            self.peft_config[adapter_name] = loaded_peft_config
+        else:
+            raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.")
+
+        # Replace the adapter with the loaded adapter
+        create_and_replace(loaded_peft_config, self, adapter_name)
+
+        adapter_state_dict = load_peft_weights(
+            peft_model_id,
+            revision=revision,
+            use_auth_token=use_auth_token,
+        )
+
+        # We need to pre-process the state dict to remove unneeded prefixes - for backward compatibility
+        processed_adapter_state_dict = {}
+        for key, value in adapter_state_dict.items():
+            if "base_model.model" in key:
+                new_key = key.replace("base_model.model.", "")
+            else:
+                new_key = key
+            processed_adapter_state_dict[new_key] = value
+
+        # Load state dict
+        incompatible_keys = set_peft_model_state_dict(self, processed_adapter_state_dict, adapter_name)
+
+        if incompatible_keys is not None:
+            # check only for unexpected keys
+            if hasattr(incompatible_keys, "unexpected_keys") and len(incompatible_keys.unexpected_keys) > 0:
+                logger.warning(
+                    f"Loading adapter weights from {peft_model_id} led to unexpected keys not found in the model: "
+                    f" {incompatible_keys.unexpected_keys}. "
+                )
+
+        # @pacman100 why this was needed?
+        if (
+            (getattr(self, "hf_device_map", None) is not None)
+            and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0)
+            and len(self.peft_config) == 1
+        ):
+            self._dispatch_accelerate_model(
+                device_map=device_map, max_memory=max_memory, offload_dir=offload_dir, offload_index=offload_index
+            )
+
+    def set_adapter(self, adapter_name: str) -> None:
+        r"""
+        Sets an adapter to switch easily between multiple adapters.
+        """
+        requires_backends(self.set_adapter, "peft")
+        if not self._hf_peft_config_loaded:
+            raise ValueError("No adapter loaded. Please load an adapter first.")
+        elif adapter_name not in self.peft_config:
+            raise ValueError(
+                f"Adapter with name {adapter_name} not found. Please pass the correct adapter name among {list(self.peft_config.keys())}"
+            )
+
+        from peft.tuners.tuners_utils import BaseTunerLayer
+
+        _adapters_has_been_set = False
+
+        for _, module in self.named_modules():
+            if isinstance(module, BaseTunerLayer):
+                module.active_adapter = adapter_name
+                _adapters_has_been_set = True
+
+        if not _adapters_has_been_set:
+            raise ValueError(
+                "Did not succeeded in setting the adapter. Please make sure you are using a model that supports adapters."
+            )
+
+    @property
+    def current_active_adapter(self) -> str:
+        r"""
+        Gets the current active adapter of the model.
+        """
+        if not is_peft_available():
+            raise ImportError("PEFT is not available. Please install PEFT to use this function: `pip install peft`.")
+
+        if not self._hf_peft_config_loaded:
+            raise ValueError("No adapter loaded. Please load an adapter first.")
+
+        from peft.tuners.tuners_utils import BaseTunerLayer
+
+        for _, module in self.named_modules():
+            if isinstance(module, BaseTunerLayer):
+                return module.active_adapter
+
+    def _dispatch_accelerate_model(
+        self,
+        device_map: str,
+        max_memory: Optional[int] = None,
+        offload_dir: Optional[str] = None,
+        offload_index: Optional[int] = None,
+    ) -> None:
+        r"""
+        Optionnal re-dispatch the model and attach new hooks to the model in case the model has been loaded with
+        accelerate (i.e. with `device_map=xxx`)
+
+        Args:
+            device_map (`str`):
+                The device map used to load the model with accelerate.
+            max_memory (`int`, `optional`):
+                The maximum memory argument to be passed to `accelerate.get_balanced_memory` method.
+            offload_dir (`str`, `optional`):
+                The offload_dir argument to be passed to `accelerate.dispatch_model` method.
+            offload_index (`int`, `optional`):
+                The offload_index argument to be passed to `accelerate.dispatch_model` method.
+        """
+        dispatch_model_kwargs = {}
+        # Safety checker for previous `accelerate` versions
+        # `offload_index` was introduced in https://github.com/huggingface/accelerate/pull/873/
+        if "offload_index" in inspect.signature(dispatch_model).parameters:
+            dispatch_model_kwargs["offload_index"] = offload_index
+
+        no_split_module_classes = self._no_split_modules
+
+        if device_map != "sequential":
+            max_memory = get_balanced_memory(
+                self,
+                max_memory=max_memory,
+                no_split_module_classes=no_split_module_classes,
+                low_zero=(device_map == "balanced_low_0"),
+            )
+        if isinstance(device_map, str):
+            device_map = infer_auto_device_map(
+                self, max_memory=max_memory, no_split_module_classes=no_split_module_classes
+            )
+        dispatch_model(
+            self,
+            device_map=device_map,
+            offload_dir=offload_dir,
+            **dispatch_model_kwargs,
+        )

src/transformers/modeling_utils.py

@@ -34,6 +34,7 @@
 from torch.nn import CrossEntropyLoss
 
 from .activations import get_activation
+from .adapters import PeftAdapterMixin
 from .configuration_utils import PretrainedConfig
 from .deepspeed import deepspeed_config, is_deepspeed_zero3_enabled
 from .dynamic_module_utils import custom_object_save
@@ -67,6 +68,7 @@
     is_bitsandbytes_available,
     is_offline_mode,
     is_optimum_available,
+    is_peft_available,
     is_remote_url,
     is_safetensors_available,
     is_torch_tpu_available,
@@ -113,6 +115,9 @@
 else:
     IS_SAGEMAKER_MP_POST_1_10 = False
 
+if is_peft_available():
+    from .utils import find_adapter_config_file
+
 
 @contextmanager
 def no_init_weights(_enable=True):
@@ -1025,7 +1030,7 @@ def floating_point_ops(
         return 6 * self.estimate_tokens(input_dict) * self.num_parameters(exclude_embeddings=exclude_embeddings)
 
 
-class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin):
+class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin, PeftAdapterMixin):
     r"""
     Base class for all models.
 
@@ -2211,6 +2216,7 @@ def from_pretrained(
         subfolder = kwargs.pop("subfolder", "")
         commit_hash = kwargs.pop("_commit_hash", None)
         variant = kwargs.pop("variant", None)
+        _adapter_model_path = kwargs.pop("_adapter_model_path", None)
 
         if use_auth_token is not None:
             warnings.warn(
@@ -2236,6 +2242,29 @@ def from_pretrained(
                 " ignored."
             )
 
+        if is_peft_available() and _adapter_model_path is None:
+            maybe_adapter_model_path = find_adapter_config_file(
+                pretrained_model_name_or_path,
+                revision=revision,
+                subfolder=subfolder,
+                use_auth_token=use_auth_token,
+                commit_hash=commit_hash,
+            )
+        elif is_peft_available() and _adapter_model_path is not None:
+            maybe_adapter_model_path = _adapter_model_path
+        else:
+            maybe_adapter_model_path = None
+
+        has_adapter_config = maybe_adapter_model_path is not None
+
+        if has_adapter_config:
+            if _adapter_model_path is not None:
+                adapter_model_id = _adapter_model_path
+            else:
+                with open(maybe_adapter_model_path, "r", encoding="utf-8") as f:
+                    adapter_model_id = pretrained_model_name_or_path
+                    pretrained_model_name_or_path = json.load(f)["base_model_name_or_path"]
+
         # change device_map into a map if we passed an int, a str or a torch.device
         if isinstance(device_map, torch.device):
             device_map = {"": device_map}
@@ -2981,6 +3010,15 @@ def from_pretrained(
                 kwargs["skip_keys"] = model._skip_keys_device_placement
             dispatch_model(model, **kwargs)
 
+        if has_adapter_config:
+            model.load_adapter(
+                adapter_model_id,
+                adapter_name="default",
+                revision=revision,
+                use_auth_token=use_auth_token,
+                commit_hash=commit_hash,
+            )
+
         if output_loading_info:
             if loading_info is None:
                 loading_info = {

src/transformers/models/auto/auto_factory.py

@@ -15,13 +15,14 @@
 """Factory function to build auto-model classes."""
 import copy
 import importlib
+import json
 import os
 import warnings
 from collections import OrderedDict
 
 from ...configuration_utils import PretrainedConfig
 from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
-from ...utils import copy_func, logging, requires_backends
+from ...utils import copy_func, find_adapter_config_file, is_peft_available, logging, requires_backends
 from .configuration_auto import AutoConfig, model_type_to_module_name, replace_list_option_in_docstrings
 
 
@@ -469,6 +470,24 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
         if token is not None:
             hub_kwargs["token"] = token
 
+        if is_peft_available():
+            revision = kwargs.get("revision", None)
+            subfolder = kwargs.get("subfolder", None)
+
+            maybe_adapter_path = find_adapter_config_file(
+                pretrained_model_name_or_path,
+                revision=revision,
+                use_auth_token=use_auth_token,
+                subfolder=subfolder,
+            )
+
+            if maybe_adapter_path is not None:
+                with open(maybe_adapter_path, "r") as f:
+                    adapter_config = json.load(f)
+
+                    kwargs["_adapter_model_path"] = pretrained_model_name_or_path
+                    pretrained_model_name_or_path = adapter_config["base_model_name_or_path"]
+
         if not isinstance(config, PretrainedConfig):
             kwargs_orig = copy.deepcopy(kwargs)
             # ensure not to pollute the config object with torch_dtype="auto" - since it's

src/transformers/utils/__init__.py

@@ -177,13 +177,11 @@
     requires_backends,
     torch_only_method,
 )
+from .peft_utils import ADAPTER_CONFIG_NAME, ADAPTER_SAFE_WEIGHTS_NAME, ADAPTER_WEIGHTS_NAME, find_adapter_config_file
 
 
 WEIGHTS_NAME = "pytorch_model.bin"
 WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json"
-ADAPTER_CONFIG_NAME = "adapter_config.json"
-ADAPTER_WEIGHTS_NAME = "adapter_model.bin"
-ADAPTER_SAFE_WEIGHTS_NAME = "adapter_model.safetensors"
 TF2_WEIGHTS_NAME = "tf_model.h5"
 TF2_WEIGHTS_INDEX_NAME = "tf_model.h5.index.json"
 TF_WEIGHTS_NAME = "model.ckpt"