[WIP] ENH: Adapter injection based on state_dict

src/peft/mapping.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 
@@ -45,7 +45,11 @@ def get_peft_config(config_dict: dict[str, Any]) -> PeftConfig:
 
 
 def inject_adapter_in_model(
-    peft_config: PeftConfig, model: torch.nn.Module, adapter_name: str = "default", low_cpu_mem_usage: bool = False
+    peft_config: PeftConfig,
+    model: torch.nn.Module,
+    adapter_name: str = "default",
+    low_cpu_mem_usage: bool = False,
+    state_dict: Optional[dict[str, torch.Tensor]] = None,
 ) -> torch.nn.Module:
     r"""
     A simple API to create and inject adapter in-place into a model. Currently the API does not support prompt learning
@@ -61,6 +65,8 @@ def inject_adapter_in_model(
             The name of the adapter to be injected, if not provided, the default adapter name is used ("default").
         low_cpu_mem_usage (`bool`, `optional`, defaults to `False`):
             Create empty adapter weights on meta device. Useful to speed up the loading process.
+
+        FIXME
     """
     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.")
@@ -73,6 +79,8 @@ def inject_adapter_in_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, low_cpu_mem_usage=low_cpu_mem_usage)
+    peft_model = tuner_cls(
+        model, peft_config, adapter_name=adapter_name, low_cpu_mem_usage=low_cpu_mem_usage, state_dict=state_dict
+    )
 
     return peft_model.model

src/peft/tuners/lora/model.py

@@ -139,8 +139,9 @@ class LoraModel(BaseTuner):
 
     prefix: str = "lora_"
 
-    def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False) -> None:
-        super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage)
+    def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False, **kwargs) -> None:
+        # FIXME add kwargs to other model classes
+        super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage, **kwargs)
 
     def _check_new_adapter_config(self, config: LoraConfig) -> None:
         """

src/peft/tuners/tuners_utils.py

@@ -154,6 +154,8 @@ class BaseTuner(nn.Module, ABC):
 
     The easiest is to check what is done in the `peft.tuners.lora.LoraModel` class.
 
+    FIXME
+
     Attributes:
         model (`torch.nn.Module`):
             The model to which the adapter tuner layers will be attached.
@@ -176,6 +178,7 @@ def __init__(
         peft_config: Union[PeftConfig, dict[str, PeftConfig]],
         adapter_name: str,
         low_cpu_mem_usage: bool = False,
+        state_dict: Optional[dict[str, torch.Tensor]] = None,
     ) -> None:
         super().__init__()
 
@@ -200,7 +203,7 @@ def __init__(
         self.active_adapter: str | list[str] = adapter_name
         self._pre_injection_hook(self.model, self.peft_config[adapter_name], adapter_name)
         if peft_config != PeftType.XLORA or peft_config[adapter_name] != PeftType.XLORA:
-            self.inject_adapter(self.model, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage)
+            self.inject_adapter(self.model, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage, state_dict=state_dict)
 
         # Copy the peft_config in the injected model.
         self.model.peft_config = self.peft_config
@@ -427,14 +430,21 @@ def _check_target_module_compatiblity(self, peft_config: PeftConfig, model: nn.M
         _check_lora_target_modules_mamba(peft_config, model, target_name)
 
     def inject_adapter(
-        self, model: nn.Module, adapter_name: str, autocast_adapter_dtype: bool = True, low_cpu_mem_usage: bool = False
+        self,
+        model: nn.Module,
+        adapter_name: str,
+        autocast_adapter_dtype: bool = True,
+        low_cpu_mem_usage: bool = False,
+        state_dict: Optional[dict[str, torch.Tensor]] = None,
     ) -> None:
         r"""
         Creates adapter layers and replaces the target modules with the adapter layers. This method is called under the
         hood by `peft.mapping.get_peft_model` if a non-prompt tuning adapter class is passed.
 
         The corresponding PEFT config is directly retrieved from the `peft_config` attribute of the BaseTuner class.
 
+        FIXME
+
         Args:
             model (`nn.Module`):
                 The model to be tuned.
@@ -444,11 +454,17 @@ def inject_adapter(
                 Whether to autocast the adapter dtype. Defaults to `True`.
             low_cpu_mem_usage (`bool`, `optional`, defaults to `False`):
                 Create empty adapter weights on meta device. Useful to speed up the loading process.
+            FIXME
 
         """
+        ###################################
+        # PREPARATION OF MODEL AND CONFIG #
+        ###################################
+
         peft_config = self.peft_config[adapter_name]
         excluded_modules = []
         unmatched_modules = []
+        targeted_modules_from_peft_config: list[str] = []  # only relevant if state_dict is passed
         # Note: If possible, all checks should be performed *at the start of this method*.
         # This way, we can raise early if something goes wrong, without leaving the model
         # in a bad (half-initialized) state.
@@ -498,11 +514,18 @@ def inject_adapter(
             if len(new_target_modules) < len(peft_config.target_modules):
                 peft_config.target_modules = new_target_modules
 
+        ###############################
+        # MATCHING & CREATING MODULES #
+        ###############################
+
         existing_adapter_map = {}
         for key, module in named_modules:
             if isinstance(module, BaseTunerLayer):
                 existing_adapter_map[key] = module
 
+        # TODO: check if this the most robust way
+        state_dict_keys = {k.rsplit(".", 2)[0] for k in state_dict} if state_dict is not None else set()
+
         for key, module in named_modules:
             if not key:
                 continue
@@ -517,18 +540,67 @@ def inject_adapter(
             if excluded_modules and excluded_modules[-1] == key:
                 continue
 
-            result = self._check_target_module_exists(peft_config, key)
-            if isinstance(result, _ExcludedModule):
-                excluded_modules.append(key)
-            elif not result:
-                unmatched_modules.append(key)
+            if state_dict is None:
+                # normal mechanism: match the modules using the peft_config
+                result = self._check_target_module_exists(peft_config, key)
+                if isinstance(result, _ExcludedModule):
+                    excluded_modules.append(key)
+                elif not result:
+                    unmatched_modules.append(key)
+                else:
+                    self.targeted_module_names.append(key)
+                    parent, target, target_name = _get_submodules(model, key)
+                    self._check_target_module_compatiblity(peft_config, model, target_name)
+                    ctx = init_empty_weights if low_cpu_mem_usage else nullcontext
+                    with ctx():
+                        self._create_and_replace(
+                            peft_config, adapter_name, target, target_name, parent, current_key=key
+                        )
             else:
-                self.targeted_module_names.append(key)
-                parent, target, target_name = _get_submodules(model, key)
-                self._check_target_module_compatiblity(peft_config, model, target_name)
-                ctx = init_empty_weights if low_cpu_mem_usage else nullcontext
-                with ctx():
-                    self._create_and_replace(peft_config, adapter_name, target, target_name, parent, current_key=key)
+                # use the state_dict to match modules instead
+                if key not in state_dict_keys:
+                    unmatched_modules.append(key)
+                else:
+                    self.targeted_module_names.append(key)
+                    parent, target, target_name = _get_submodules(model, key)
+                    self._check_target_module_compatiblity(peft_config, model, target_name)
+                    ctx = init_empty_weights if low_cpu_mem_usage else nullcontext
+                    with ctx():
+                        self._create_and_replace(
+                            peft_config, adapter_name, target, target_name, parent, current_key=key
+                        )
+
+                # still record what would have been matched via the config so that the two results can be compared
+                if self._check_target_module_exists(peft_config, key):
+                    targeted_modules_from_peft_config.append(key)
+                    self.targeted_module_names.append(key)
+
+        ####################
+        # CHECK FOR ERRORS #
+        ####################
+
+        if state_dict is not None:
+            # in case that the state_dict was used as source of truth and it resulted in different outcomes than what
+            # would have been matched with the PEFT config, warn the user about that.
+            targeted_set_from_peft_config = set(targeted_modules_from_peft_config)
+            targeted_set_from_state_dict = set(self.targeted_module_names)
+            diff_peft_config = targeted_set_from_peft_config - targeted_set_from_state_dict
+            diff_state_dict = targeted_set_from_state_dict - targeted_set_from_peft_config
+            error_msg = ""
+            if diff_peft_config or diff_state_dict:
+                error_msg = (
+                    "While injecting the PEFT adapters, an inconsistency was discovered between the PEFT config and "
+                    "the provided state_dict. This is not necessarily an issue and can be ignored if this was the "
+                    "intent. "
+                )
+            if diff_peft_config:
+                error_msg += f"The PEFT config contained these additional target modules: {sorted(diff_peft_config)}. "
+            if diff_state_dict:
+                error_msg += f"The state_dict contained these additional target modules: {sorted(diff_state_dict)}. "
+            if error_msg:
+                # FIXME for debugging purposes, raise here
+                1/0
+                warnings.warn(error_msg)
 
         if not self.targeted_module_names and not uses_dummy_target_modules:
             if excluded_modules and not unmatched_modules:
@@ -578,6 +650,10 @@ def inject_adapter(
                 "See for example https://github.com/huggingface/peft/issues/2018."
             )
 
+        ################
+        # HOUSEKEEPING #
+        ################
+
         # It's important to set the adapter here (again), because otherwise it can happen that if a 2nd adapter is
         # added, and it targets different layer(s) than the first adapter (which is active), then those different
         # layers will be activated, which we don't want.