diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index 391be4a..31b4aa5 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -1,15 +1,15 @@ -from transformers import AutoModelForSeq2SeqLM -from peft import get_peft_model, AdaLoraConfig, AdaLoraModel, TaskType -import torch -from datasets import load_dataset import os -os.environ["TOKENIZERS_PARALLELISM"] = "false" -from transformers import AutoTokenizer -from torch.utils.data import DataLoader -from transformers import default_data_collator, get_linear_schedule_with_warmup -from tqdm import tqdm +import torch from datasets import load_dataset +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, default_data_collator, get_linear_schedule_with_warmup + +from peft import AdaLoraConfig, PeftConfig, PeftModel, TaskType, get_peft_model + + +os.environ["TOKENIZERS_PARALLELISM"] = "false" device = "cuda" model_name_or_path = "facebook/bart-base" @@ -20,18 +20,23 @@ text_column = "sentence" label_column = "text_label" max_length = 128 lr = 1e-3 -num_epochs = 8 +num_epochs = 8 batch_size = 8 # creating model peft_config = AdaLoraConfig( - init_r=12, target_r=8, - beta1=0.85, beta2=0.85, - tinit=200, tfinal=1000, deltaT=10, - lora_alpha=32, lora_dropout=0.1, - task_type=TaskType.SEQ_2_SEQ_LM, - inference_mode=False + init_r=12, + target_r=8, + beta1=0.85, + beta2=0.85, + tinit=200, + tfinal=1000, + deltaT=10, + lora_alpha=32, + lora_dropout=0.1, + task_type=TaskType.SEQ_2_SEQ_LM, + inference_mode=False, ) model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) @@ -98,7 +103,7 @@ model.base_model.peft_config.total_step = len(train_dataloader) * num_epochs # training and evaluation model = model.to(device) -global_step = 0 +global_step = 0 for epoch in range(num_epochs): model.train() total_loss = 0 @@ -110,8 +115,8 @@ for epoch in range(num_epochs): loss.backward() optimizer.step() lr_scheduler.step() - # Update the importance of low-rank matrices - # and allocate the budget accordingly. + # Update the importance of low-rank matrices + # and allocate the budget accordingly. model.base_model.update_and_allocate(global_step) optimizer.zero_grad() global_step += 1 @@ -158,8 +163,6 @@ ckpt = f"{peft_model_id}/adapter_model.bin" # get_ipython().system('du -h $ckpt') -from peft import PeftModel, PeftConfig - peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}" config = PeftConfig.from_pretrained(peft_model_id) @@ -177,5 +180,3 @@ with torch.no_grad(): outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=10) print(outputs) print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)) - - diff --git a/src/peft/__init__.py b/src/peft/__init__.py index 3a68691..dc666ce 100644 --- a/src/peft/__init__.py +++ b/src/peft/__init__.py @@ -30,8 +30,8 @@ from .peft_model import ( from .tuners import ( LoraConfig, LoraModel, - AdaLoraConfig, - AdaLoraModel, + AdaLoraConfig, + AdaLoraModel, PrefixEncoder, PrefixTuningConfig, PromptEmbedding, diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 7d10f0f..260335c 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -20,7 +20,7 @@ from .peft_model import ( PeftModelForSequenceClassification, PeftModelForTokenClassification, ) -from .tuners import LoraConfig, AdaLoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig +from .tuners import AdaLoraConfig, LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from .utils import PromptLearningConfig @@ -36,7 +36,7 @@ PEFT_TYPE_TO_CONFIG_MAPPING = { "PREFIX_TUNING": PrefixTuningConfig, "P_TUNING": PromptEncoderConfig, "LORA": LoraConfig, - "ADALORA": AdaLoraConfig, + "ADALORA": AdaLoraConfig, } TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { @@ -145,6 +145,7 @@ def _prepare_lora_config(peft_config, model_config): peft_config.merge_weights = True return peft_config + def _prepare_adalora_config(peft_config, model_config): if peft_config.target_modules is None: if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING: diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 2d80efc..45a5b75 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -28,9 +28,7 @@ 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, AdaLoraConfig, AdaLoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder +from .tuners import AdaLoraConfig, AdaLoraModel, LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder from .utils import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, WEIGHTS_NAME, diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 1e7be53..f98b3c6 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -1,19 +1,14 @@ import importlib -import math import re -import warnings -import numpy as np -from dataclasses import asdict, dataclass, field -from enum import Enum -from typing import List, Optional, Union +from dataclasses import dataclass, field +from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F -from transformers.pytorch_utils import Conv1D -from ..utils import PeftConfig, PeftType, transpose -from .lora import LoraConfig, LoraModel, LoraLayer, mark_only_lora_as_trainable +from ..utils import PeftType, transpose +from .lora import LoraConfig, LoraLayer, LoraModel, mark_only_lora_as_trainable def is_bnb_available(): @@ -30,46 +25,37 @@ class AdaLoraConfig(LoraConfig): This is the configuration class to store the configuration of a [`~peft.AdaLora`]. Args: - target_r (`int`): The target average rank of incremental matrix. - init_r (`int`): The initial rank for each incremental matrix. - tinit (`int`): The steps of initial fine-tuning warmup. - tfinal (`int`): The step of final fine-tuning. - deltaT (`int`): The time internval between two budget allocations. + target_r (`int`): The target average rank of incremental matrix. + init_r (`int`): The initial rank for each incremental matrix. + tinit (`int`): The steps of initial fine-tuning warmup. + tfinal (`int`): The step of final fine-tuning. + deltaT (`int`): The time internval between two budget allocations. beta1 (`float`): The hyperparameter of EMA for sensitivity smoothing. - beta2 (`float`): The hyperparameter of EMA for undertainty quantification. - orth_reg_weight (`float`): The coefficient of orthogonal regularization. - total_step (`int`): The total training steps that should be specified before training. - rank_pattern (`list`): The allocated rank for each weight matrix by RankAllocator. + beta2 (`float`): The hyperparameter of EMA for undertainty quantification. + orth_reg_weight (`float`): The coefficient of orthogonal regularization. + total_step (`int`): The total training steps that should be specified before training. + rank_pattern (`list`): The allocated rank for each weight matrix by RankAllocator. """ + target_r: int = field(default=8, metadata={"help": "Target Lora matrix dimension."}) init_r: int = field(default=12, metadata={"help": "Intial Lora matrix dimension."}) tinit: int = field(default=0, metadata={"help": "The steps of initial warmup."}) tfinal: int = field(default=0, metadata={"help": "The steps of final warmup."}) deltaT: int = field(default=1, metadata={"help": "Step interval of rank allocation."}) beta1: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) - beta2: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) - orth_reg_weight: float = field( - default=0.5, - metadata={"help": "The orthogonal regularization coefficient."} - ) - total_step: Optional[int] = field( - default=None, - metadata={"help": "The total training steps."} - ) - rank_pattern: Optional[dict] = field( - default=None, - metadata={"help":"The saved rank pattern."} - ) + beta2: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) + orth_reg_weight: float = field(default=0.5, metadata={"help": "The orthogonal regularization coefficient."}) + total_step: Optional[int] = field(default=None, metadata={"help": "The total training steps."}) + rank_pattern: Optional[dict] = field(default=None, metadata={"help": "The saved rank pattern."}) def __post_init__(self): self.peft_type = PeftType.ADALORA - class AdaLoraModel(LoraModel): """ - Creates AdaLoRA (Adaptive LoRA) model from a pretrained transformers model. - Paper: https://openreview.net/pdf?id=lq62uWRJjiY + Creates AdaLoRA (Adaptive LoRA) model from a pretrained transformers model. Paper: + https://openreview.net/pdf?id=lq62uWRJjiY Args: model ([`transformers.PreTrainedModel`]): The model to be adapted. @@ -80,13 +66,12 @@ class AdaLoraModel(LoraModel): Example:: - >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import AdaLoraModel, AdaLoraConfig + >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import AdaLoraModel, AdaLoraConfig >>> config = AdaLoraConfig( peft_type="ADALORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"], lora_dropout=0.01, ) - >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") - >>> model = AdaLoraModel(config, model) + >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> model = AdaLoraModel(config, model) **Attributes**: - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. @@ -101,10 +86,7 @@ class AdaLoraModel(LoraModel): mark_only_lora_as_trainable(self.model, self.peft_config.bias) self.rankallocator = RankAllocator(config, self.model) if config.enable_lora is not None: - raise NotImplementedError( - "MergedLinear has not been implemented for AdaLoRA." - ) - + raise NotImplementedError("MergedLinear has not been implemented for AdaLoRA.") def _find_and_replace(self): loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) @@ -159,26 +141,25 @@ class AdaLoraModel(LoraModel): return getattr(self.model, name) def forward(self, *args, **kwargs): - outputs = self.model.forward(*args, **kwargs) + outputs = self.model.forward(*args, **kwargs) - # Calculate the orthogonal regularization + # Calculate the orthogonal regularization orth_reg_weight = self.peft_config.orth_reg_weight - assert orth_reg_weight > 0 + assert orth_reg_weight > 0 if hasattr(outputs, "loss"): - regu_loss = 0 - num_param = 0 - for n,p in self.model.named_parameters(): + regu_loss = 0 + num_param = 0 + for n, p in self.model.named_parameters(): if "lora_A" in n or "lora_B" in n: - para_cov = p @ p.T if "lora_A" in n else p.T @ p + para_cov = p @ p.T if "lora_A" in n else p.T @ p I = torch.eye(*para_cov.size(), out=torch.empty_like(para_cov)) I.requires_grad = False num_param += 1 - regu_loss += torch.norm(para_cov-I, p="fro") - regu_loss = regu_loss / num_param - outputs.loss += orth_reg_weight * regu_loss - return outputs - + regu_loss += torch.norm(para_cov - I, p="fro") + regu_loss = regu_loss / num_param + outputs.loss += orth_reg_weight * regu_loss + return outputs def _prepare_new_module(self, target, rank_idx): if isinstance(rank_idx, list): @@ -187,7 +168,7 @@ class AdaLoraModel(LoraModel): rank_idx = rank_idx.view(-1) rank = rank_idx.sum().item() else: - raise ValueError(f"Unexcepted type of rank_idx") + raise ValueError("Unexcepted type of rank_idx") kwargs = { "r": rank, "lora_alpha": self.peft_config.lora_alpha, @@ -218,68 +199,64 @@ class AdaLoraModel(LoraModel): if rank > 0: new_module.lora_E.copy_(target.lora_E[rank_idx]) new_module.lora_A.copy_(target.lora_A[rank_idx]) - new_module.lora_B.copy_(target.lora_B[:,rank_idx]) - # The scaling is exactly as the previous + new_module.lora_B.copy_(target.lora_B[:, rank_idx]) + # The scaling is exactly as the previous new_module.ranknum.copy_(target.ranknum) return new_module def resize_modules_by_rank_pattern(self, rank_pattern): - for name,rank_idx in rank_pattern.items(): - key = ".".join(name.split(".")[0:-1]) - parent, target, target_name = self._get_submodules(key) + for name, rank_idx in rank_pattern.items(): + key = ".".join(name.split(".")[0:-1]) + parent, target, target_name = self._get_submodules(key) new_module = self._prepare_new_module(target, rank_idx) self._replace_module(parent, target_name, new_module, target) def update_and_allocate(self, global_step): - # Update the importance score and allocate the budget + # Update the importance score and allocate the budget if global_step < self.peft_config.total_step - self.peft_config.tfinal: budget, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step) if rank_pattern: - self.peft_config.rank_pattern = rank_pattern - # Finalize the budget allocation - elif global_step == self.peft_config.total_step - self.peft_config.tfinal: - budget, rank_pattern = self.rankallocator.update_and_allocate( - self.model, global_step, force_mask=True - ) + self.peft_config.rank_pattern = rank_pattern + # Finalize the budget allocation + elif global_step == self.peft_config.total_step - self.peft_config.tfinal: + budget, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True) self.resize_modules_by_rank_pattern(rank_pattern) self.peft_config.rank_pattern = rank_pattern - self.rankallocator.reset_ipt() - # Pass the function and do forward propagation - else: + self.rankallocator.reset_ipt() + # Pass the function and do forward propagation + else: return None - class SVDLinear(nn.Linear, LoraLayer): # SVD-based adaptation by a dense layer def __init__( - self, - in_features: int, - out_features: int, - r: int = 0, - lora_alpha: int = 1, - lora_dropout: float = 0., - fan_in_fan_out: bool = False, + self, + in_features: int, + out_features: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + fan_in_fan_out: bool = False, merge_weights: bool = True, - **kwargs + **kwargs, ): nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, - merge_weights=merge_weights) - + LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + self.fan_in_fan_out = fan_in_fan_out # Actual trainable parameters if r > 0: # Right singular vectors self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) - # Singular values - self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) + # Singular values + self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) # Left singular vectors self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) # The current rank self.ranknum = nn.Parameter(self.weight.new_zeros(1), requires_grad=False) self.ranknum.data.fill_(float(self.r)) - self.scaling = self.lora_alpha if self.lora_alpha>0 else float(self.r) + self.scaling = self.lora_alpha if self.lora_alpha > 0 else float(self.r) # Freezing the pre-trained weight matrix self.weight.requires_grad = False self.ranknum.requires_grad = False @@ -289,7 +266,7 @@ class SVDLinear(nn.Linear, LoraLayer): def reset_parameters(self): nn.Linear.reset_parameters(self) - if hasattr(self, 'lora_A'): + if hasattr(self, "lora_A"): nn.init.zeros_(self.lora_E) nn.init.normal_(self.lora_A, mean=0.0, std=0.02) nn.init.normal_(self.lora_B, mean=0.0, std=0.02) @@ -299,19 +276,19 @@ class SVDLinear(nn.Linear, LoraLayer): if self.merge_weights and self.merged: # Make sure that the weights are not merged if self.r > 0: - self.weight.data -= transpose( - self.lora_B @ (self.lora_A * self.lora_E) - ) * self.scaling/(self.ranknum+1e-5) + self.weight.data -= ( + transpose(self.lora_B @ (self.lora_A * self.lora_E)) * self.scaling / (self.ranknum + 1e-5) + ) self.merged = False - + def eval(self): nn.Linear.eval(self) if self.merge_weights and not self.merged: # Merge the weights and mark it if self.r > 0: - self.weight.data += transpose( - self.lora_B @ (self.lora_A * self.lora_E) - ) * self.scaling/(self.ranknum+1e-5) + self.weight.data += ( + transpose(self.lora_B @ (self.lora_A * self.lora_E)) * self.scaling / (self.ranknum + 1e-5) + ) self.merged = True def forward(self, x: torch.Tensor): @@ -319,16 +296,19 @@ class SVDLinear(nn.Linear, LoraLayer): result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if self.r > 0: result += ( - self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T - ) * self.scaling / (self.ranknum+1e-5) + (self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T) + * self.scaling + / (self.ranknum + 1e-5) + ) return result else: return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if is_bnb_available(): + class SVDLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): - # Low-rank matrix for SVD-based adaptation + # Low-rank matrix for SVD-based adaptation def __init__( self, in_features, @@ -353,14 +333,14 @@ if is_bnb_available(): if r > 0: # Right singular vectors self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) - # Singular values - self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) + # Singular values + self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) # Left singular vectors self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) # The current rank self.ranknum = nn.Parameter(self.weight.new_zeros(1), requires_grad=False) self.ranknum.data.fill_(float(self.r)) - self.scaling = self.lora_alpha if self.lora_alpha>0 else float(self.r) + self.scaling = self.lora_alpha if self.lora_alpha > 0 else float(self.r) # Freezing the pre-trained weight matrix self.weight.requires_grad = False self.ranknum.requires_grad = False @@ -385,34 +365,33 @@ if is_bnb_available(): if x.dtype != torch.float32: x = x.float() output = ( - self.lora_dropout(x) @ (self.lora_A*self.lora_E).T @ self.lora_B.T /(self.ranknum+1e-5) - ).to(expected_dtype) * self.scaling + self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T / (self.ranknum + 1e-5) + ).to(expected_dtype) * self.scaling result += output else: output = ( - self.lora_dropout(x) @ (self.lora_A*self.lora_E).T @ self.lora_B.T /(self.ranknum+1e-5) - ) * self.scaling + self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T / (self.ranknum + 1e-5) + ) * self.scaling result += output return result - class RankAllocator(object): """ - The RankAllocator for AdaLoraModel. - Paper: https://openreview.net/pdf?id=lq62uWRJjiY + The RankAllocator for AdaLoraModel. Paper: https://openreview.net/pdf?id=lq62uWRJjiY Args: config ([`AdaLoraConfig`]): The configuration of the AdaLora model. - model: the model that we apply AdaLoRA to. + model: the model that we apply AdaLoRA to. """ + def __init__(self, peft_config, model): self.peft_config = peft_config - self.beta1 = peft_config.beta1 - self.beta2 = peft_config.beta2 - assert (self.beta1>0 and self.beta1<1) - assert (self.beta2>0 and self.beta2<1) + self.beta1 = peft_config.beta1 + self.beta2 = peft_config.beta2 + assert self.beta1 > 0 and self.beta1 < 1 + assert self.beta2 > 0 and self.beta2 < 1 self.reset_ipt() self._set_budget_scheduler(model) @@ -421,129 +400,125 @@ class RankAllocator(object): self.peft_config.total_step = total_step def reset_ipt(self): - self.ipt = {} + self.ipt = {} self.exp_avg_ipt = {} self.exp_avg_unc = {} def _set_budget_scheduler(self, model): - self.init_bgt = 0 - self.name_set = set() - for n,p in model.named_parameters(): - if "lora_A" in n: - self.init_bgt += p.size(0) + self.init_bgt = 0 + self.name_set = set() + for n, p in model.named_parameters(): + if "lora_A" in n: + self.init_bgt += p.size(0) self.name_set.add(n.replace("lora_A", "%s")) - self.name_set = list(sorted(self.name_set)) - # The total final rank budget - self.target_bgt = self.peft_config.target_r * len(self.name_set) + self.name_set = sorted(self.name_set) + # The total final rank budget + self.target_bgt = self.peft_config.target_r * len(self.name_set) - def budget_schedule(self, step:int): - tinit = self.peft_config.tinit - tfinal = self.peft_config.tfinal - total_step = self.peft_config.total_step - # Initial warmup - if step <= tinit: - budget = self.init_bgt - mask_ind = False - # Final fine-tuning - elif step > total_step - tfinal: - budget = self.target_bgt - mask_ind = True - else: - # Budget decreasing with a cubic scheduler - mul_coeff = 1 - (step-tinit) / (total_step-tfinal-tinit) - budget = int( - (self.init_bgt-self.target_bgt)*(mul_coeff**3)+self.target_bgt - ) - mask_ind = True if step % self.peft_config.deltaT == 0 else False - return budget, mask_ind + def budget_schedule(self, step: int): + tinit = self.peft_config.tinit + tfinal = self.peft_config.tfinal + total_step = self.peft_config.total_step + # Initial warmup + if step <= tinit: + budget = self.init_bgt + mask_ind = False + # Final fine-tuning + elif step > total_step - tfinal: + budget = self.target_bgt + mask_ind = True + else: + # Budget decreasing with a cubic scheduler + mul_coeff = 1 - (step - tinit) / (total_step - tfinal - tinit) + budget = int((self.init_bgt - self.target_bgt) * (mul_coeff**3) + self.target_bgt) + mask_ind = True if step % self.peft_config.deltaT == 0 else False + return budget, mask_ind - def update_ipt(self, model): - # Update the sensitivity and uncertainty for every weight - for n,p in model.named_parameters(): - if "lora_" in n: + def update_ipt(self, model): + # Update the sensitivity and uncertainty for every weight + for n, p in model.named_parameters(): + if "lora_" in n: if n not in self.ipt: - self.ipt[n] = torch.zeros_like(p) - self.exp_avg_ipt[n] = torch.zeros_like(p) - self.exp_avg_unc[n] = torch.zeros_like(p) + self.ipt[n] = torch.zeros_like(p) + self.exp_avg_ipt[n] = torch.zeros_like(p) + self.exp_avg_unc[n] = torch.zeros_like(p) with torch.no_grad(): self.ipt[n] = (p * p.grad).abs().detach() - # Sensitivity smoothing - self.exp_avg_ipt[n] = self.beta1 * self.exp_avg_ipt[n] + \ - (1 - self.beta1)*self.ipt[n] - # Uncertainty quantification - self.exp_avg_unc[n] = self.beta2 * self.exp_avg_unc[n] + \ - (1-self.beta2)*(self.ipt[n]-self.exp_avg_ipt[n]).abs() + # Sensitivity smoothing + self.exp_avg_ipt[n] = self.beta1 * self.exp_avg_ipt[n] + (1 - self.beta1) * self.ipt[n] + # Uncertainty quantification + self.exp_avg_unc[n] = ( + self.beta2 * self.exp_avg_unc[n] + (1 - self.beta2) * (self.ipt[n] - self.exp_avg_ipt[n]).abs() + ) def _element_score(self, n): return self.exp_avg_ipt[n] * self.exp_avg_unc[n] def _combine_ipt(self, ipt_E, ipt_AB): ipt_AB = ipt_AB.sum(dim=1, keepdim=False) - sum_ipt = ipt_E.view(-1) + ipt_AB.view(-1) - return sum_ipt + sum_ipt = ipt_E.view(-1) + ipt_AB.view(-1) + return sum_ipt - def mask_to_budget(self, model, budget): + def mask_to_budget(self, model, budget): value_ipt = {} - vector_ipt = {} + vector_ipt = {} triplet_ipt = {} # Get the importance score for A, E, B - for n,p in model.named_parameters(): - if "lora_A" in n: + for n, p in model.named_parameters(): + if "lora_A" in n: entry_ipt = self._element_score(n) comb_ipt = torch.mean(entry_ipt, dim=1, keepdim=True) name_m = n.replace("lora_A", "%s") - if name_m not in vector_ipt: + if name_m not in vector_ipt: vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) - if "lora_B" in n: + if "lora_B" in n: entry_ipt = self._element_score(n) comb_ipt = torch.mean(entry_ipt, dim=0, keepdim=False).view(-1, 1) name_m = n.replace("lora_B", "%s") - if name_m not in vector_ipt: + if name_m not in vector_ipt: vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) if "lora_E" in n: - entry_ipt = self._element_score(n) + entry_ipt = self._element_score(n) name_m = n.replace("lora_E", "%s") value_ipt[name_m] = entry_ipt all_score = [] - # Calculate the score for each triplet - for name_m in vector_ipt: - ipt_E = value_ipt[name_m] + # Calculate the score for each triplet + for name_m in vector_ipt: + ipt_E = value_ipt[name_m] ipt_AB = torch.cat(vector_ipt[name_m], dim=1) sum_ipt = self._combine_ipt(ipt_E, ipt_AB) - name_E = name_m%"lora_E" + name_E = name_m % "lora_E" triplet_ipt[name_E] = sum_ipt.view(-1, 1) all_score.append(sum_ipt.view(-1)) # Get the threshold by ranking ipt mask_threshold = torch.kthvalue( - torch.cat(all_score), - k = self.init_bgt - budget, + torch.cat(all_score), + k=self.init_bgt - budget, )[0].item() rank_pattern = {} - # Mask the unimportant triplets + # Mask the unimportant triplets with torch.no_grad(): - for n,p in model.named_parameters(): - if "lora_E" in n: - p.masked_fill_(triplet_ipt[n]<=mask_threshold, 0.0) - rank_pattern[n] = (~(triplet_ipt[n]<=mask_threshold)).view(-1).tolist() + for n, p in model.named_parameters(): + if "lora_E" in n: + p.masked_fill_(triplet_ipt[n] <= mask_threshold, 0.0) + rank_pattern[n] = (~(triplet_ipt[n] <= mask_threshold)).view(-1).tolist() return rank_pattern def update_and_allocate(self, model, global_step, force_mask=False): - # # Update the importance score and allocate the budget + # # Update the importance score and allocate the budget if global_step < self.peft_config.total_step - self.peft_config.tfinal: self.update_ipt(model) budget, mask_ind = self.budget_schedule(global_step) - # Allocate the budget according to importance scores + # Allocate the budget according to importance scores if mask_ind or force_mask: rank_pattern = self.mask_to_budget(model, budget) else: - rank_pattern = None + rank_pattern = None return budget, rank_pattern - - diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index cf4c813..0b10717 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -72,10 +72,10 @@ def set_peft_model_state_dict(model, peft_model_state_dict): """ if model.peft_config.peft_type == PeftType.ADALORA: rank_pattern = model.peft_config.rank_pattern - if rank_pattern: + if rank_pattern: model.base_model.resize_modules_by_rank_pattern(rank_pattern) model.load_state_dict(peft_model_state_dict, strict=False) - + if model.peft_config.peft_type not in (PeftType.LORA, PeftType.ADALORA): model.prompt_encoder.embedding.load_state_dict( {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True