wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-10 12:20:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Run make style and make quality

Browse Source
This commit is contained in:
Qingru Zhang
2023-04-05 20:52:12 +00:00
committed by zqingru
parent 4f8c134102
commit 072da6d9d6
6 changed files with 185 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/conditional_generation/peft_adalora_seq2seq.py
+24 -23
View File
@@ -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))
src/peft/__init__.py
+2 -2
View File
@@ -30,8 +30,8 @@ from .peft_model import (
from .tuners import (
LoraConfig,
LoraModel,
AdaLoraConfig,
AdaLoraModel,
AdaLoraConfig,
AdaLoraModel,
PrefixEncoder,
PrefixTuningConfig,
PromptEmbedding,
src/peft/mapping.py
+3 -2
View File
@@ -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:
src/peft/peft_model.py
+1 -3
View File
@@ -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,
src/peft/tuners/adalora.py
+153 -178
View File
@@ -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
src/peft/utils/save_and_load.py
+2 -2
View File
@@ -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