[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,
+    create_and_replace,
+)
 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,13 @@ def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None, **kwargs
 
         loaded_attributes = cls.from_json_file(config_file)
 
-        config = cls(**class_kwargs)
+        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 +174,13 @@ 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
+
 
 @dataclass
 class PeftConfig(PeftConfigMixin):
@@ -227,3 +223,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,7 @@
 
 from typing import TYPE_CHECKING, Any, Dict
 
+from .config import PeftConfig
 from .peft_model import (
     PeftModel,
     PeftModelForCausalLM,
@@ -27,6 +28,7 @@
     PeftModelForTokenClassification,
 )
 from .tuners import (
+    TUNERS_MAPPING,
     AdaLoraConfig,
     AdaptionPromptConfig,
     IA3Config,
@@ -35,14 +37,12 @@
     PromptEncoderConfig,
     PromptTuningConfig,
 )
-from .utils import PromptLearningConfig, _prepare_prompt_learning_config
+from .utils import _get_submodules, _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,
@@ -89,10 +89,51 @@ 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)
+
+
+# TODO: docstring and typehints
+def create_and_replace(peft_config, model, adapter_name):
+    if not isinstance(peft_config, PeftConfig):
+        raise ValueError(f"peft_config must be an instance of PeftConfig got {type(peft_config)} instead.")
+
+    peft_type = peft_config.peft_type
+
+    if peft_type not in TUNERS_MAPPING:
+        raise ValueError(
+            f"Task type {peft_type} is not supported. Supported task types are {list(TUNERS_MAPPING.keys())}"
+        )
+    tuner_cls = TUNERS_MAPPING[peft_type]
+
+    # TODO: test that
+    for module in model.modules():
+        if not getattr(module, "_is_peft_tuner_layer", False):
+            module.requires_grad_(False)
+
+    is_target_modules_in_base_model = False
+    key_list = [key for key, _ in model.named_modules()]
+
+    for key in key_list:
+        if not tuner_cls._check_target_module_exists(peft_config, key):
+            continue
+
+        is_target_modules_in_base_model = True
+        parent, target, target_name = _get_submodules(model, key)
+
+        optionnal_kwargs = {
+            "loaded_in_8bit": getattr(model, "is_loaded_in_8bit", False),
+            "loaded_in_4bit": getattr(model, "is_loaded_in_4bit", False),
+            "current_key": key,
+        }
+
+        tuner_cls.create_and_replace(peft_config, adapter_name, target, target_name, parent, **optionnal_kwargs)
+
+    if not is_target_modules_in_base_model:
+        raise ValueError(
+            f"Target modules {peft_config.target_modules} not found in the base model. "
+            f"Please check the target modules and try again."
+        )