[PEFT] Peft integration alternative design

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,136 @@
+# 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 os
+from typing import Optional
+
+from ..utils import ADAPTER_CONFIG_NAME, cached_file, logging, requires_backends
+
+
+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
+    """
+
+    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,
+    ):
+        """
+        Load adapter weights from file. Requires peft as a backend to load the adapter weights
+        """
+        requires_backends(self.load_adapter, "peft")
+
+        from peft import LoraConfig, PeftModel, create_and_replace
+        from peft.utils import set_peft_model_state_dict
+        from peft.utils.other import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING
+
+        self.peft_config = {}
+
+        adapter_config_file = self._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."
+            )
+
+        # TODO: automatically infer the correct config class
+        loaded_peft_config = LoraConfig.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
+
+        # TODO: constraint this to single adapter
+        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.peft_type, loaded_peft_config, self, adapter_name)
+
+        # TODO: move that to peft.utils
+        adapter_state_dict = PeftModel._get_peft_state_dict(
+            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}. "
+                )
+
+    def _find_adapter_config_file(
+        self,
+        model_id: str,
+        revision: str = None,
+        use_auth_token: Optional[str] = None,
+        commit_hash: Optional[str] = None,
+    ) -> Optional[str]:
+        r"""
+        Simply checks if the model stored on the Hub or locally is an adapter model or not, return the path the the
+        adapter config file if it is, None otherwise.
+        """
+        adapter_cached_filename = None
+        if os.path.isdir(model_id):
+            list_remote_files = os.listdir(model_id)
+            if ADAPTER_CONFIG_NAME in list_remote_files:
+                adapter_cached_filename = os.path.join(model_id, ADAPTER_CONFIG_NAME)
+        else:
+            adapter_cached_filename = cached_file(
+                model_id,
+                ADAPTER_CONFIG_NAME,
+                revision=revision,
+                use_auth_token=use_auth_token,
+                _commit_hash=commit_hash,
+                _raise_exceptions_for_missing_entries=False,
+                _raise_exceptions_for_connection_errors=False,
+            )
+
+        return adapter_cached_filename

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
@@ -1021,7 +1022,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.
 

src/transformers/utils/import_utils.py

@@ -989,6 +989,11 @@ def is_jieba_available():
 jieba`. Please note that you may need to restart your runtime after installation.
 """
 
+PEFT_IMPORT_ERROR = """
+{0} requires the peft library but it was not found in your environment. You can install it with pip: `pip install
+peft`. Please note that you may need to restart your runtime after installation.
+"""
+
 BACKENDS_MAPPING = OrderedDict(
     [
         ("bs4", (is_bs4_available, BS4_IMPORT_ERROR)),
@@ -1022,6 +1027,7 @@ def is_jieba_available():
         ("decord", (is_decord_available, DECORD_IMPORT_ERROR)),
         ("cython", (is_cython_available, CYTHON_IMPORT_ERROR)),
         ("jieba", (is_jieba_available, JIEBA_IMPORT_ERROR)),
+        ("peft", (is_peft_available, PEFT_IMPORT_ERROR)),
     ]
 )