add modules_to_save to LoraConfig and other fixes

1. Add `modules_to_save` to LoraConfig
2. Using PeftModel for LoraConfig instead of task-specific classes because LoRA is task agnostic.

src/peft/mapping.py

@@ -14,13 +14,14 @@
 # limitations under the License.
 
 from .peft_model import (
+    PeftModel,
     PeftModelForCausalLM,
     PeftModelForSeq2SeqLM,
     PeftModelForSequenceClassification,
     PeftModelForTokenClassification,
 )
 from .tuners import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
-from .utils import PeftType
+from .utils import PromptLearningConfig
 
 
 MODEL_TYPE_TO_PEFT_MODEL_MAPPING = {
@@ -133,11 +134,9 @@ def get_peft_model(model, peft_config):
     """
 
     model_config = model.config.to_dict()
-    if peft_config.peft_type != PeftType.LORA:
-        peft_config = _prepare_prompt_learning_config(peft_config, model_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)
-
+    if not isinstance(peft_config, PromptLearningConfig):
+        peft_config = _prepare_lora_config(peft_config, model_config)
+        return PeftModel(model, peft_config)
+    peft_config = _prepare_prompt_learning_config(peft_config, model_config)
     return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config)

src/peft/peft_model.py

@@ -30,6 +30,7 @@ from .utils import (
     WEIGHTS_NAME,
     PeftConfig,
     PeftType,
+    PromptLearningConfig,
     TaskType,
     _set_trainable,
     get_peft_model_state_dict,
@@ -52,14 +53,14 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
         - **peft_config** ([`PeftConfig`]) -- The configuration of the Peft model.
         - **modules_to_save** (`list` of `str`) -- The list of sub-module names to save when
         saving the model.
-        - **prompt_encoder** ([`PromptEncoder`]) -- The prompt encoder used for Peft if `peft_config.peft_type
-        != PeftType.LORA`.
+        - **prompt_encoder** ([`PromptEncoder`]) -- The prompt encoder used for Peft if
+        `isinstance(self.peft_config, PromptLearningConfig)`.
         - **prompt_tokens** (`torch.Tensor`) -- The virtual prompt tokens used for Peft if
-        `peft_config.peft_type != PeftType.LORA`.
+        `isinstance(self.peft_config, PromptLearningConfig)`.
         - **transformer_backbone_name** (`str`) -- The name of the transformer
-        backbone in the base model if `peft_config.peft_type != PeftType.LORA`.
+        backbone in the base model if `isinstance(self.peft_config, PromptLearningConfig)`.
         - **word_embeddings** (`torch.nn.Embedding`) -- The word embeddings of the transformer backbone
-        in the base model if `peft_config.peft_type != PeftType.LORA`.
+        in the base model if `isinstance(self.peft_config, PromptLearningConfig)`.
     """
 
     def __init__(self, model, peft_config: PeftConfig):
@@ -68,10 +69,13 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
         self.base_model = model
         self.config = self.base_model.config
         self.modules_to_save = None
-        if peft_config.peft_type != PeftType.LORA:
+        if isinstance(self.peft_config, PromptLearningConfig):
             self._setup_prompt_encoder()
         else:
             self.base_model = LoraModel(peft_config, model)
+        if getattr(self.peft_config, "modules_to_save", None) is not None:
+            self.modules_to_save = self.peft_config.modules_to_save
+            _set_trainable(self.base_model)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     def save_pretrained(self, save_directory, **kwargs):
@@ -123,7 +127,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
         # 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)
+        if not isinstance(config, PromptLearningConfig):
+            model = cls(model, config)
+        else:
+            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)):
@@ -291,7 +298,7 @@ class PeftModelForSequenceClassification(PeftModel):
     ):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        if self.peft_config.peft_type == PeftType.LORA:
+        if not isinstance(self.peft_config, PromptLearningConfig):
             return self.base_model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -448,7 +455,7 @@ class PeftModelForCausalLM(PeftModel):
         return_dict=None,
         **kwargs,
     ):
-        if self.peft_config.peft_type == PeftType.LORA:
+        if not isinstance(self.peft_config, PromptLearningConfig):
             return self.base_model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -497,7 +504,7 @@ class PeftModelForCausalLM(PeftModel):
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
     def generate(self, **kwargs):
-        if self.peft_config.peft_type == PeftType.LORA:
+        if not isinstance(self.peft_config, PromptLearningConfig):
             return self.base_model.generate(**kwargs)
         else:
             if "input_ids" not in kwargs:
@@ -579,7 +586,7 @@ class PeftModelForSeq2SeqLM(PeftModel):
         return_dict=None,
         **kwargs,
     ):
-        if self.peft_config.peft_type == PeftType.LORA:
+        if not isinstance(self.peft_config, PromptLearningConfig):
             return self.base_model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -647,7 +654,7 @@ class PeftModelForSeq2SeqLM(PeftModel):
             return self.base_model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs)
 
     def generate(self, **kwargs):
-        if self.peft_config.peft_type == PeftType.LORA:
+        if not isinstance(self.peft_config, PromptLearningConfig):
             return self.base_model.generate(**kwargs)
         else:
             if "input_ids" not in kwargs:
@@ -735,7 +742,7 @@ class PeftModelForTokenClassification(PeftModel):
     ):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        if self.peft_config.peft_type == PeftType.LORA:
+        if not isinstance(self.peft_config, PromptLearningConfig):
             return self.base_model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,

src/peft/tuners/lora.py

@@ -53,6 +53,8 @@ class LoraConfig(PeftConfig):
         fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out)
         enable_lora ( `List[bool]`): Used with `lora.MergedLinear`.
         bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only'
+        modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable
+            and saved in the final checkpoint.
     """
 
     r: int = field(default=8, metadata={"help": "Lora attention dimension"})
@@ -68,6 +70,14 @@ class LoraConfig(PeftConfig):
     )
     enable_lora: Optional[List[bool]] = field(default=None, metadata={"help": "Used with `lora.MergedLinear`."})
     bias: str = field(default="none", metadata={"help": "Bias type for Lora. Can be 'none', 'all' or 'lora_only'"})
+    modules_to_save: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "List of modules apart from LoRA layers to be set as trainable and saved in the final checkpoint. "
+            "For example, in Sequence Classification or Token Classification tasks, "
+            "the final layer `classifier/score` are randomly initialized and as such need to be trainable and saved."
+        },
+    )
 
     def __post_init__(self):
         self.peft_type = PeftType.LORA