adapt for other models

src/peft/mapping.py

@@ -138,4 +138,6 @@ def get_peft_model(model, peft_config):
     else:
         peft_config = _prepare_lora_config(peft_config, model_config)
 
+    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)
+
     return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config)

src/peft/peft_model.py

@@ -14,18 +14,31 @@
 # limitations under the License.
 
 import inspect
+import os
 import warnings
 
 import torch
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput
+from transformers.utils import PushToHubMixin
+
+from huggingface_hub import hf_hub_download
 
 from .tuners import LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder
-from .utils import PeftConfig, PeftType, TaskType, _set_trainable, shift_tokens_right
+from .utils import (
+    WEIGHTS_NAME,
+    PeftConfig,
+    PeftType,
+    TaskType,
+    _set_trainable,
+    get_peft_model_state_dict,
+    set_peft_model_state_dict,
+    shift_tokens_right,
+)
 
 
-class PeftModel(torch.nn.Module):
+class PeftModel(PushToHubMixin, torch.nn.Module):
     """
     Parameter-Efficient Fine-Tuning Model. Base model encompassing various Peft methods.
 
@@ -61,6 +74,73 @@ class PeftModel(torch.nn.Module):
             self.base_model = LoraModel(peft_config, model)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+    def save_pretrained(self, save_directory, **kwargs):
+        r"""
+        Args:
+        This function saves the adapter model and the adapter configuration files to a directory, so that it can be
+        re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub`
+        method.
+            save_directory (`str`):
+                Directory where the adapter model and configuration files will be saved (will be created if it does not
+                exist).
+            **kwargs:
+                Additional keyword arguments passed along to the `push_to_hub` method.
+        """
+        if os.path.isfile(save_directory):
+            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save the config
+        if self.peft_config.base_model_name_or_path is None:
+            self.peft_config.base_model_name_or_path = self.base_model.__dict__.get("name_or_path", None)
+        self.peft_config.inference_mode = True
+        self.peft_config.save_pretrained(save_directory)
+
+        for param in self.parameters():
+            param.requires_grad = False  # freeze the model
+
+        # save only the trainable weights
+        output_state_dict = get_peft_model_state_dict(self, kwargs.get("state_dict", None))
+        torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
+
+    @classmethod
+    def from_pretrained(cls, model, model_id, **kwargs):
+        r"""
+        Args:
+        Instantiate a `LoraModel` from a pretrained Lora configuration and weights.
+            model (`transformers.PreTrainedModel`):
+                The model to be adapted. The model should be initialized with the `from_pretrained` method. from
+                `transformers` library.
+            model_id (`str`):
+                The name of the Lora configuration to use. Can be either:
+                    - A string, the `model id` of a Lora configuration hosted inside a model repo on
+                        huggingface Hub
+                    - A path to a directory containing a Lora configuration file saved using the
+                        `save_pretrained` method, e.g., ``./my_lora_config_directory/``.
+        """
+        from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PEFT_TYPE_TO_CONFIG_MAPPING
+
+        # load the config
+        config = PEFT_TYPE_TO_CONFIG_MAPPING[PeftConfig.from_pretrained(model_id).peft_type].from_pretrained(model_id)
+
+        model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config)
+
+        # load weights if any
+        if os.path.exists(os.path.join(model_id, WEIGHTS_NAME)):
+            filename = os.path.join(model_id, WEIGHTS_NAME)
+        else:
+            try:
+                filename = hf_hub_download(model_id, WEIGHTS_NAME)
+            except:  # noqa
+                raise ValueError(
+                    f"Can't find weights for {model_id} in {model_id} or in the Hugging Face Hub. "
+                    f"Please check that the file {WEIGHTS_NAME} is present at {model_id}."
+                )
+
+        adapters_weights = torch.load(filename)
+        # load the weights into the model
+        return set_peft_model_state_dict(model, adapters_weights)
+
     def _setup_prompt_encoder(self):
         num_transformer_submodules = 0
         transformer_backbone = None

src/peft/tuners/lora.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 import importlib
 import math
-import os
 import warnings
 from dataclasses import asdict, dataclass, field
 from enum import Enum
@@ -24,12 +23,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers.pytorch_utils import Conv1D
-from transformers.utils import PushToHubMixin
 
 import bitsandbytes as bnb
-from huggingface_hub import hf_hub_download
 
-from ..utils import WEIGHTS_NAME, PeftConfig, PeftType, get_peft_model_state_dict, transpose
+from ..utils import PeftConfig, PeftType, transpose
 
 
 def is_loralib_available():
@@ -76,7 +73,7 @@ class LoraConfig(PeftConfig):
         self.peft_type = PeftType.LORA
 
 
-class LoraModel(PushToHubMixin, torch.nn.Module):
+class LoraModel(torch.nn.Module):
     """
     Creates Low Rank Adapter (Lora) model from a pretrained transformers model.
 
@@ -165,72 +162,6 @@ class LoraModel(PushToHubMixin, torch.nn.Module):
             new_module.state = old_module.state
             new_module.to(old_module.weight.device)
 
-    def save_pretrained(self, save_directory, **kwargs):
-        r"""
-        This function saves the adapter model and the adapter configuration files to a directory, so that it can be
-        re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub`
-        method.
-
-        Args:
-            save_directory (`str`):
-                Directory where the adapter model and configuration files will be saved (will be created if it does not
-                exist).
-            **kwargs:
-                Additional keyword arguments passed along to the `push_to_hub` method.
-        """
-        if os.path.isfile(save_directory):
-            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
-        os.makedirs(save_directory, exist_ok=True)
-
-        # save the config
-        self.peft_config.save_pretrained(save_directory)
-
-        for param in self.parameters():
-            param.requires_grad = False  # freeze the model
-
-        # save only the trainable weights
-        output_state_dict = get_peft_model_state_dict(self)
-        torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME))
-
-    @classmethod
-    def from_pretrained(cls, model, lora_id, **kwargs):
-        r"""
-        Instantiate a `LoraModel` from a pretrained Lora configuration and weights.
-
-        Args:
-            model (`transformers.PreTrainedModel`):
-                The model to be adapted. The model should be initialized with the `from_pretrained` method. from
-                `transformers` library.
-            lora_id (`str`):
-                The name of the Lora configuration to use. Can be either:
-                    - A string, the `model id` of a Lora configuration hosted inside a model repo on
-                        huggingface Hub
-                    - A path to a directory containing a Lora configuration file saved using the
-                        `save_pretrained` method, e.g., ``./my_lora_config_directory/``.
-        """
-        # load the config
-        config = LoraConfig.from_pretrained(lora_id)
-
-        model = cls(config, model)
-
-        # load weights if any
-        if os.path.exists(os.path.join(lora_id, WEIGHTS_NAME)):
-            filename = os.path.join(lora_id, WEIGHTS_NAME)
-        else:
-            try:
-                filename = hf_hub_download(lora_id, WEIGHTS_NAME)
-            except:  # noqa
-                raise ValueError(
-                    f"Can't find weights for {lora_id} in {lora_id} or in the Hugging Face Hub. "
-                    f"Please check that the file {WEIGHTS_NAME} is present at {lora_id}."
-                )
-
-        adapters_weights = torch.load(filename)
-        # load the weights into the model
-        model.load_state_dict(adapters_weights, strict=False)
-
-        return model
-
     def __getattr__(self, name: str):
         """Forward missing attributes to the wrapped module."""
         try:

tests/test_lora.py

@@ -17,36 +17,62 @@ import torch
 import tempfile
 import unittest
 
-from peft import LoraModel, LoraConfig, get_peft_model_state_dict
+from peft import PeftConfig, PeftModel, LoraConfig, get_peft_model_state_dict, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
 
 from transformers import AutoModelForCausalLM
 
 class LoraTestMixin:
     checkpoints_to_test = [
-        "trl-internal-testing/tiny-random-OPTForCausalLM",
+        "hf-internal-testing/tiny-random-OPTForCausalLM",
     ]
+    config_classes = (
+        LoraConfig,
+        # PrefixTuningConfig,
+        # PromptEncoderConfig,
+        # PromptTuningConfig,
+    )
+    config_kwargs = (
+        dict(
+            r = 8,
+            lora_alpha=32,
+            target_modules=["q_proj", "v_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+        ),
+        # dict(
+        #     encoder_hidden_size=32,
+        #     task_type="CAUSAL_LM",
+        # ),
+        # dict(
+        #     encoder_hidden_size=32,
+        #     task_type="CAUSAL_LM",
+        # ),
+        # dict(
+        #     task_type="CAUSAL_LM",
+        # )
 
-class LoraTester(unittest.TestCase, LoraTestMixin):
+    )
+
+class PeftModelTester(unittest.TestCase, LoraTestMixin):
     r"""
-    Test if the LoraModel behaves as expected. This includes:
+    Test if the PeftModel behaves as expected. This includes:
     - test if the model has the expected methods
     """
     def test_attributes_lora_model(self):
         for model_id in self.checkpoints_to_test:
             model = AutoModelForCausalLM.from_pretrained(model_id)
             
-            config = LoraConfig(
-                r = 8,
-                lora_alpha=32,
-                target_modules=["q_proj", "v_proj"],
-                lora_dropout=0.05,
-                bias="none",
-            )
-            model = LoraModel(config, model)
+            for i, config_cls in enumerate(self.config_classes):
+                config = config_cls(
+                    base_model_name_or_path=model_id,
+                    **self.config_kwargs[i],
+                )
+                model = PeftModel(model, config)
 
-            self.assertTrue(hasattr(model, 'save_pretrained'))
-            self.assertTrue(hasattr(model, 'from_pretrained'))
-            self.assertTrue(hasattr(model, 'push_to_hub'))
+                self.assertTrue(hasattr(model, 'save_pretrained'))
+                self.assertTrue(hasattr(model, 'from_pretrained'))
+                self.assertTrue(hasattr(model, 'push_to_hub'))
     
     def test_save_pretrained(self):
         r"""
@@ -62,42 +88,38 @@ class LoraTester(unittest.TestCase, LoraTestMixin):
         for model_id in self.checkpoints_to_test:
             model = AutoModelForCausalLM.from_pretrained(model_id)
             
-            config = LoraConfig(
-                r = 8,
-                lora_alpha=32,
-                target_modules=["q_proj", "v_proj"],
-                lora_dropout=0.05,
-                bias="none",
-            )
-            model = LoraModel(config, model)
+            for i, config_cls in enumerate(self.config_classes):
+                config = config_cls(
+                    base_model_name_or_path=model_id,
+                    **self.config_kwargs[i],
+                )
+                model = PeftModel(model, config)
 
-            with tempfile.TemporaryDirectory() as tmp_dirname:
-                model.save_pretrained(tmp_dirname)
+                with tempfile.TemporaryDirectory() as tmp_dirname:
+                    model.save_pretrained(tmp_dirname)
 
-                model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id)
-                model_from_pretrained = LoraModel.from_pretrained(model_from_pretrained, tmp_dirname)
-                
-                # check if the state dicts are equal
-                state_dict = get_peft_model_state_dict(model)
-                state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained)
+                    model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id)
+                    model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname)
+                    
+                    # check if the state dicts are equal
+                    state_dict = get_peft_model_state_dict(model)
+                    state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained)
 
-                # check if same keys
-                self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys())
+                    # check if same keys
+                    self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys())
 
-                # check if tensors equal
-                for key in state_dict.keys():
-                    self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key]))
+                    # check if tensors equal
+                    for key in state_dict.keys():
+                        self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key]))
 
-                # check if `adapter_model.bin` is present
-                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin")))
+                    # check if `adapter_model.bin` is present
+                    self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin")))
 
-                # check if `adapter_config.json` is present
-                self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json")))
-
-                # check if `pytorch_model.bin` is not present
-                self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin")))
-
-                # check if `config.json` is not present
-                self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))
+                    # check if `adapter_config.json` is present
+                    self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json")))
 
+                    # check if `pytorch_model.bin` is not present
+                    self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin")))
 
+                    # check if `config.json` is not present
+                    self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))