[core] PEFT refactor + introducing inject_adapter_in_model public method

docs/source/package_reference/config.mdx

@@ -4,7 +4,7 @@ The configuration classes stores the configuration of a [`PeftModel`], PEFT adap
 
 ## PeftConfigMixin
 
-[[autodoc]] utils.config.PeftConfigMixin
+[[autodoc]] config.PeftConfigMixin
     - all
 
 ## PeftConfig

src/peft/__init__.py

@@ -28,7 +28,13 @@
     AutoPeftModelForQuestionAnswering,
     AutoPeftModelForFeatureExtraction,
 )
-from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PEFT_TYPE_TO_CONFIG_MAPPING, get_peft_config, get_peft_model
+from .mapping import (
+    MODEL_TYPE_TO_PEFT_MODEL_MAPPING,
+    PEFT_TYPE_TO_CONFIG_MAPPING,
+    get_peft_config,
+    get_peft_model,
+    inject_adapter_in_model,
+)
 from .peft_model import (
     PeftModel,
     PeftModelForCausalLM,
@@ -58,14 +64,14 @@
 )
 from .utils import (
     TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING,
-    PeftConfig,
     PeftType,
-    PromptLearningConfig,
     TaskType,
     bloom_model_postprocess_past_key_value,
     get_peft_model_state_dict,
     prepare_model_for_int8_training,
     prepare_model_for_kbit_training,
     set_peft_model_state_dict,
     shift_tokens_right,
+    load_peft_weights,
 )
+from .config import PeftConfig, PromptLearningConfig

src/peft/auto.py

@@ -27,6 +27,7 @@
     AutoModelForTokenClassification,
 )
 
+from .config import PeftConfig
 from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING
 from .peft_model import (
     PeftModel,
@@ -37,7 +38,6 @@
     PeftModelForSequenceClassification,
     PeftModelForTokenClassification,
 )
-from .utils import PeftConfig
 
 
 class _BaseAutoPeftModel:

src/peft/utils/config.py → src/peft/config.py

@@ -12,7 +12,6 @@
 # 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 enum
 import inspect
 import json
 import os
@@ -22,26 +21,7 @@
 from huggingface_hub import hf_hub_download
 from transformers.utils import PushToHubMixin
 
-from .other import CONFIG_NAME
-
-
-class PeftType(str, enum.Enum):
-    PROMPT_TUNING = "PROMPT_TUNING"
-    P_TUNING = "P_TUNING"
-    PREFIX_TUNING = "PREFIX_TUNING"
-    LORA = "LORA"
-    ADALORA = "ADALORA"
-    ADAPTION_PROMPT = "ADAPTION_PROMPT"
-    IA3 = "IA3"
-
-
-class TaskType(str, enum.Enum):
-    SEQ_CLS = "SEQ_CLS"
-    SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM"
-    CAUSAL_LM = "CAUSAL_LM"
-    TOKEN_CLS = "TOKEN_CLS"
-    QUESTION_ANS = "QUESTION_ANS"
-    FEATURE_EXTRACTION = "FEATURE_EXTRACTION"
+from .utils import CONFIG_NAME, PeftType, TaskType
 
 
 @dataclass
@@ -102,6 +82,9 @@ def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None, **kwargs
             kwargs (additional keyword arguments, *optional*):
                 Additional keyword arguments passed along to the child class initialization.
         """
+        # Avoid circular dependency .. TODO: fix this with a larger refactor
+        from peft.mapping import PEFT_TYPE_TO_CONFIG_MAPPING
+
         path = (
             os.path.join(pretrained_model_name_or_path, subfolder)
             if subfolder is not None
@@ -122,7 +105,27 @@ def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None, **kwargs
 
         loaded_attributes = cls.from_json_file(config_file)
 
-        config = cls(**class_kwargs)
+        # TODO: this hack is needed to fix the following issue (on commit 702f937):
+        # if someone saves a default config and loads it back with `PeftConfig` class it yields to
+        # not loading the correct config class.
+
+        # from peft import AdaLoraConfig, PeftConfig
+        # peft_config = AdaLoraConfig()
+        # print(peft_config)
+        # >>> AdaLoraConfig(peft_type=<PeftType.ADALORA: 'ADALORA'>, auto_mapping=None, base_model_name_or_path=None,
+        # revision=None, task_type=None, inference_mode=False, r=8, target_modules=None, lora_alpha=8, lora_dropout=0.0, ...
+        #
+        # peft_config.save_pretrained("./test_config")
+        # peft_config = PeftConfig.from_pretrained("./test_config")
+        # print(peft_config)
+        # >>> PeftConfig(peft_type='ADALORA', auto_mapping=None, base_model_name_or_path=None, revision=None, task_type=None, inference_mode=False)
+        if "peft_type" in loaded_attributes:
+            peft_type = loaded_attributes["peft_type"]
+            config_cls = PEFT_TYPE_TO_CONFIG_MAPPING[peft_type]
+        else:
+            config_cls = cls
+
+        config = config_cls(**class_kwargs)
 
         for key, value in loaded_attributes.items():
             if hasattr(config, key):
@@ -185,6 +188,18 @@ def _get_peft_type(
         loaded_attributes = cls.from_json_file(config_file)
         return loaded_attributes["peft_type"]
 
+    @property
+    def is_prompt_learning(self):
+        r"""
+        Utility method to check if the configuration is for prompt learning.
+        """
+        return False
+
+    @property
+    def is_adaption_prompt(self) -> bool:
+        """Return True if this is an adaption prompt config."""
+        return False
+
 
 @dataclass
 class PeftConfig(PeftConfigMixin):
@@ -227,3 +242,10 @@ class PromptLearningConfig(PeftConfig):
     )
     num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"})
     num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})
+
+    @property
+    def is_prompt_learning(self):
+        r"""
+        Utility method to check if the configuration is for prompt learning.
+        """
+        return True

src/peft/mapping.py

@@ -17,6 +17,9 @@
 
 from typing import TYPE_CHECKING, Any, Dict
 
+import torch
+
+from .config import PeftConfig
 from .peft_model import (
     PeftModel,
     PeftModelForCausalLM,
@@ -28,21 +31,22 @@
 )
 from .tuners import (
     AdaLoraConfig,
+    AdaLoraModel,
     AdaptionPromptConfig,
     IA3Config,
+    IA3Model,
     LoraConfig,
+    LoraModel,
     PrefixTuningConfig,
     PromptEncoderConfig,
     PromptTuningConfig,
 )
-from .utils import PromptLearningConfig, _prepare_prompt_learning_config
+from .utils import _prepare_prompt_learning_config
 
 
 if TYPE_CHECKING:
     from transformers import PreTrainedModel
 
-    from .utils.config import PeftConfig
-
 
 MODEL_TYPE_TO_PEFT_MODEL_MAPPING = {
     "SEQ_CLS": PeftModelForSequenceClassification,
@@ -63,6 +67,12 @@
     "IA3": IA3Config,
 }
 
+PEFT_TYPE_TO_TUNER_MAPPING = {
+    "LORA": LoraModel,
+    "ADALORA": AdaLoraModel,
+    "IA3": IA3Model,
+}
+
 
 def get_peft_config(config_dict: Dict[str, Any]):
     """
@@ -89,10 +99,38 @@ def get_peft_model(model: PreTrainedModel, peft_config: PeftConfig, adapter_name
 
     peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)
 
-    if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not isinstance(
-        peft_config, PromptLearningConfig
-    ):
+    if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not peft_config.is_prompt_learning:
         return PeftModel(model, peft_config, adapter_name=adapter_name)
-    if isinstance(peft_config, PromptLearningConfig):
+    if peft_config.is_prompt_learning:
         peft_config = _prepare_prompt_learning_config(peft_config, model_config)
     return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config, adapter_name=adapter_name)
+
+
+def inject_adapter_in_model(peft_config: PeftConfig, model: torch.nn.Module, adapter_name: str):
+    r"""
+    A simple API to create and inject adapter in-place into a model. Currently the API does not support prompt learning
+    methods and adaption prompt. Make sure to have the correct `target_names` set in the `peft_config` object. The API
+    calls `get_peft_model` under the hood but would be restricted only to non-prompt learning methods.
+
+    Args:
+        peft_config (`PeftConfig`):
+            Configuration object containing the parameters of the Peft model.
+        model (`torch.nn.Module`):
+            The input model where the adapter will be injected.
+        adapter_name (`str`):
+            The name of the adapter to be injected.
+    """
+    if peft_config.is_prompt_learning or peft_config.is_adaption_prompt:
+        raise ValueError("`create_and_replace` does not support prompt learning and adaption prompt yet.")
+
+    if peft_config.peft_type not in PEFT_TYPE_TO_TUNER_MAPPING.keys():
+        raise ValueError(
+            f"`inject_adapter_in_model` does not support {peft_config.peft_type} yet. Please use `get_peft_model`."
+        )
+
+    tuner_cls = PEFT_TYPE_TO_TUNER_MAPPING[peft_config.peft_type]
+
+    # By instantiating a peft model we are injecting randomly initialized LoRA layers into the model's modules.
+    peft_model = tuner_cls(model, peft_config, adapter_name=adapter_name)
+
+    return peft_model.model