mirror of
https://github.com/wassname/peft.git
synced 2026-09-10 12:20:21 +08:00
refine the key of rank pattern
This commit is contained in:
+11
-17
@@ -98,7 +98,7 @@ class AdaLoraModel(LoraModel):
|
||||
self.model = model
|
||||
self._find_and_replace()
|
||||
mark_only_lora_as_trainable(self.model, self.peft_config.bias)
|
||||
self.rankallocator = RankAllocator(config, self.named_parameters())
|
||||
self.rankallocator = RankAllocator(config, self.model)
|
||||
|
||||
def _find_and_replace(self):
|
||||
loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False)
|
||||
@@ -236,7 +236,7 @@ class AdaLoraModel(LoraModel):
|
||||
|
||||
def resize_modules_by_rank_pattern(self, rank_pattern):
|
||||
for name,rank_idx in rank_pattern.items():
|
||||
key = ".".join(name.split(".")[1:-1])
|
||||
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)
|
||||
@@ -245,12 +245,14 @@ class AdaLoraModel(LoraModel):
|
||||
def update_and_allocate(self, global_step):
|
||||
# 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, global_step)
|
||||
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, global_step, force_mask=True)
|
||||
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()
|
||||
@@ -415,10 +417,10 @@ class RankAllocator(object):
|
||||
|
||||
Args:
|
||||
config ([`AdaLoraConfig`]): The configuration of the AdaLora model.
|
||||
param_iterator: the parameter iterator to initalize the rankallocator.
|
||||
model: the model that we apply AdaLoRA to.
|
||||
|
||||
"""
|
||||
def __init__(self, peft_config, param_iterator):
|
||||
def __init__(self, peft_config, model):
|
||||
self.peft_config = peft_config
|
||||
self.beta1 = peft_config.beta1
|
||||
self.beta2 = peft_config.beta2
|
||||
@@ -426,23 +428,20 @@ class RankAllocator(object):
|
||||
assert (self.beta2>0 and self.beta2<1)
|
||||
|
||||
self.reset_ipt()
|
||||
self._set_budget_scheduler(param_iterator)
|
||||
|
||||
self._set_budget_scheduler(model)
|
||||
|
||||
def set_total_step(self, total_step):
|
||||
self.peft_config.total_step = total_step
|
||||
|
||||
|
||||
def reset_ipt(self):
|
||||
self.ipt = {}
|
||||
self.exp_avg_ipt = {}
|
||||
self.exp_avg_unc = {}
|
||||
|
||||
|
||||
def _set_budget_scheduler(self, param_iterator):
|
||||
def _set_budget_scheduler(self, model):
|
||||
self.init_bgt = 0
|
||||
self.name_set = set()
|
||||
for n,p in param_iterator:
|
||||
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"))
|
||||
@@ -450,7 +449,6 @@ class RankAllocator(object):
|
||||
# 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
|
||||
@@ -472,7 +470,6 @@ class RankAllocator(object):
|
||||
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():
|
||||
@@ -490,17 +487,14 @@ class RankAllocator(object):
|
||||
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
|
||||
|
||||
|
||||
def mask_to_budget(self, model, budget):
|
||||
value_ipt = {}
|
||||
vector_ipt = {}
|
||||
|
||||
Reference in New Issue
Block a user