From d04f6661eec9dc2811cabe6c4405e64a9d49a07f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 1 Feb 2023 15:41:35 +0530 Subject: [PATCH 1/7] 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 | 13 ++++++------- src/peft/peft_model.py | 33 ++++++++++++++++++++------------- src/peft/tuners/lora.py | 10 ++++++++++ 3 files changed, 36 insertions(+), 20 deletions(-) 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 From d53a631608e1e3de3eaac288d5c30e7ed14e4abb Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 1 Feb 2023 15:59:24 +0530 Subject: [PATCH 2/7] fixes --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 89f0e0d..eba63a6 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -75,7 +75,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): 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) + _set_trainable(self) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def save_pretrained(self, save_directory, **kwargs): From 915a5db0c60b2d9fc7a64d37ced24bf0a348408b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 1 Feb 2023 16:25:42 +0530 Subject: [PATCH 3/7] fixes --- src/peft/peft_model.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index eba63a6..af17075 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -96,7 +96,11 @@ class PeftModel(PushToHubMixin, torch.nn.Module): # 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.base_model_name_or_path = ( + self.base_model.__dict__.get("name_or_path", None) + if isinstance(self.peft_config, PromptLearningConfig) + else self.base_model.model.__dict__.get("name_or_path", None) + ) self.peft_config.inference_mode = True self.peft_config.save_pretrained(save_directory) From fcd213708d7c2509bef42fb2532029f041d284c9 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 1 Feb 2023 17:17:14 +0530 Subject: [PATCH 4/7] fixes --- src/peft/peft_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index af17075..38993bf 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -76,6 +76,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): 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) + self.forward = self.base_model.forward self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def save_pretrained(self, save_directory, **kwargs): From c884daf96abc005e6d0cca44c00515da974d6cda Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 1 Feb 2023 19:18:38 +0530 Subject: [PATCH 5/7] getting rid to forward call linking --- src/peft/mapping.py | 9 ++++++--- src/peft/peft_model.py | 3 ++- src/peft/tuners/lora.py | 3 ++- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/peft/mapping.py b/src/peft/mapping.py index e885ea4..14dd6f2 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -21,7 +21,7 @@ from .peft_model import ( PeftModelForTokenClassification, ) from .tuners import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig -from .utils import PromptLearningConfig +from .utils import PeftType, PromptLearningConfig MODEL_TYPE_TO_PEFT_MODEL_MAPPING = { @@ -135,8 +135,11 @@ def get_peft_model(model, peft_config): model_config = model.config.to_dict() peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None) - if not isinstance(peft_config, PromptLearningConfig): + if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): peft_config = _prepare_lora_config(peft_config, model_config) return PeftModel(model, peft_config) - peft_config = _prepare_prompt_learning_config(peft_config, model_config) + if not isinstance(peft_config, PromptLearningConfig): + peft_config = _prepare_lora_config(peft_config, model_config) + else: + 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 38993bf..7403d6b 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -76,7 +76,6 @@ class PeftModel(PushToHubMixin, torch.nn.Module): 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) - self.forward = self.base_model.forward self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def save_pretrained(self, save_directory, **kwargs): @@ -246,6 +245,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module): def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" + if name == "forward": + return getattr(self.base_model, name) try: return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 567ecc6..f4676d2 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -115,7 +115,6 @@ class LoraModel(torch.nn.Module): self.model = model self._find_and_replace() mark_only_lora_as_trainable(self.model, self.peft_config.bias) - self.forward = self.model.forward def _find_and_replace(self): kwargs = { @@ -174,6 +173,8 @@ class LoraModel(torch.nn.Module): def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" + if name == "forward": + return getattr(self.model, name) try: return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: From c37ee25be79a81dbba80e34df5e215c3b6b4b75d Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 1 Feb 2023 19:35:19 +0530 Subject: [PATCH 6/7] trying diff approaches --- src/peft/mapping.py | 2 +- src/peft/peft_model.py | 11 +++++++++-- src/peft/tuners/lora.py | 3 +-- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 14dd6f2..68de0c2 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -21,7 +21,7 @@ from .peft_model import ( PeftModelForTokenClassification, ) from .tuners import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig -from .utils import PeftType, PromptLearningConfig +from .utils import PromptLearningConfig MODEL_TYPE_TO_PEFT_MODEL_MAPPING = { diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 7403d6b..ef8892a 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -245,13 +245,20 @@ class PeftModel(PushToHubMixin, torch.nn.Module): def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" - if name == "forward": - return getattr(self.base_model, name) try: return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: return getattr(self.base_model, name) + def forward(self, *args, **kwargs): + """ + Forward pass of the model. + """ + if isinstance(self.peft_config, PromptLearningConfig): + return self.base_model(*args, **kwargs) + else: + return self.base_model.model(*args, **kwargs) + class PeftModelForSequenceClassification(PeftModel): """ diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index f4676d2..567ecc6 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -115,6 +115,7 @@ class LoraModel(torch.nn.Module): self.model = model self._find_and_replace() mark_only_lora_as_trainable(self.model, self.peft_config.bias) + self.forward = self.model.forward def _find_and_replace(self): kwargs = { @@ -173,8 +174,6 @@ class LoraModel(torch.nn.Module): def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" - if name == "forward": - return getattr(self.model, name) try: return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: From 44d8e72ca8d5af9a0b7ef93c773c31edbc90c2fd Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 2 Feb 2023 13:19:14 +0530 Subject: [PATCH 7/7] fixes --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index ef8892a..32937ae 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -131,7 +131,7 @@ 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) - if not isinstance(config, PromptLearningConfig): + if config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): model = cls(model, config) else: model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config)