diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py new file mode 100644 index 0000000..31b4aa5 --- /dev/null +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -0,0 +1,182 @@ +import os + +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" +tokenizer_name_or_path = "facebook/bart-base" + +checkpoint_name = "financial_sentiment_analysis_lora_v1.pt" +text_column = "sentence" +label_column = "text_label" +max_length = 128 +lr = 1e-3 +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, +) + +model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) +model = get_peft_model(model, peft_config) +model.print_trainable_parameters() + + +# loading dataset +dataset = load_dataset("financial_phrasebank", "sentences_allagree") +dataset = dataset["train"].train_test_split(test_size=0.1) +dataset["validation"] = dataset["test"] +del dataset["test"] + +classes = dataset["train"].features["label"].names +dataset = dataset.map( + lambda x: {"text_label": [classes[label] for label in x["label"]]}, + batched=True, + num_proc=1, +) + + +# data preprocessing +tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) + + +def preprocess_function(examples): + inputs = examples[text_column] + targets = examples[label_column] + model_inputs = tokenizer(inputs, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt") + labels = tokenizer(targets, max_length=3, padding="max_length", truncation=True, return_tensors="pt") + labels = labels["input_ids"] + labels[labels == tokenizer.pad_token_id] = -100 + model_inputs["labels"] = labels + return model_inputs + + +processed_datasets = dataset.map( + preprocess_function, + batched=True, + num_proc=1, + remove_columns=dataset["train"].column_names, + load_from_cache_file=False, + desc="Running tokenizer on dataset", +) + +train_dataset = processed_datasets["train"] +eval_dataset = processed_datasets["validation"] + +train_dataloader = DataLoader( + train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True +) +eval_dataloader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True) + + +# optimizer and lr scheduler +optimizer = torch.optim.AdamW(model.parameters(), lr=lr) +lr_scheduler = get_linear_schedule_with_warmup( + optimizer=optimizer, + num_warmup_steps=0, + num_training_steps=(len(train_dataloader) * num_epochs), +) +model.base_model.peft_config.total_step = len(train_dataloader) * num_epochs + + +# training and evaluation +model = model.to(device) +global_step = 0 +for epoch in range(num_epochs): + model.train() + total_loss = 0 + for step, batch in enumerate(tqdm(train_dataloader)): + batch = {k: v.to(device) for k, v in batch.items()} + outputs = model(**batch) + loss = outputs.loss + total_loss += loss.detach().float() + loss.backward() + optimizer.step() + lr_scheduler.step() + # 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 + + model.eval() + eval_loss = 0 + eval_preds = [] + for step, batch in enumerate(tqdm(eval_dataloader)): + batch = {k: v.to(device) for k, v in batch.items()} + with torch.no_grad(): + outputs = model(**batch) + loss = outputs.loss + eval_loss += loss.detach().float() + eval_preds.extend( + tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(), skip_special_tokens=True) + ) + + eval_epoch_loss = eval_loss / len(train_dataloader) + eval_ppl = torch.exp(eval_epoch_loss) + train_epoch_loss = total_loss / len(eval_dataloader) + train_ppl = torch.exp(train_epoch_loss) + print(f"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}") + + +# print accuracy +correct = 0 +total = 0 +for pred, true in zip(eval_preds, dataset["validation"]["text_label"]): + if pred.strip() == true.strip(): + correct += 1 + total += 1 +accuracy = correct / total * 100 +print(f"{accuracy=} % on the evaluation dataset") +print(f"{eval_preds[:10]=}") +print(f"{dataset['validation']['text_label'][:10]=}") + + +# saving model +peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}" +model.save_pretrained(peft_model_id) + + +ckpt = f"{peft_model_id}/adapter_model.bin" +# get_ipython().system('du -h $ckpt') + + +peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}" + +config = PeftConfig.from_pretrained(peft_model_id) +model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path) +model = PeftModel.from_pretrained(model, peft_model_id) + + +model.eval() +i = 13 +inputs = tokenizer(dataset["validation"][text_column][i], return_tensors="pt") +print(dataset["validation"][text_column][i]) +print(inputs) + +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/examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py b/examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py index 0e47f87..c6b9130 100644 --- a/examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py +++ b/examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py @@ -102,7 +102,8 @@ class TorchTracemalloc: def main(): accelerator = Accelerator() - model_name_or_path = "bigscience/T0_3B" + # model_name_or_path = "bigscience/T0_3B" + model_name_or_path = "facebook/bart-large" dataset_name = "twitter_complaints" peft_config = LoraConfig( task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1 diff --git a/src/peft/__init__.py b/src/peft/__init__.py index e141347..dc666ce 100644 --- a/src/peft/__init__.py +++ b/src/peft/__init__.py @@ -30,6 +30,8 @@ from .peft_model import ( from .tuners import ( LoraConfig, LoraModel, + AdaLoraConfig, + AdaLoraModel, PrefixEncoder, PrefixTuningConfig, PromptEmbedding, diff --git a/src/peft/mapping.py b/src/peft/mapping.py index c814655..baa05d3 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, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig +from .tuners import AdaLoraConfig, LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from .utils import PromptLearningConfig @@ -36,6 +36,7 @@ PEFT_TYPE_TO_CONFIG_MAPPING = { "PREFIX_TUNING": PrefixTuningConfig, "P_TUNING": PromptEncoderConfig, "LORA": LoraConfig, + "ADALORA": AdaLoraConfig, } diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index fbe98e4..2b7b286 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -28,7 +28,7 @@ from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput from transformers.utils import PushToHubMixin -from .tuners import LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder +from .tuners import AdaLoraModel, LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder from .utils import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, WEIGHTS_NAME, @@ -49,6 +49,7 @@ PEFT_TYPE_TO_MODEL_MAPPING = { PeftType.PROMPT_TUNING: PromptEmbedding, PeftType.P_TUNING: PromptEncoder, PeftType.PREFIX_TUNING: PrefixEncoder, + PeftType.ADALORA: AdaLoraModel, } @@ -81,7 +82,6 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.base_model = model self.config = self.base_model.config self.modules_to_save = None - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.peft_config = {} self.active_adapter = adapter_name self.peft_type = peft_config.peft_type @@ -320,7 +320,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.modules_to_save = self.modules_to_save.update(peft_config.modules_to_save) _set_trainable(self, adapter_name) - def load_adapter(self, model_id, adapter_name, **kwargs): + def load_adapter(self, model_id, adapter_name, is_trainable=False, **kwargs): from .mapping import PEFT_TYPE_TO_CONFIG_MAPPING if adapter_name not in self.peft_config: @@ -328,6 +328,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): peft_config = PEFT_TYPE_TO_CONFIG_MAPPING[ PeftConfig.from_pretrained(model_id, subfolder=kwargs.get("subfolder", None)).peft_type ].from_pretrained(model_id, subfolder=kwargs.get("subfolder", None)) + if isinstance(peft_config, PromptLearningConfig) and is_trainable: + raise ValueError("Cannot set a prompt learning adapter to trainable when loading pretrained adapter.") + else: + peft_config[adapter_name].inference_mode = not is_trainable self.add_adapter(adapter_name, peft_config) # load weights if any @@ -389,6 +393,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): remove_hook_from_submodules(self.prompt_encoder) add_hook_to_module(self.get_base_model(), hook) + # Set model in evaluation mode to deactivate Dropout modules by default + self.eval() + def set_adapter(self, adapter_name): """ Sets the active adapter. diff --git a/src/peft/tuners/__init__.py b/src/peft/tuners/__init__.py index 8f93079..146366b 100644 --- a/src/peft/tuners/__init__.py +++ b/src/peft/tuners/__init__.py @@ -17,7 +17,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .lora import LoraConfig, LoraModel +from .adalora import AdaLoraConfig, AdaLoraModel from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType from .prefix_tuning import PrefixEncoder, PrefixTuningConfig from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit -from .lora import LoraConfig, LoraModel diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py new file mode 100644 index 0000000..f98b3c6 --- /dev/null +++ b/src/peft/tuners/adalora.py @@ -0,0 +1,524 @@ +import importlib +import re +from dataclasses import dataclass, field +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ..utils import PeftType, transpose +from .lora import LoraConfig, LoraLayer, LoraModel, mark_only_lora_as_trainable + + +def is_bnb_available(): + return importlib.util.find_spec("bitsandbytes") is not None + + +if is_bnb_available(): + import bitsandbytes as bnb + + +@dataclass +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. + 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. + """ + + 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."}) + + 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 + + Args: + model ([`transformers.PreTrainedModel`]): The model to be adapted. + config ([`AdaLoraConfig`]): The configuration of the AdaLora model. + + Returns: + `torch.nn.Module`: The AdaLora model. + + Example:: + + >>> 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) + + **Attributes**: + - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. + - **peft_config** ([`AdaLoraConfig`]): The configuration of the AdaLora model. + """ + + def __init__(self, config, model): + 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.rankallocator = RankAllocator(config, self.model) + if config.enable_lora is not None: + 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) + if loaded_in_8bit and not is_bnb_available(): + raise ImportError( + "To use Lora with 8-bit quantization, please install the `bitsandbytes` package. " + "You can install it with `pip install bitsandbytes`." + ) + 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, + } + 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) + else: + target_module_found = any(key.endswith(target_key) for target_key in self.peft_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) + 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, + } + ) + 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) + 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"Please check the target modules and try again." + ) + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + return getattr(self.model, name) + + def forward(self, *args, **kwargs): + outputs = self.model.forward(*args, **kwargs) + + # Calculate the orthogonal regularization + orth_reg_weight = self.peft_config.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: + 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 + + def _prepare_new_module(self, target, rank_idx): + if isinstance(rank_idx, list): + rank = sum(rank_idx) + elif isinstance(rank_idx, torch.Tensor): + rank_idx = rank_idx.view(-1) + rank = rank_idx.sum().item() + else: + raise ValueError("Unexcepted type of rank_idx") + 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, + } + bias = target.bias is not None + loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) + 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(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 = new_module.to(target.weight.device) + + with torch.no_grad(): + new_module.weight.copy_(target.weight) + 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]) + # 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) + 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 + 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.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: + 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.0, + fan_in_fan_out: bool = False, + merge_weights: bool = True, + **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) + + 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) + + 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) + ) + 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) + ) + 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 + def __init__( + self, + in_features, + out_features, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + **kwargs, + ): + bnb.nn.Linear8bitLt.__init__( + self, + in_features, + out_features, + bias=kwargs.get("bias", True), + has_fp16_weights=kwargs.get("has_fp16_weights", True), + memory_efficient_backward=kwargs.get("memory_efficient_backward", False), + 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() + + 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) + + def forward(self, x: torch.Tensor): + result = super().forward(x) + + if self.disable_adapters: + return result + elif self.r > 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 + 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 + return result + + +class RankAllocator(object): + """ + 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. + + """ + + 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.reset_ipt() + 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, 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.name_set.add(n.replace("lora_A", "%s")) + 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 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) + 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() + ) + + 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 = {} + triplet_ipt = {} + # Get the importance score for A, E, B + 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: + vector_ipt[name_m] = [comb_ipt] + else: + vector_ipt[name_m].append(comb_ipt) + 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: + 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) + 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] + 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" + 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, + )[0].item() + + rank_pattern = {} + # 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() + return rank_pattern + + def update_and_allocate(self, model, global_step, force_mask=False): + # # 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 + if mask_ind or force_mask: + rank_pattern = self.mask_to_budget(model, budget) + else: + rank_pattern = None + return budget, rank_pattern diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index 34e98e9..31c4669 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -29,6 +29,7 @@ class PeftType(str, enum.Enum): P_TUNING = "P_TUNING" PREFIX_TUNING = "PREFIX_TUNING" LORA = "LORA" + ADALORA = "ADALORA" class TaskType(str, enum.Enum): diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index a258c32..dc0a391 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -30,7 +30,7 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): config = model.peft_config[adapter_name] if state_dict is None: state_dict = model.state_dict() - if config.peft_type == PeftType.LORA: + if config.peft_type in (PeftType.LORA, PeftType.ADALORA): # to_return = lora_state_dict(model, bias=model.peft_config.bias) # adapted from `https://github.com/microsoft/LoRA/blob/main/loralib/utils.py` # to be used directly with the state dict which is necessary when using DeepSpeed or FSDP @@ -49,6 +49,12 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): to_return[bias_name] = state_dict[bias_name] else: raise NotImplementedError + if config.peft_type == PeftType.ADALORA: + rank_pattern = config.rank_pattern + if rank_pattern is not None: + rank_pattern = {k.replace(f"{adapter_name}.", ""): v for k, v in rank_pattern.items()} + config.rank_pattern = rank_pattern + to_return = {k: v for k, v in to_return.items() if (("lora_" in k and adapter_name in k) or ("bias" in k))} elif isinstance(config, PromptLearningConfig): to_return = {} @@ -89,7 +95,7 @@ def set_peft_model_state_dict(model, peft_model_state_dict, adapter_name="defaul else: state_dict = peft_model_state_dict - if config.peft_type == PeftType.LORA: + if config.peft_type in (PeftType.LORA, PeftType.ADALORA): peft_model_state_dict = {} for k, v in state_dict.items(): if "lora_" in k: @@ -98,10 +104,15 @@ def set_peft_model_state_dict(model, peft_model_state_dict, adapter_name="defaul peft_model_state_dict[k] = v else: peft_model_state_dict[k] = v + if config.peft_type == PeftType.ADALORA: + rank_pattern = config.rank_pattern + if rank_pattern is not None: + model.resize_modules_by_rank_pattern(rank_pattern, adapter_name) elif isinstance(config, PromptLearningConfig): peft_model_state_dict = state_dict else: raise NotImplementedError + model.load_state_dict(peft_model_state_dict, strict=False) if isinstance(config, PromptLearningConfig): model.prompt_encoder[adapter_name].embedding.load_state_dict(