diff --git a/src/peft/tuners/adaption_prompt/model.py b/src/peft/tuners/adaption_prompt/model.py index 682df120ea..e547a1ca37 100644 --- a/src/peft/tuners/adaption_prompt/model.py +++ b/src/peft/tuners/adaption_prompt/model.py @@ -57,7 +57,7 @@ def __init__(self, model, configs: Dict, adapter_name: str): self._enabled = True self.forward = self.model.forward self.add_adapter(adapter_name, configs[adapter_name]) - self._mark_only_adaption_prompts_as_trainable() + self._mark_only_adaption_prompts_as_trainable(self.model) def add_adapter(self, adapter_name: str, config: AdaptionPromptConfig) -> None: """Add an adapter with the given name and config.""" @@ -146,9 +146,9 @@ def _remove_adapted_attentions(self, adapter_name: str) -> None: setattr(par, config.target_modules, attn.model) self._cached_adapters[adapter_name] = adapted_attentions - def _mark_only_adaption_prompts_as_trainable(self) -> None: + def _mark_only_adaption_prompts_as_trainable(self, model: nn.Module) -> None: """Freeze all parameters of the model except the adaption prompts.""" - for n, p in self.model.named_parameters(): + for n, p in model.named_parameters(): if not is_adaption_prompt_trainable(n): p.requires_grad = False diff --git a/src/peft/tuners/ia3/model.py b/src/peft/tuners/ia3/model.py index a6617e6ef8..845f8d2ef3 100644 --- a/src/peft/tuners/ia3/model.py +++ b/src/peft/tuners/ia3/model.py @@ -21,6 +21,7 @@ from typing import List, Optional import torch +from torch import nn from transformers.pytorch_utils import Conv1D from peft.import_utils import is_bnb_4bit_available, is_bnb_available @@ -147,8 +148,8 @@ def _create_new_module(ia3_config, adapter_name, target, **kwargs): def _check_target_module_exists(ia3_config, key): return check_target_module_exists(ia3_config, key) - def _mark_only_adapters_as_trainable(self) -> None: - for n, p in self.model.named_parameters(): + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): if self.prefix not in n: p.requires_grad = False diff --git a/src/peft/tuners/lora/model.py b/src/peft/tuners/lora/model.py index f6b6453b62..419884f7b9 100644 --- a/src/peft/tuners/lora/model.py +++ b/src/peft/tuners/lora/model.py @@ -26,6 +26,7 @@ from typing import List, Optional import torch +from torch import nn from tqdm import tqdm from transformers.pytorch_utils import Conv1D @@ -221,8 +222,8 @@ def _replace_module(self, parent, child_name, new_module, child): weight = child.qweight if hasattr(child, "qweight") else child.weight module.to(weight.device) - def _mark_only_adapters_as_trainable(self) -> None: - for n, p in self.model.named_parameters(): + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): if self.prefix not in n: p.requires_grad = False @@ -232,11 +233,11 @@ def _mark_only_adapters_as_trainable(self) -> None: continue if bias == "all": - for n, p in self.model.named_parameters(): + for n, p in model.named_parameters(): if "bias" in n: p.requires_grad = True elif bias == "lora_only": - for m in self.model.modules(): + for m in model.modules(): if isinstance(m, LoraLayer) and hasattr(m, "bias") and m.bias is not None: m.bias.requires_grad = True else: diff --git a/src/peft/tuners/lycoris_utils.py b/src/peft/tuners/lycoris_utils.py index 925fdff0f5..9b46981c5e 100644 --- a/src/peft/tuners/lycoris_utils.py +++ b/src/peft/tuners/lycoris_utils.py @@ -272,8 +272,8 @@ def _create_new_module(cls, config: LycorisConfig, adapter_name: str, target: nn return new_module - def _mark_only_adapters_as_trainable(self) -> None: - for n, p in self.model.named_parameters(): + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): if self.prefix not in n: p.requires_grad = False diff --git a/src/peft/tuners/mixed/model.py b/src/peft/tuners/mixed/model.py index e8c98bc4f7..24bdb9f888 100644 --- a/src/peft/tuners/mixed/model.py +++ b/src/peft/tuners/mixed/model.py @@ -131,8 +131,8 @@ def _replace_module(self, parent, child_name, new_module, child) -> None: if "ranknum" in name: module.to(child.weight.device) - def _mark_only_adapters_as_trainable(self) -> None: - for n, p in self.model.named_parameters(): + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): if not any(prefix in n for prefix in PREFIXES): p.requires_grad = False @@ -142,12 +142,12 @@ def _mark_only_adapters_as_trainable(self) -> None: continue if bias == "all": - for n, p in self.model.named_parameters(): + for n, p in model.named_parameters(): if "bias" in n: p.requires_grad = True elif bias == "lora_only": # TODO: check if this is needed for other supported types - for m in self.model.modules(): + for m in model.modules(): if isinstance(m, Layers) and hasattr(m, "bias") and m.bias is not None: m.bias.requires_grad = True else: diff --git a/src/peft/tuners/tuners_utils.py b/src/peft/tuners/tuners_utils.py index f262ec5596..b8080fb2d8 100644 --- a/src/peft/tuners/tuners_utils.py +++ b/src/peft/tuners/tuners_utils.py @@ -167,7 +167,7 @@ def _create_and_replace( ... @abstractmethod - def _mark_only_adapters_as_trainable(self): + def _mark_only_adapters_as_trainable(self, model: nn.Module): r""" A helper method to mark only the adapter layers as trainable (i.e. module.requires_grad = False) This needs to be overriden for all tuner classes to match the correct key names. @@ -252,10 +252,10 @@ def inject_adapter(self, model: nn.Module, adapter_name: str): f"Please check the target modules and try again." ) - self._mark_only_adapters_as_trainable() + self._mark_only_adapters_as_trainable(model) if self.peft_config[adapter_name].inference_mode: - for n, p in self.model.named_parameters(): + for n, p in model.named_parameters(): if adapter_name in n: p.requires_grad = False