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

making adalora compatible with multiple adapters

Browse Source
This commit is contained in:
Sourab Mangrulkar
2023-04-06 19:06:33 +05:30
parent 75808eb2a6
commit 7397160435
7 changed files with 328 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/peft/tuners/adalora.py
+289 -144
View File
@@ -1,14 +1,27 @@
import importlib
import re
import warnings
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 PeftType, transpose
from .lora import LoraConfig, LoraLayer, LoraModel, mark_only_lora_as_trainable
from ..utils import (
TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING,
PeftType,
_freeze_adapter,
_get_submodules,
transpose,
)
from .lora import (
LoraConfig,
LoraLayer,
LoraModel,
mark_only_lora_as_trainable,
)
def is_bnb_available():
@@ -78,17 +91,42 @@ class AdaLoraModel(LoraModel):
- **peft_config** ([`AdaLoraConfig`]): The configuration of the AdaLora model.
"""
def __init__(self, config, model):
def __init__(self, model, config, adapter_name):
nn.Module.__init__(self)
self.peft_config = config
self.model = model
self._find_and_replace()
mark_only_lora_as_trainable(self.model, self.peft_config.bias)
self.peft_config = config
self.rankallocator = RankAllocator(config, self.model)
if config.enable_lora is not None:
raise NotImplementedError("MergedLinear has not been implemented for AdaLoRA.")
self.add_adapter(adapter_name, self.peft_config[adapter_name])
def _find_and_replace(self):
def add_adapter(self, adapter_name, config=None):
if config is not None:
config = self._prepare_adalora_config(config, self.model.config.to_dict())
self.peft_config[adapter_name] = config
self._find_and_replace(adapter_name)
if len(self.peft_config) > 1 and self.peft_config[adapter_name].bias != "none":
raise ValueError(
"AdaLoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters."
)
traininable_mode_counter = 0
for config in self.peft_config.values():
if not config.inference_mode:
traininable_mode_counter += 1
if traininable_mode_counter > 1:
raise ValueError(
"AdaLoraModel supports only 1 trainable adapter. "
"When using multiple adapters, set inference_mode to True for all adapters except the one you want to train."
)
if self.peft_config[adapter_name].inference_mode:
_freeze_adapter(self.model, adapter_name)
else:
self.trainable_adapter_name = adapter_name
mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias)
self.rankallocator = RankAllocator(self.model, self.peft_config[adapter_name], self.trainable_adapter_name)
def _find_and_replace(self, adapter_name):
lora_config = self.peft_config[adapter_name]
loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False)
if loaded_in_8bit and not is_bnb_available():
raise ImportError(
@@ -97,39 +135,74 @@ class AdaLoraModel(LoraModel):
)
is_target_modules_in_base_model = False
kwargs = {
"r": self.peft_config.init_r,
"lora_alpha": self.peft_config.lora_alpha,
"lora_dropout": self.peft_config.lora_dropout,
"fan_in_fan_out": self.peft_config.fan_in_fan_out,
"merge_weights": self.peft_config.merge_weights or self.peft_config.inference_mode,
"r": lora_config.r,
"lora_alpha": lora_config.lora_alpha,
"lora_dropout": lora_config.lora_dropout,
"fan_in_fan_out": lora_config.fan_in_fan_out,
"init_lora_weights": lora_config.init_lora_weights,
}
key_list = [key for key, _ in self.model.named_modules()]
for key in key_list:
if isinstance(self.peft_config.target_modules, str):
target_module_found = re.fullmatch(self.peft_config.target_modules, key)
if isinstance(lora_config.target_modules, str):
target_module_found = re.fullmatch(lora_config.target_modules, key)
else:
target_module_found = any(key.endswith(target_key) for target_key in self.peft_config.target_modules)
target_module_found = any(key.endswith(target_key) for target_key in lora_config.target_modules)
if target_module_found:
if not is_target_modules_in_base_model:
is_target_modules_in_base_model = True
parent, target, target_name = self._get_submodules(key)
parent, target, target_name = _get_submodules(self.model, key)
bias = target.bias is not None
if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt):
kwargs.update(
{
"has_fp16_weights": target.state.has_fp16_weights,
"memory_efficient_backward": target.state.memory_efficient_backward,
"threshold": target.state.threshold,
"index": target.index,
}
if isinstance(target, LoraLayer):
target.update_layer(
adapter_name,
lora_config.r,
lora_config.lora_alpha,
lora_config.lora_dropout,
lora_config.init_lora_weights,
)
new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs)
elif isinstance(target, torch.nn.Linear):
new_module = SVDLinear(target.in_features, target.out_features, bias=bias, **kwargs)
self._replace_module(parent, target_name, new_module, target)
else:
if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt):
kwargs.update(
{
"has_fp16_weights": target.state.has_fp16_weights,
"memory_efficient_backward": target.state.memory_efficient_backward,
"threshold": target.state.threshold,
"index": target.index,
}
)
new_module = SVDLinear8bitLt(
adapter_name, target.in_features, target.out_features, bias=bias, **kwargs
)
else:
if isinstance(target, torch.nn.Linear):
in_features, out_features = target.in_features, target.out_features
if kwargs["fan_in_fan_out"]:
warnings.warn(
"fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. "
"Setting fan_in_fan_out to False."
)
kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = False
elif isinstance(target, Conv1D):
in_features, out_features = (
target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape
)
if not kwargs["fan_in_fan_out"]:
warnings.warn(
"fan_in_fan_out is set to False but the target module is `Conv1D`. "
"Setting fan_in_fan_out to True."
)
kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = True
else:
raise ValueError(
f"Target module {target} is not supported. "
f"Currently, only `torch.nn.Linear` and `Conv1D` are supported."
)
new_module = SVDLinear(adapter_name, in_features, out_features, bias=bias, **kwargs)
self._replace_module(parent, target_name, new_module, target)
if not is_target_modules_in_base_model:
raise ValueError(
f"Target modules {self.peft_config.target_modules} not found in the base model. "
f"Target modules {lora_config.target_modules} not found in the base model. "
f"Please check the target modules and try again."
)
@@ -144,14 +217,14 @@ class AdaLoraModel(LoraModel):
outputs = self.model.forward(*args, **kwargs)
# Calculate the orthogonal regularization
orth_reg_weight = self.peft_config.orth_reg_weight
orth_reg_weight = self.peft_config[self.trainable_adapter_name].orth_reg_weight
assert orth_reg_weight > 0
if hasattr(outputs, "loss"):
regu_loss = 0
num_param = 0
for n, p in self.model.named_parameters():
if "lora_A" in n or "lora_B" in n:
if ("lora_A" in n or "lora_B" in n) and self.trainable_adapter_name in n:
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
@@ -161,7 +234,7 @@ class AdaLoraModel(LoraModel):
outputs.loss += orth_reg_weight * regu_loss
return outputs
def _prepare_new_module(self, target, rank_idx):
def _prepare_new_module(self, target, rank_idx, adapter_name):
if isinstance(rank_idx, list):
rank = sum(rank_idx)
elif isinstance(rank_idx, torch.Tensor):
@@ -169,12 +242,14 @@ class AdaLoraModel(LoraModel):
rank = rank_idx.sum().item()
else:
raise ValueError("Unexcepted type of rank_idx")
lora_config = self.peft_config[adapter_name]
kwargs = {
"r": rank,
"lora_alpha": self.peft_config.lora_alpha,
"lora_dropout": self.peft_config.lora_dropout,
"fan_in_fan_out": self.peft_config.fan_in_fan_out,
"merge_weights": self.peft_config.merge_weights or self.peft_config.inference_mode,
"lora_alpha": lora_config.lora_alpha,
"lora_dropout": lora_config.lora_dropout,
"fan_in_fan_out": lora_config.fan_in_fan_out,
"init_lora_weights": lora_config.init_lora_weights,
}
bias = target.bias is not None
loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False)
@@ -187,9 +262,9 @@ class AdaLoraModel(LoraModel):
"index": target.index,
}
)
new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs)
new_module = SVDLinear8bitLt(adapter_name, target.in_features, target.out_features, bias=bias, **kwargs)
elif isinstance(target, torch.nn.Linear):
new_module = SVDLinear(target.in_features, target.out_features, bias=bias, **kwargs)
new_module = SVDLinear(adapter_name, target.in_features, target.out_features, bias=bias, **kwargs)
new_module = new_module.to(target.weight.device)
with torch.no_grad():
@@ -197,120 +272,195 @@ class AdaLoraModel(LoraModel):
if bias:
new_module.bias.copy_(target.bias)
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])
new_module.lora_E[adapter_name].copy_(target.lora_E[rank_idx])
new_module.lora_A[adapter_name].copy_(target.lora_A[rank_idx])
new_module.lora_B[adapter_name].copy_(target.lora_B[:, rank_idx])
# The scaling is exactly as the previous
new_module.ranknum.copy_(target.ranknum)
new_module.ranknum[adapter_name].copy_(target.ranknum)
return new_module
def resize_modules_by_rank_pattern(self, rank_pattern):
def resize_modules_by_rank_pattern(self, rank_pattern, adapter_name):
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)
key = f"{key}.{adapter_name}" if adapter_name not in key else key
parent, target, target_name = _get_submodules(key)
new_module = self._prepare_new_module(target, rank_idx, adapter_name)
self._replace_module(parent, target_name, new_module, target)
def update_and_allocate(self, global_step):
lora_config = self.peft_config[self.trainable_adapter_name]
# 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 global_step < lora_config.total_step - lora_config.tfinal:
_, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step)
if rank_pattern:
self.peft_config.rank_pattern = rank_pattern
lora_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
elif global_step == lora_config.total_step - lora_config.tfinal:
_, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True)
self.resize_modules_by_rank_pattern(rank_pattern, self.trainable_adapter_name)
lora_config.rank_pattern = rank_pattern
self.rankallocator.reset_ipt()
# Pass the function and do forward propagation
else:
return None
@staticmethod
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:
raise ValueError("Please specify `target_modules` in `peft_config`")
peft_config.target_modules = TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING[
model_config["model_type"]
]
if peft_config.inference_mode:
peft_config.merge_weights = True
return peft_config
class SVDLinear(nn.Linear, LoraLayer):
class AdaLoraLayer(LoraLayer):
def __init__(
self,
in_features: int,
out_features: int,
):
super().__init__(in_features, out_features)
self.lora_E = nn.ParameterDict({})
self.lora_A = nn.ParameterDict({})
self.lora_B = nn.ParameterDict({})
self.ranknum = nn.ParameterDict({})
def update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
self.r[adapter_name] = r
self.lora_alpha[adapter_name] = lora_alpha
if lora_dropout > 0.0:
lora_dropout_layer = nn.Dropout(p=lora_dropout)
else:
def lora_dropout_layer(x):
return x
self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer}))
# Actual trainable parameters
if r > 0:
# Right singular vectors
self.lora_A.update(
nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros((r, self.in_features)))})
)
# Singular values
self.lora_E.update(nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros(r, 1))}))
# Left singular vectors
self.lora_B.update(
nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros((self.out_features, r)))})
)
# The current rank
self.ranknum.update(
nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros(1), requires_grad=False)})
)
self.ranknum[adapter_name].data.fill_(float(self.r))
self.ranknum[adapter_name].requires_grad = False
self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r)
if init_lora_weights:
self.reset_lora_parameters(adapter_name)
self.to(self.weight.device)
def reset_lora_parameters(self, adapter_name):
if adapter_name in self.lora_A.keys():
nn.init.zeros_(self.lora_E[adapter_name])
nn.init.normal_(self.lora_A[adapter_name], mean=0.0, std=0.02)
nn.init.normal_(self.lora_B[adapter_name], mean=0.0, std=0.02)
class SVDLinear(nn.Linear, AdaLoraLayer):
# SVD-based adaptation by a dense layer
def __init__(
self,
adapter_name: str,
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,
):
init_lora_weights = kwargs.pop("init_lora_weights", True)
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)
AdaLoraLayer.__init__(self, in_features=in_features, out_features=out_features)
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
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))
# 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)
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
self.ranknum.requires_grad = False
self.reset_parameters()
if fan_in_fan_out:
self.weight.data = self.weight.data.T
def reset_parameters(self):
nn.Linear.reset_parameters(self)
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)
self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
self.active_adapter = adapter_name
def train(self, mode: bool = True):
nn.Linear.train(self, mode)
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.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)
def merge(self):
if self.active_adapter not in self.lora_A.keys():
return
if self.merged:
warnings.warn("Already merged. Nothing to do.")
return
if self.r[self.active_adapter] > 0:
self.weight.data += (
transpose(
self.lora_B[self.active_adapter]
@ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter])
)
* self.scaling[self.active_adapter]
/ (self.ranknum[self.active_adapter] + 1e-5)
)
self.merged = True
def forward(self, x: torch.Tensor):
if self.r > 0 and not self.merged:
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)
def unmerge(self):
if self.active_adapter not in self.lora_A.keys():
return
if not self.merged:
warnings.warn("Already unmerged. Nothing to do.")
return
if self.r[self.active_adapter] > 0:
self.weight.data -= (
transpose(
self.lora_B[self.active_adapter]
@ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter])
)
return result
else:
* self.scaling[self.active_adapter]
/ (self.ranknum[self.active_adapter] + 1e-5)
)
self.merged = False
def forward(self, x: torch.Tensor):
if self.active_adapter not in self.lora_A.keys():
return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
if self.disable_adapters:
if self.r[self.active_adapter] > 0 and self.merged:
self.unmerge()
result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
elif self.r[self.active_adapter] > 0 and not self.merged:
result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
result += (
(
self.lora_dropout[self.active_adapter](x)
@ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]).T
@ self.lora_B[self.active_adapter].T
)
* self.scaling[self.active_adapter]
/ (self.ranknum[self.active_adapter] + 1e-5)
)
else:
result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
return result
if is_bnb_available():
class SVDLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer):
class SVDLinear8bitLt(bnb.nn.Linear8bitLt, AdaLoraLayer):
# Low-rank matrix for SVD-based adaptation
def __init__(
self,
adapter_name,
in_features,
out_features,
r: int = 0,
@@ -328,51 +478,45 @@ if is_bnb_available():
threshold=kwargs.get("threshold", 0.0),
index=kwargs.get("index", None),
)
LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False)
# 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))
# 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)
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
self.ranknum.requires_grad = False
self.reset_parameters()
AdaLoraLayer.__init__(self, in_features=in_features, out_features=out_features)
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
def reset_parameters(self):
if hasattr(self, "lora_A"):
# initialize A the same way as the default for nn.Linear and B to zero
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)
init_lora_weights = kwargs.pop("init_lora_weights", True)
self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
self.active_adapter = adapter_name
def forward(self, x: torch.Tensor):
result = super().forward(x)
if self.disable_adapters:
if self.disable_adapters or self.active_adapter not in self.lora_A.keys():
return result
elif self.r > 0:
elif self.r[self.active_adapter] > 0:
if not torch.is_autocast_enabled():
expected_dtype = result.dtype
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
result += output
(
self.lora_dropout[self.active_adapter](x)
@ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]).T
@ self.lora_B[self.active_adapter].T
).to(expected_dtype)
* self.scaling[self.active_adapter]
/ (self.ranknum[self.active_adapter] + 1e-5)
)
else:
output = (
self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T / (self.ranknum + 1e-5)
) * self.scaling
result += output
(
self.lora_dropout[self.active_adapter](x)
@ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]).T
@ self.lora_B[self.active_adapter].T
)
* self.scaling[self.active_adapter]
/ (self.ranknum[self.active_adapter] + 1e-5)
)
result += output
return result
@@ -386,8 +530,9 @@ class RankAllocator(object):
"""
def __init__(self, peft_config, model):
def __init__(self, model, peft_config, adapter_name):
self.peft_config = peft_config
self.adapter_name = adapter_name
self.beta1 = peft_config.beta1
self.beta2 = peft_config.beta2
assert self.beta1 > 0 and self.beta1 < 1
@@ -408,7 +553,7 @@ class RankAllocator(object):
self.init_bgt = 0
self.name_set = set()
for n, p in model.named_parameters():
if "lora_A" in n:
if f"lora_A.{self.adapter_name}" in n:
self.init_bgt += p.size(0)
self.name_set.add(n.replace("lora_A", "%s"))
self.name_set = sorted(self.name_set)
@@ -437,7 +582,7 @@ class RankAllocator(object):
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 "lora_" in n and self.adapter_name in n:
if n not in self.ipt:
self.ipt[n] = torch.zeros_like(p)
self.exp_avg_ipt[n] = torch.zeros_like(p)
@@ -465,7 +610,7 @@ class RankAllocator(object):
triplet_ipt = {}
# Get the importance score for A, E, B
for n, p in model.named_parameters():
if "lora_A" in n:
if f"lora_A.{self.adapter_name}" 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")
@@ -473,7 +618,7 @@ class RankAllocator(object):
vector_ipt[name_m] = [comb_ipt]
else:
vector_ipt[name_m].append(comb_ipt)
if "lora_B" in n:
if f"lora_B.{self.adapter_name}" 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")
@@ -481,7 +626,7 @@ class RankAllocator(object):
vector_ipt[name_m] = [comb_ipt]
else:
vector_ipt[name_m].append(comb_ipt)
if "lora_E" in n:
if f"lora_E.{self.adapter_name}" in n:
entry_ipt = self._element_score(n)
name_m = n.replace("lora_E", "%s")
value_ipt[name_m] = entry_ipt
@@ -506,7 +651,7 @@ class RankAllocator(object):
# Mask the unimportant triplets
with torch.no_grad():
for n, p in model.named_parameters():
if "lora_E" in n:
if f"lora_E.{self.adapter_name}" 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
src/peft/tuners/lora.py
+7 -20
View File
@@ -29,6 +29,7 @@ from ..utils import (
TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING,
PeftConfig,
PeftType,
_freeze_adapter,
_get_submodules,
transpose,
)
@@ -52,8 +53,6 @@ class LoraConfig(PeftConfig):
target_modules (`Union[List[str],str]`): The names of the modules to apply Lora to.
lora_alpha (`float`): The alpha parameter for Lora scaling.
lora_dropout (`float`): The dropout probability for Lora layers.
merge_weights (`bool`):
Whether to merge the weights of the Lora layers with the base transformer model in `eval` mode.
fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out).
For example, gpt-2 uses `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`.:
bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only'
@@ -71,9 +70,6 @@ class LoraConfig(PeftConfig):
)
lora_alpha: int = field(default=None, metadata={"help": "Lora alpha"})
lora_dropout: float = field(default=None, metadata={"help": "Lora dropout"})
merge_weights: bool = field(
default=False, metadata={"help": "Merge weights of the original model and the Lora model"}
)
fan_in_fan_out: bool = field(
default=False,
metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"},
@@ -147,7 +143,10 @@ class LoraModel(torch.nn.Module):
raise ValueError(
"LoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters."
)
mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias)
if self.peft_config[adapter_name].inference_mode:
_freeze_adapter(self.model, adapter_name)
else:
mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias)
def _find_and_replace(self, adapter_name):
lora_config = self.peft_config[adapter_name]
@@ -158,14 +157,11 @@ class LoraModel(torch.nn.Module):
"You can install it with `pip install bitsandbytes`."
)
is_target_modules_in_base_model = False
is_hf_device_map_available = hasattr(self.model, "hf_device_map")
kwargs = {
"r": lora_config.r,
"lora_alpha": lora_config.lora_alpha,
"lora_dropout": lora_config.lora_dropout,
"fan_in_fan_out": lora_config.fan_in_fan_out,
"merge_weights": (lora_config.merge_weights or lora_config.inference_mode)
and not is_hf_device_map_available,
"init_lora_weights": lora_config.init_lora_weights,
}
key_list = [key for key, _ in self.model.named_modules()]
@@ -360,7 +356,6 @@ def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None:
class LoraLayer:
def __init__(
self,
merge_weights: bool,
in_features: int,
out_features: int,
):
@@ -372,7 +367,6 @@ class LoraLayer:
self.lora_B = nn.ModuleDict({})
# Mark the weight as unmerged
self.merged = False
self.merge_weights = merge_weights
self.disable_adapters = False
self.in_features = in_features
self.out_features = out_features
@@ -415,13 +409,12 @@ class Linear(nn.Linear, LoraLayer):
lora_alpha: int = 1,
lora_dropout: float = 0.0,
fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
merge_weights: bool = True,
**kwargs,
):
init_lora_weights = kwargs.pop("init_lora_weights", True)
nn.Linear.__init__(self, in_features, out_features, **kwargs)
LoraLayer.__init__(self, merge_weights=merge_weights, in_features=in_features, out_features=out_features)
LoraLayer.__init__(self, in_features=in_features, out_features=out_features)
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
@@ -436,9 +429,6 @@ class Linear(nn.Linear, LoraLayer):
def merge(self):
if self.active_adapter not in self.lora_A.keys():
return
if not self.merge_weights:
warnings.warn("Nothing to merge. Set merge_weights to True to enable merging.")
return
if self.merged:
warnings.warn("Already merged. Nothing to do.")
return
@@ -455,9 +445,6 @@ class Linear(nn.Linear, LoraLayer):
def unmerge(self):
if self.active_adapter not in self.lora_A.keys():
return
if not self.merge_weights:
warnings.warn("Nothing to unmerge. Set merge_weights to True to enable (un)merging.")
return
if not self.merged:
warnings.warn("Already unmerged. Nothing to do.")
return
@@ -515,7 +502,7 @@ if is_bnb_available():
threshold=kwargs.get("threshold", 0.0),
index=kwargs.get("index", None),
)
LoraLayer.__init__(self, merge_weights=False, in_features=in_features, out_features=out_features)
LoraLayer.__init__(self, in_features=in_features, out_features=out_features)
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
src/peft/utils/__init__.py
+2
View File
@@ -21,6 +21,7 @@ from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType
from .other import (
TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING,
TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING,
TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING,
CONFIG_NAME,
WEIGHTS_NAME,
_set_trainable,
@@ -30,5 +31,6 @@ from .other import (
transpose,
_get_submodules,
_set_adapter,
_freeze_adapter,
)
from .save_and_load import get_peft_model_state_dict, set_peft_model_state_dict
src/peft/utils/other.py
+26
View File
@@ -134,6 +134,12 @@ def _get_submodules(model, key):
return parent, target, target_name
def _freeze_adapter(model, adapter_name):
for n, p in model.named_parameters():
if adapter_name in n:
p.requires_grad = False
def _set_trainable(model, adapter_name):
key_list = [key for key, _ in model.named_modules()]
for key in key_list:
@@ -199,6 +205,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
"bart": ["q_proj", "v_proj"],
"gpt2": ["c_attn"],
"bloom": ["query_key_value"],
"blip-2": ["q", "v", "q_proj", "v_proj"],
"opt": ["q_proj", "v_proj"],
"gptj": ["q_proj", "v_proj"],
"gpt_neox": ["query_key_value"],
@@ -214,6 +221,25 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
"chatglm": ["query_key_value"],
}
TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING = {
"t5": ["q", "k", "v", "o", "wi", "wo"],
"mt5": ["q", "k", "v", "o", "wi_0", "wi_1", "wo"],
"bart": ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"],
# "gpt2": ["c_attn"],
# "bloom": ["query_key_value"],
"opt": ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"],
# "gptj": ["q_proj", "v_proj"],
# "gpt_neox": ["query_key_value"],
# "gpt_neo": ["q_proj", "v_proj"],
# "bert": ["query", "value"],
"roberta": ["query", "key", "value", "dense"],
# "xlm-roberta": ["query", "value"],
# "electra": ["query", "value"],
"deberta-v2": ["query_proj", "key_proj", "value_proj", "dense"],
# "deberta": ["in_proj"],
# "layoutlm": ["query", "value"],
}
TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = {
"bloom": bloom_model_postprocess_past_key_value,
}
tests/test_decoder_models.py
+1 -1
View File
@@ -72,7 +72,7 @@ class PeftDecoderModelTester(unittest.TestCase, PeftCommonTester):
PeftTestConfigManager.get_grid_parameters(
{
"model_ids": PEFT_DECODER_MODELS_TO_TEST,
"lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]},
"lora_kwargs": {"init_lora_weights": [False]},
"task_type": "CAUSAL_LM",
},
)
tests/test_encoder_decoder_models.py
+1 -1
View File
@@ -74,7 +74,7 @@ class PeftEncoderDecoderModelTester(unittest.TestCase, PeftCommonTester):
PeftTestConfigManager.get_grid_parameters(
{
"model_ids": PEFT_ENCODER_DECODER_MODELS_TO_TEST,
"lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]},
"lora_kwargs": {"init_lora_weights": [False]},
"task_type": "SEQ_2_SEQ_LM",
},
)
tests/testing_common.py
+2 -6
View File
@@ -268,12 +268,8 @@ class PeftCommonTester:
logits_transformers = transformers_model(**dummy_input)[0]
if config.merge_weights:
self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4))
self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10))
else:
self.assertFalse(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4))
self.assertTrue(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10))
self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4))
self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10))
with tempfile.TemporaryDirectory() as tmp_dirname:
model.save_pretrained(tmp_dirname)