diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 18e292b..e885ea4 100644 --- a/src/peft/mapping.py +++ b/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) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index f3e01c1..89f0e0d 100644 --- a/src/peft/peft_model.py +++ b/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, diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index a1eb3b1..567ecc6 100644 --- a/src/peft/tuners/lora.py +++ b/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