From 81eec9ba70e2b6f754350bf91cbb265bc9d2b99e Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Mon, 27 Feb 2023 21:08:55 -0500 Subject: [PATCH 01/76] train script --- .../peft_lora_seq2seq_accelerate_ds_zero3_offload.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 cef9773..3c80bfb 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 From 26b84e6fd9081b2e022690451b8cc3ed59c2f386 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Tue, 28 Feb 2023 23:14:25 -0500 Subject: [PATCH 02/76] add adalora example --- .../peft_adalora_seq2seq.py | 172 ++++++++++++++++++ src/peft/utils/config.py | 1 + 2 files changed, 173 insertions(+) create mode 100644 examples/conditional_generation/peft_adalora_seq2seq.py diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py new file mode 100644 index 0000000..19875c4 --- /dev/null +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -0,0 +1,172 @@ +from transformers import AutoModelForSeq2SeqLM +from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, 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 +from datasets import load_dataset + +device = "cuda" +model_name_or_path = "bigscience/mt0-large" +tokenizer_name_or_path = "bigscience/mt0-large" + +checkpoint_name = "financial_sentiment_analysis_lora_v1.pt" +text_column = "sentence" +label_column = "text_label" +max_length = 128 +lr = 1e-3 +num_epochs = 3 +batch_size = 8 + + +# creating model +peft_config = LoraConfig(task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1) + +model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) +model = get_peft_model(model, peft_config) +model.print_trainable_parameters() +model + + +# 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, +) + +dataset["train"][0] + + +# 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), +) + + +# training and evaluation +model = model.to(device) + +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() + optimizer.zero_grad() + + 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') + + +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) +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/src/peft/utils/config.py b/src/peft/utils/config.py index f0587fe..5697da6 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -30,6 +30,7 @@ class PeftType(str, enum.Enum): P_TUNING = "P_TUNING" PREFIX_TUNING = "PREFIX_TUNING" LORA = "LORA" + ADALORA = "ADALORA" class TaskType(str, enum.Enum): From be86f904907e849a45b9acb7d8f975a59d523237 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Tue, 28 Feb 2023 23:18:19 -0500 Subject: [PATCH 03/76] Implement the AdaLoRA --- src/peft/tuners/__init__.py | 1 + src/peft/tuners/adalora.py | 497 ++++++++++++++++++++++++++++++++++++ 2 files changed, 498 insertions(+) create mode 100644 src/peft/tuners/adalora.py diff --git a/src/peft/tuners/__init__.py b/src/peft/tuners/__init__.py index 38b7926..146366b 100644 --- a/src/peft/tuners/__init__.py +++ b/src/peft/tuners/__init__.py @@ -18,6 +18,7 @@ # 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 diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py new file mode 100644 index 0000000..e8d5f86 --- /dev/null +++ b/src/peft/tuners/adalora.py @@ -0,0 +1,497 @@ +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 + +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 + + +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: + r (`int`): Lora attention dimension + 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) + enable_lora ( `List[bool]`): Used with `lora.MergedLinear`. + bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only' + modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable + and saved in the final checkpoint. + """ + 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."} + ) + + def __post_init__(self): + self.peft_type = PeftType.ADALORA + + + +class AdaLoraModel(LoraModel): + """ + Creates Adaptive LoRA (AdaLora) model from a pretrained transformers model. + + Args: + model ([`transformers.PreTrainedModel`]): The model to be adapted. + config ([`LoraConfig`]): The configuration of the Lora model. + + Returns: + `torch.nn.Module`: The Lora model. + + Example:: + + >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import LoraModel, LoraConfig >>> + config = LoraConfig( + peft_type="LORA", 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") >>> lora_model = LoraModel(config, model) + + **Attributes**: + - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. + - **peft_config** ([`LoraConfig`]): The configuration of the Lora model. + """ + + def __init__(self, config, model): + # super().__init__() + 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.forward = self.model.forward + self.rankallocator = RankAllocator(config, self.named_parameters()) + + 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) and self.peft_config.enable_lora is None: + 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) and self.peft_config.enable_lora is None: + new_module = SVDLinear(target.in_features, target.out_features, bias=bias, **kwargs) + # TODO: Implement the MergedLinear of SVD Adapattion + # elif self.peft_config.enable_lora is not None: + # kwargs.update({"enable_lora": self.peft_config.enable_lora}) + # if isinstance(target, Conv1D): + # in_features, out_features = target.weight.shape + # else: + # 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 not a Conv1D. " + # "Setting fan_in_fan_out to False." + # ) + # kwargs["fan_in_fan_out"] = False + # new_module = MergedLinear(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"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 = None + 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 + if regu_loss is None: + regu_loss = torch.norm(para_cov-I, p="fro") + else: + regu_loss += torch.norm(para_cov-I, p="fro") + + outputs.loss += orth_reg_weight * regu_loss + return outputs + + + + +class SVDLinear(nn.Linear, LoRALayer): + # SVD-based adaptation for 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, + 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): + # def T(w): + # return w.T if self.fan_in_fan_out else w + 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): + # def T(w): + # return w.T if self.fan_in_fan_out else w + 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): + # def T(w): + # return w.T if self.fan_in_fan_out else w + 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): + # Lora implemented in a dense layer + 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.lora_A = nn.Linear(in_features, r, bias=False) + # self.lora_B = nn.Linear(r, out_features, bias=False) + # self.scaling = self.lora_alpha / self.r + # # Freezing the pre-trained weight matrix + # self.weight.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.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) + # nn.init.zeros_(self.lora_B.weight) + + 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 + # output = self.lora_B(self.lora_A(self.lora_dropout(x))).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 + # output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling + result += output + return result + + + +class RankAllocator(object): + def __init__(self, peft_config, param_iterator): + self.peft_config = peft_config + + self.ipt = {} + self.exp_avg_ipt = {} + self.exp_avg_unc = {} + self.cat_ipt = {} + + 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._set_budget_scheduler(param_iterator) + + + def set_total_step(self, total_step): + self.peft_config.total_step = total_step + + + def _set_budget_scheduler(self, param_iterator): + self.init_bgt = 0 + self.name_set = set() + for n,p in param_iterator: + 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) + + + 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 warmup + elif step > self.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): + 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() + self.exp_avg_ipt[n] = self.beta1 * self.exp_avg_ipt[n] + \ + (1 - self.beta1)*self.ipt[n] + 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 = {} + for n,p in model.named_parameters(): + if "lora_A" in n: + ipt_score = self._element_score(n) + comb_ipt = torch.mean(ipt_score, 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: + ipt_score = self._element_score(n) + comb_ipt = torch.mean(ipt_score, 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: + ipt_score = self._element_score(n) + name_m = n.replace("lora_E", "%s") + value_ipt[name_m] = ipt_score + + all_score = [] + 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)) + + mask_threshold = torch.kthvalue( + torch.cat(all_score), + k = self.init_bgt - budget, + )[0].item() + + with torch.no_grad(): + for n,p in model.named_parameters(): + if "lora_E" in n: + p.data.masked_fill_(triplet_ipt[n]<=mask_threshold, 0.0) + return mask_threshold + + def update_and_mask(self, model, global_step): + if global_step < self.peft_config.total_step - self.tfinal: + self.update_ipt(model) + budget, mask_ind = self.budget_schedule(global_step) + if mask_ind: + mask_threshold = self.mask_to_budget(model, budget) + else: + mask_threshold = None + + return budget, mask_threshold + + From 4acd81142905902cc3e3df5dac67707b96bc4d92 Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Wed, 1 Mar 2023 02:43:27 -0500 Subject: [PATCH 04/76] target module mapping for adalora --- src/peft/mapping.py | 38 ++++++++++++++++++++++++++++++++++++-- src/peft/tuners/adalora.py | 8 ++++++-- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 68de0c2..afc8bbe 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 LoraConfig, AdaLoraConfig 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, } TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { @@ -57,6 +58,25 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "layoutlm": ["query", "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"], +} + def get_peft_config(config_dict): """ @@ -123,6 +143,18 @@ 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: + 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 len(peft_config.target_modules) == 1: + peft_config.fan_in_fan_out = True + # peft_config.enable_lora = [True, False, True] + if peft_config.inference_mode: + peft_config.merge_weights = True + return peft_config + def get_peft_model(model, peft_config): """ @@ -138,7 +170,9 @@ def get_peft_model(model, peft_config): if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): peft_config = _prepare_lora_config(peft_config, model_config) return PeftModel(model, peft_config) - if not isinstance(peft_config, PromptLearningConfig): + if isinstance(peft_config, AdaLoraConfig): + peft_config = _prepare_adalora_config(peft_config, model_config) + elif not isinstance(peft_config, PromptLearningConfig): peft_config = _prepare_lora_config(peft_config, model_config) else: peft_config = _prepare_prompt_learning_config(peft_config, model_config) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index e8d5f86..f53df06 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -167,7 +167,7 @@ class AdaLoraModel(LoraModel): 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 @@ -189,6 +189,10 @@ class AdaLoraModel(LoraModel): outputs.loss += orth_reg_weight * regu_loss return outputs + def update_and_allocate(self, global_step): + self.rankallocator.update_and_allocate(self, global_step) + + @@ -483,7 +487,7 @@ class RankAllocator(object): p.data.masked_fill_(triplet_ipt[n]<=mask_threshold, 0.0) return mask_threshold - def update_and_mask(self, model, global_step): + def update_and_allocate(self, model, global_step): if global_step < self.peft_config.total_step - self.tfinal: self.update_ipt(model) budget, mask_ind = self.budget_schedule(global_step) From 6a03e43cbc97c4d871350f5890251195e798f8e5 Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Wed, 1 Mar 2023 02:47:33 -0500 Subject: [PATCH 05/76] peft import --- src/peft/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/peft/__init__.py b/src/peft/__init__.py index 3dd7acf..f25f032 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, From 510f172c58ff200d9c74eba636058c1dd7c8da56 Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Wed, 1 Mar 2023 21:26:07 +0000 Subject: [PATCH 06/76] adalora example --- .../peft_lora_seq2seq.ipynb | 2 +- src/peft/mapping.py | 2 +- src/peft/peft_model.py | 4 ++- src/peft/tuners/adalora.py | 26 +++++++++---------- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/examples/conditional_generation/peft_lora_seq2seq.ipynb b/examples/conditional_generation/peft_lora_seq2seq.ipynb index f22d3c6..bf864d4 100644 --- a/examples/conditional_generation/peft_lora_seq2seq.ipynb +++ b/examples/conditional_generation/peft_lora_seq2seq.ipynb @@ -473,7 +473,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.4" + "version": "3.9.16" }, "vscode": { "interpreter": { diff --git a/src/peft/mapping.py b/src/peft/mapping.py index afc8bbe..ceb220e 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -20,7 +20,7 @@ from .peft_model import ( PeftModelForSequenceClassification, PeftModelForTokenClassification, ) -from .tuners import LoraConfig, AdaLoraConfig PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig +from .tuners import LoraConfig, AdaLoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from .utils import PromptLearningConfig diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 4703059..ed92162 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -29,7 +29,7 @@ from transformers.utils import PushToHubMixin from huggingface_hub import hf_hub_download -from .tuners import LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder +from .tuners import LoraModel, AdaLoraConfig, AdaLoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder from .utils import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, WEIGHTS_NAME, @@ -76,6 +76,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.modules_to_save = None if isinstance(self.peft_config, PromptLearningConfig): self._setup_prompt_encoder() + elif isinstance(self.peft_config, AdaLoraConfig): + self.base_model = AdaLoraModel(peft_config, model) else: self.base_model = LoraModel(peft_config, model) if getattr(self.peft_config, "modules_to_save", None) is not None: diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index f53df06..9405d92 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -13,7 +13,7 @@ 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 .lora import LoraConfig, LoraModel, LoraLayer, mark_only_lora_as_trainable def is_bnb_available(): @@ -69,26 +69,26 @@ class AdaLoraModel(LoraModel): Args: model ([`transformers.PreTrainedModel`]): The model to be adapted. - config ([`LoraConfig`]): The configuration of the Lora model. + config ([`AdaLoraConfig`]): The configuration of the AdaLora model. Returns: - `torch.nn.Module`: The Lora model. + `torch.nn.Module`: The AdaLora model. Example:: - >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import LoraModel, LoraConfig >>> - config = LoraConfig( - peft_type="LORA", 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") >>> lora_model = LoraModel(config, model) + >>> 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") >>> adalora_model = AdaLoraModel(config, model) **Attributes**: - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. - - **peft_config** ([`LoraConfig`]): The configuration of the Lora model. + - **peft_config** ([`AdaLoraConfig`]): The configuration of the AdaLora model. """ def __init__(self, config, model): - # super().__init__() nn.Module.__init__(self) self.peft_config = config self.model = model @@ -194,9 +194,7 @@ class AdaLoraModel(LoraModel): - - -class SVDLinear(nn.Linear, LoRALayer): +class SVDLinear(nn.Linear, LoraLayer): # SVD-based adaptation for a dense layer def __init__( self, @@ -210,7 +208,7 @@ class SVDLinear(nn.Linear, LoRALayer): **kwargs ): nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, + 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 From 1a3680d8a74ae677a49d3d89daa203bc384f26af Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Wed, 1 Mar 2023 21:52:33 +0000 Subject: [PATCH 07/76] test for adalora example --- .../peft_adalora_seq2seq.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index 19875c4..5163f65 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -1,5 +1,5 @@ from transformers import AutoModelForSeq2SeqLM -from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, TaskType +from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, AdaLoraConfig, AdaLoraModel, TaskType import torch from datasets import load_dataset import os @@ -12,20 +12,24 @@ from tqdm import tqdm from datasets import load_dataset device = "cuda" -model_name_or_path = "bigscience/mt0-large" -tokenizer_name_or_path = "bigscience/mt0-large" +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 = 3 +num_epochs = 1 batch_size = 8 # creating model -peft_config = LoraConfig(task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1) +peft_config = AdaLoraConfig( + r=8, 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) @@ -89,6 +93,7 @@ lr_scheduler = get_linear_schedule_with_warmup( 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 From 35cd771c975df2469c37261aaec6344c265a9691 Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Wed, 1 Mar 2023 21:55:31 +0000 Subject: [PATCH 08/76] example --- examples/conditional_generation/peft_adalora_seq2seq.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index 5163f65..ec6b91a 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -26,7 +26,7 @@ batch_size = 8 # creating model peft_config = AdaLoraConfig( - r=8, lora_alpha=32, lora_dropout=0.1 + r=8, lora_alpha=32, lora_dropout=0.1, task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False ) @@ -151,7 +151,7 @@ model.save_pretrained(peft_model_id) ckpt = f"{peft_model_id}/adapter_model.bin" -get_ipython().system('du -h $ckpt') +# get_ipython().system('du -h $ckpt') from peft import PeftModel, PeftConfig From 7471035885a4ff19081a77a98da9735bb33db75c Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Thu, 2 Mar 2023 01:04:48 +0000 Subject: [PATCH 09/76] finish the testing and debugging --- .../peft_adalora_seq2seq.py | 25 ++++-- src/peft/peft_model.py | 2 + src/peft/tuners/adalora.py | 89 ++++++++----------- src/peft/utils/save_and_load.py | 4 +- 4 files changed, 58 insertions(+), 62 deletions(-) diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index ec6b91a..b3626f7 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -26,7 +26,10 @@ batch_size = 8 # creating model peft_config = AdaLoraConfig( - r=8, lora_alpha=32, lora_dropout=0.1, + init_r=12, target_r=1, + beta1=0.85, beta2=0.85, + tinit=0, tfinal=230, deltaT=1, + lora_alpha=32, lora_dropout=0.1, task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False ) @@ -98,19 +101,23 @@ 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() - optimizer.zero_grad() + with torch.autograd.set_detect_anomaly(True): + outputs = model(**batch) + loss = outputs.loss + total_loss += loss.detach().float() + loss.backward() + optimizer.step() + lr_scheduler.step() + + model.base_model.update_and_allocate(global_step) + optimizer.zero_grad() + global_step += 1 model.eval() eval_loss = 0 diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index ed92162..a3df6f6 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -180,6 +180,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module): hook = AlignDevicesHook(io_same_device=True) if model.peft_config.peft_type == PeftType.LORA: add_hook_to_module(model.base_model.model, hook) + elif model.peft_config.peft_type == PeftType.ADALORA: + add_hook_to_module(model.base_model.model, hook) else: remove_hook_from_submodules(model.prompt_encoder) add_hook_to_module(model.base_model, hook) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 9405d92..ad91f6a 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -30,17 +30,15 @@ class AdaLoraConfig(LoraConfig): This is the configuration class to store the configuration of a [`~peft.AdaLora`]. Args: - r (`int`): Lora attention dimension - 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) - enable_lora ( `List[bool]`): Used with `lora.MergedLinear`. - bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only' - modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable - and saved in the final checkpoint. + 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. """ target_r: int = field(default=8, metadata={"help": "Target Lora matrix dimension."}) init_r: int = field(default=12, metadata={"help": "Intial Lora matrix dimension."}) @@ -81,7 +79,8 @@ class AdaLoraModel(LoraModel): 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") >>> adalora_model = AdaLoraModel(config, model) + >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") + >>> model = AdaLoraModel(config, model) **Attributes**: - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. @@ -94,7 +93,6 @@ class AdaLoraModel(LoraModel): self.model = model self._find_and_replace() mark_only_lora_as_trainable(self.model, self.peft_config.bias) - # self.forward = self.model.forward self.rankallocator = RankAllocator(config, self.named_parameters()) def _find_and_replace(self): @@ -173,7 +171,7 @@ class AdaLoraModel(LoraModel): assert orth_reg_weight > 0 if hasattr(outputs, "loss"): - regu_loss = None + regu_loss = 0 num_param = 0 for n,p in self.model.named_parameters(): if "lora_A" in n or "lora_B" in n: @@ -181,12 +179,9 @@ class AdaLoraModel(LoraModel): I = torch.eye(*para_cov.size(), out=torch.empty_like(para_cov)) I.requires_grad = False num_param += 1 - if regu_loss is None: - regu_loss = torch.norm(para_cov-I, p="fro") - else: - regu_loss += torch.norm(para_cov-I, p="fro") - - outputs.loss += orth_reg_weight * regu_loss + 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 update_and_allocate(self, global_step): @@ -195,7 +190,7 @@ class AdaLoraModel(LoraModel): class SVDLinear(nn.Linear, LoraLayer): - # SVD-based adaptation for a dense layer + # SVD-based adaptation by a dense layer def __init__( self, in_features: int, @@ -239,8 +234,6 @@ class SVDLinear(nn.Linear, LoraLayer): nn.init.normal_(self.lora_B, mean=0.0, std=0.02) def train(self, mode: bool = True): - # def T(w): - # return w.T if self.fan_in_fan_out else w nn.Linear.train(self, mode) if self.merge_weights and self.merged: # Make sure that the weights are not merged @@ -251,8 +244,6 @@ class SVDLinear(nn.Linear, LoraLayer): self.merged = False def eval(self): - # def T(w): - # return w.T if self.fan_in_fan_out else w nn.Linear.eval(self) if self.merge_weights and not self.merged: # Merge the weights and mark it @@ -263,8 +254,6 @@ class SVDLinear(nn.Linear, LoraLayer): self.merged = True def forward(self, x: torch.Tensor): - # def T(w): - # return w.T if self.fan_in_fan_out else w 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: @@ -277,9 +266,8 @@ class SVDLinear(nn.Linear, LoraLayer): if is_bnb_available(): - class SVDLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): - # Lora implemented in a dense layer + # Low-rank matrix for SVD-based adaptation def __init__( self, in_features, @@ -315,20 +303,11 @@ if is_bnb_available(): # Freezing the pre-trained weight matrix self.weight.requires_grad = False self.ranknum.requires_grad = False - - # self.lora_A = nn.Linear(in_features, r, bias=False) - # self.lora_B = nn.Linear(r, out_features, bias=False) - # self.scaling = self.lora_alpha / self.r - # # Freezing the pre-trained weight matrix - # self.weight.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.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - # nn.init.zeros_(self.lora_B.weight) - 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) @@ -347,13 +326,11 @@ if is_bnb_available(): 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 - # output = self.lora_B(self.lora_A(self.lora_dropout(x))).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 - # output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling result += output return result @@ -400,8 +377,8 @@ class RankAllocator(object): if step <= tinit: budget = self.init_bgt mask_ind = False - # Final warmup - elif step > self.total_step - tfinal: + # Final fine-tuning + elif step > total_step - tfinal: budget = self.target_bgt mask_ind = True else: @@ -415,6 +392,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 n not in self.ipt: @@ -423,8 +401,10 @@ class RankAllocator(object): 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() @@ -443,29 +423,31 @@ class RankAllocator(object): 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: - ipt_score = self._element_score(n) - comb_ipt = torch.mean(ipt_score, dim=1, keepdim=True) + 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: - ipt_score = self._element_score(n) - comb_ipt = torch.mean(ipt_score, dim=0, keepdim=False).view(-1, 1) + 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: - ipt_score = self._element_score(n) + entry_ipt = self._element_score(n) name_m = n.replace("lora_E", "%s") - value_ipt[name_m] = ipt_score + 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) @@ -474,11 +456,13 @@ class RankAllocator(object): 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() + # Mask the unimportant triplets with torch.no_grad(): for n,p in model.named_parameters(): if "lora_E" in n: @@ -486,14 +470,17 @@ class RankAllocator(object): return mask_threshold def update_and_allocate(self, model, global_step): - if global_step < self.peft_config.total_step - self.tfinal: + # Update the importance score and allocate the budget + if global_step < self.peft_config.total_step - self.peft_config.tfinal: self.update_ipt(model) + # TODO: Finalize the budget distribution by replacing with new Linear. budget, mask_ind = self.budget_schedule(global_step) + print("budget:", budget) if mask_ind: - mask_threshold = self.mask_to_budget(model, budget) + mask_threshold = self.mask_to_budget(model, budget) + print("mask threshold:", mask_threshold) else: mask_threshold = None - return budget, mask_threshold diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index c6596c7..86e388b 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -29,7 +29,7 @@ def get_peft_model_state_dict(model, state_dict=None): """ if state_dict is None: state_dict = model.state_dict() - if model.peft_config.peft_type == PeftType.LORA: + if model.peft_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 directly with the state dict which is necessary when using DeepSpeed or FSDP @@ -72,7 +72,7 @@ def set_peft_model_state_dict(model, peft_model_state_dict): """ model.load_state_dict(peft_model_state_dict, strict=False) - if model.peft_config.peft_type != PeftType.LORA: + 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 ) From 0a0c6ea6eac9a0e6aa420abb55a7a0ed418cf442 Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Thu, 2 Mar 2023 01:08:41 +0000 Subject: [PATCH 10/76] adalora training example --- .../peft_adalora_seq2seq.py | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index b3626f7..49fa497 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -1,5 +1,5 @@ from transformers import AutoModelForSeq2SeqLM -from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, AdaLoraConfig, AdaLoraModel, TaskType +from peft import get_peft_model, AdaLoraConfig, AdaLoraModel, TaskType import torch from datasets import load_dataset import os @@ -20,15 +20,15 @@ text_column = "sentence" label_column = "text_label" max_length = 128 lr = 1e-3 -num_epochs = 1 +num_epochs = 8 batch_size = 8 # creating model peft_config = AdaLoraConfig( - init_r=12, target_r=1, + init_r=12, target_r=8, beta1=0.85, beta2=0.85, - tinit=0, tfinal=230, deltaT=1, + tinit=200, tfinal=1000, deltaT=10, lora_alpha=32, lora_dropout=0.1, task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False @@ -107,17 +107,17 @@ for epoch in range(num_epochs): total_loss = 0 for step, batch in enumerate(tqdm(train_dataloader)): batch = {k: v.to(device) for k, v in batch.items()} - with torch.autograd.set_detect_anomaly(True): - outputs = model(**batch) - loss = outputs.loss - total_loss += loss.detach().float() - loss.backward() - optimizer.step() - lr_scheduler.step() - - model.base_model.update_and_allocate(global_step) - optimizer.zero_grad() - global_step += 1 + 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 From fa65b95b9e9613f3adb705af1151ea5823bc76a0 Mon Sep 17 00:00:00 2001 From: QingruZhang Date: Thu, 2 Mar 2023 01:11:14 +0000 Subject: [PATCH 11/76] update comment --- src/peft/tuners/adalora.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index ad91f6a..61bd512 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -63,7 +63,8 @@ class AdaLoraConfig(LoraConfig): class AdaLoraModel(LoraModel): """ - Creates Adaptive LoRA (AdaLora) model from a pretrained transformers model. + 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. From c21afbe868734c0af8bd4577c4c7acdf366b96d1 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 18:56:24 +0530 Subject: [PATCH 12/76] multi adapter for training and inference Might have breaking changes --- src/peft/mapping.py | 43 +-- src/peft/peft_model.py | 284 ++++++++++------ src/peft/tuners/__init__.py | 2 +- src/peft/tuners/lora.py | 563 ++++++++++++------------------- src/peft/utils/__init__.py | 6 +- src/peft/utils/adapters_utils.py | 18 - src/peft/utils/config.py | 16 +- src/peft/utils/other.py | 83 ++++- src/peft/utils/save_and_load.py | 57 +++- 9 files changed, 524 insertions(+), 548 deletions(-) delete mode 100644 src/peft/utils/adapters_utils.py diff --git a/src/peft/mapping.py b/src/peft/mapping.py index dbb9f36..c814655 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -38,27 +38,6 @@ PEFT_TYPE_TO_CONFIG_MAPPING = { "LORA": LoraConfig, } -TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { - "t5": ["q", "v"], - "mt5": ["q", "v"], - "bart": ["q_proj", "v_proj"], - "gpt2": ["c_attn"], - "bloom": ["query_key_value"], - "opt": ["q_proj", "v_proj"], - "gptj": ["q_proj", "v_proj"], - "gpt_neox": ["query_key_value"], - "gpt_neo": ["q_proj", "v_proj"], - "bert": ["query", "value"], - "roberta": ["query", "value"], - "xlm-roberta": ["query", "value"], - "electra": ["query", "value"], - "deberta-v2": ["query_proj", "value_proj"], - "deberta": ["in_proj"], - "layoutlm": ["query", "value"], - "llama": ["q_proj", "v_proj"], - "chatglm": ["query_key_value"], -} - def get_peft_config(config_dict): """ @@ -113,19 +92,6 @@ def _prepare_prompt_learning_config(peft_config, model_config): return peft_config -def _prepare_lora_config(peft_config, model_config): - if peft_config.target_modules is None: - if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: - raise ValueError("Please specify `target_modules` in `peft_config`") - peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] - if len(peft_config.target_modules) == 1: - peft_config.fan_in_fan_out = True - peft_config.enable_lora = [True, False, True] - if peft_config.inference_mode: - peft_config.merge_weights = True - return peft_config - - def get_peft_model(model, peft_config): """ Returns a Peft model object from a model and a config. @@ -137,11 +103,10 @@ def get_peft_model(model, peft_config): model_config = model.config.to_dict() peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None) - if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): - peft_config = _prepare_lora_config(peft_config, model_config) + if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not isinstance( + peft_config, PromptLearningConfig + ): return PeftModel(model, peft_config) - if not isinstance(peft_config, PromptLearningConfig): - peft_config = _prepare_lora_config(peft_config, model_config) - else: + if isinstance(peft_config, PromptLearningConfig): peft_config = _prepare_prompt_learning_config(peft_config, model_config) return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index f73a66a..cbc8a89 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -36,6 +36,7 @@ from .utils import ( PeftType, PromptLearningConfig, TaskType, + _set_adapter, _set_trainable, get_peft_model_state_dict, set_peft_model_state_dict, @@ -43,6 +44,14 @@ from .utils import ( ) +PEFT_TYPE_TO_MODEL_MAPPING = { + PeftType.LORA: LoraModel, + PeftType.PROMPT_TUNING: PromptEmbedding, + PeftType.P_TUNING: PromptEncoder, + PeftType.PREFIX_TUNING: PrefixEncoder, +} + + class PeftModel(PushToHubMixin, torch.nn.Module): """ Parameter-Efficient Fine-Tuning Model. Base model encompassing various Peft methods. @@ -67,20 +76,19 @@ class PeftModel(PushToHubMixin, torch.nn.Module): in the base model if `isinstance(self.peft_config, PromptLearningConfig)`. """ - def __init__(self, model, peft_config: PeftConfig): + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): super().__init__() - self.peft_config = peft_config self.base_model = model self.config = self.base_model.config self.modules_to_save = None - if isinstance(self.peft_config, PromptLearningConfig): - self._setup_prompt_encoder() - else: - self.base_model = LoraModel(peft_config, model) - if getattr(self.peft_config, "modules_to_save", None) is not None: - self.modules_to_save = self.peft_config.modules_to_save - _set_trainable(self) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.peft_config = {} + self.active_adapter = adapter_name + if not isinstance(peft_config, PromptLearningConfig): + self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( + self.base_model, peft_config, adapter_name + ) + self.add_adapter(adapter_name, peft_config) def save_pretrained(self, save_directory, **kwargs): r""" @@ -98,27 +106,30 @@ class PeftModel(PushToHubMixin, torch.nn.Module): raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file") os.makedirs(save_directory, exist_ok=True) - # save only the trainable weights - output_state_dict = get_peft_model_state_dict(self, kwargs.get("state_dict", None)) - torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME)) + for adapter_name, peft_config in self.peft_config.items(): + # save only the trainable weights + output_state_dict = get_peft_model_state_dict(self, adapter_name, kwargs.get("state_dict", None)) + output_dir = os.path.join(save_directory, adapter_name) if adapter_name != "default" else save_directory + os.makedirs(output_dir, exist_ok=True) + torch.save(output_state_dict, os.path.join(output_dir, WEIGHTS_NAME)) - # save the config and change the inference mode to `True` - if self.peft_config.base_model_name_or_path is None: - self.peft_config.base_model_name_or_path = ( - self.base_model.__dict__.get("name_or_path", None) - if isinstance(self.peft_config, PromptLearningConfig) - else self.base_model.model.__dict__.get("name_or_path", None) - ) - inference_mode = self.peft_config.inference_mode - self.peft_config.inference_mode = True - self.peft_config.save_pretrained(save_directory) - self.peft_config.inference_mode = inference_mode + # save the config and change the inference mode to `True` + if peft_config.base_model_name_or_path is None: + peft_config.base_model_name_or_path = ( + self.base_model.__dict__.get("name_or_path", None) + if isinstance(self.peft_config, PromptLearningConfig) + else self.base_model.model.__dict__.get("name_or_path", None) + ) + inference_mode = self.peft_config.inference_mode + peft_config.inference_mode = True + peft_config.save_pretrained(output_dir) + peft_config.inference_mode = inference_mode @classmethod - def from_pretrained(cls, model, model_id, **kwargs): + def from_pretrained(cls, model, model_id, adapter_name="default", **kwargs): r""" Args: - Instantiate a `LoraModel` from a pretrained Lora configuration and weights. + Instantiate a `PeftModel` from a pretrained Peft configuration and weights. model (`transformers.PreTrainedModel`): The model to be adapted. The model should be initialized with the `from_pretrained` method. from `transformers` library. @@ -132,58 +143,26 @@ class PeftModel(PushToHubMixin, torch.nn.Module): from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PEFT_TYPE_TO_CONFIG_MAPPING # load the config - config = PEFT_TYPE_TO_CONFIG_MAPPING[PeftConfig.from_pretrained(model_id).peft_type].from_pretrained(model_id) + 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 getattr(model, "hf_device_map", None) is not None: + if (getattr(model, "hf_device_map", None) is not None) and len( + set(model.hf_device_map.values()).intersection({"cpu", "disk"}) + ) > 0: remove_hook_from_submodules(model) if config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): - model = cls(model, config) + model = cls(model, config, adapter_name) else: - model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config) - - # load weights if any - if os.path.exists(os.path.join(model_id, WEIGHTS_NAME)): - filename = os.path.join(model_id, WEIGHTS_NAME) - else: - try: - filename = hf_hub_download(model_id, WEIGHTS_NAME) - except: # noqa - raise ValueError( - f"Can't find weights for {model_id} in {model_id} or in the Hugging Face Hub. " - f"Please check that the file {WEIGHTS_NAME} is present at {model_id}." - ) - - adapters_weights = torch.load( - filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu") - ) - # load the weights into the model - model = set_peft_model_state_dict(model, adapters_weights) - if getattr(model, "hf_device_map", None) is not None: - device_map = kwargs.get("device_map", "auto") - max_memory = kwargs.get("max_memory", None) - no_split_module_classes = model._no_split_modules - if device_map != "sequential": - max_memory = get_balanced_memory( - model, - max_memory=max_memory, - no_split_module_classes=no_split_module_classes, - low_zero=(device_map == "balanced_low_0"), - ) - if isinstance(device_map, str): - device_map = infer_auto_device_map( - model, max_memory=max_memory, no_split_module_classes=no_split_module_classes - ) - model = dispatch_model(model, device_map=device_map) - hook = AlignDevicesHook(io_same_device=True) - if model.peft_config.peft_type == PeftType.LORA: - add_hook_to_module(model.base_model.model, hook) - else: - remove_hook_from_submodules(model.prompt_encoder) - add_hook_to_module(model.base_model, hook) + model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config, adapter_name) + model.load_adapter(model_id, adapter_name, **kwargs) return model - def _setup_prompt_encoder(self): + def _setup_prompt_encoder(self, adapter_name): + config = self.peft_config[adapter_name] + self.prompt_encoder = torch.nn.ModuleDict({}) + self.prompt_tokens = {} transformer_backbone = None for name, module in self.base_model.named_children(): for param in module.parameters(): @@ -194,51 +173,50 @@ class PeftModel(PushToHubMixin, torch.nn.Module): transformer_backbone = module self.transformer_backbone_name = name - if self.peft_config.num_transformer_submodules is None: - self.peft_config.num_transformer_submodules = ( - 2 if self.peft_config.task_type == TaskType.SEQ_2_SEQ_LM else 1 - ) + if config.num_transformer_submodules is None: + config.num_transformer_submodules = 2 if config.task_type == TaskType.SEQ_2_SEQ_LM else 1 for named_param, value in list(transformer_backbone.named_parameters()): if value.shape[0] == self.base_model.config.vocab_size: self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", "")) break - if self.peft_config.peft_type == PeftType.PROMPT_TUNING: - prompt_encoder = PromptEmbedding(self.peft_config, self.word_embeddings) - elif self.peft_config.peft_type == PeftType.P_TUNING: - prompt_encoder = PromptEncoder(self.peft_config) - elif self.peft_config.peft_type == PeftType.PREFIX_TUNING: - prompt_encoder = PrefixEncoder(self.peft_config) + if config.peft_type == PeftType.PROMPT_TUNING: + prompt_encoder = PromptEmbedding(config, self.word_embeddings) + elif config.peft_type == PeftType.P_TUNING: + prompt_encoder = PromptEncoder(config) + elif config.peft_type == PeftType.PREFIX_TUNING: + prompt_encoder = PrefixEncoder(config) else: raise ValueError("Not supported") - self.prompt_encoder = prompt_encoder - self.prompt_tokens = torch.arange( - self.peft_config.num_virtual_tokens * self.peft_config.num_transformer_submodules + self.prompt_encoder.update(torch.nn.ModuleDict({adapter_name: prompt_encoder})) + self.prompt_tokens[adapter_name] = torch.arange( + config.num_virtual_tokens * config.num_transformer_submodules ).long() - def get_prompt_embedding_to_save(self): + def get_prompt_embedding_to_save(self, adapter_name): """ Returns the prompt embedding to save when saving the model. Only applicable when `peft_config.peft_type != PeftType.LORA`. """ - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.device) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: - prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens] - prompt_embeddings = self.prompt_encoder(prompt_tokens) + prompt_tokens = self.prompt_tokens[adapter_name].unsqueeze(0).expand(1, -1).to(self.device) + if self.peft_config[adapter_name].peft_type == PeftType.PREFIX_TUNING: + prompt_tokens = prompt_tokens[:, : self.peft_config[adapter_name].num_virtual_tokens] + prompt_embeddings = self.prompt_encoder[adapter_name](prompt_tokens) return prompt_embeddings[0].detach().cpu() def get_prompt(self, batch_size): """ Returns the virtual prompts to use for Peft. Only applicable when `peft_config.peft_type != PeftType.LORA`. """ - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.device) + prompt_encoder = self.prompt_encoder[self.active_adapter] + prompt_tokens = self.prompt_tokens[self.active_adapter].unsqueeze(0).expand(batch_size, -1).to(self.device) if self.peft_config.peft_type == PeftType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens] if self.peft_config.inference_mode: - past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + past_key_values = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: - past_key_values = self.prompt_encoder(prompt_tokens) + past_key_values = prompt_encoder(prompt_tokens) past_key_values = past_key_values.view( batch_size, self.peft_config.num_virtual_tokens, @@ -257,9 +235,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): return past_key_values else: if self.peft_config.inference_mode: - prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + prompts = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: - prompts = self.prompt_encoder(prompt_tokens) + prompts = prompt_encoder(prompt_tokens) return prompts def print_trainable_parameters(self): @@ -299,13 +277,13 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Disables the adapter module. """ - if isinstance(self.peft_config, PromptLearningConfig): + if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): old_forward = self.forward self.forward = self.base_model.forward else: self.base_model.disable_adapter_layers() yield - if isinstance(self.peft_config, PromptLearningConfig): + if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): self.forward = old_forward else: self.base_model.enable_adapter_layers() @@ -314,7 +292,91 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Returns the base model. """ - return self.base_model if isinstance(self.peft_config, PromptLearningConfig) else self.base_model.model + return ( + self.base_model + if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig) + else self.base_model.model + ) + + def add_adapter(self, adapter_name, peft_config): + self.peft_config[adapter_name] = peft_config + if isinstance(peft_config, PromptLearningConfig): + self._setup_prompt_encoder(adapter_name) + else: + self.base_model.add_adapter(adapter_name, peft_config) + if getattr(peft_config, "modules_to_save", None) is not None: + if self.modules_to_save is None: + self.modules_to_save = set(peft_config.modules_to_save) + else: + 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): + from .mapping import PEFT_TYPE_TO_CONFIG_MAPPING + + if adapter_name not in self.peft_config: + # load the config + 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)) + self.add_adapter(adapter_name, peft_config) + + # load weights if any + if kwargs.get("subfolder", None) is not None: + path = os.path.join(model_id, kwargs["subfolder"]) + if os.path.exists(os.path.join(path, WEIGHTS_NAME)): + filename = os.path.join(path, WEIGHTS_NAME) + else: + try: + filename = hf_hub_download(model_id, WEIGHTS_NAME, subfolder=kwargs.get("subfolder", None)) + except: # noqa + raise ValueError( + f"Can't find weights for {model_id} in {model_id} or in the Hugging Face Hub. " + f"Please check that the file {WEIGHTS_NAME} is present at {model_id}." + ) + + adapters_weights = torch.load( + filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu") + ) + # load the weights into the model + set_peft_model_state_dict(self, adapter_name, adapters_weights) + if ( + (getattr(self, "hf_device_map", None) is not None) + and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0) + and len(self.peft_config == 1) + ): + device_map = kwargs.get("device_map", "auto") + max_memory = kwargs.get("max_memory", None) + no_split_module_classes = self._no_split_modules + if device_map != "sequential": + max_memory = get_balanced_memory( + self, + max_memory=max_memory, + no_split_module_classes=no_split_module_classes, + low_zero=(device_map == "balanced_low_0"), + ) + if isinstance(device_map, str): + device_map = infer_auto_device_map( + self, max_memory=max_memory, no_split_module_classes=no_split_module_classes + ) + dispatch_model(self, device_map=device_map) + hook = AlignDevicesHook(io_same_device=True) + if not isinstance(self.peft_config[adapter_name]) == PeftType.LORA: + add_hook_to_module(self.base_model.model, hook) + else: + remove_hook_from_submodules(self.prompt_encoder) + add_hook_to_module(self.base_model, hook) + + def set_adapter(self, adapter_name): + """ + Sets the active adapter. + """ + if adapter_name not in self.peft_config: + raise ValueError(f"Adapter {adapter_name} not found.") + self.active_adapter = adapter_name + if not isinstance(self.peft_config[adapter_name], PromptLearningConfig): + self.base_model.set_adapter(adapter_name) + _set_adapter(self, adapter_name) class PeftModelForSequenceClassification(PeftModel): @@ -343,9 +405,12 @@ class PeftModelForSequenceClassification(PeftModel): params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) - self.modules_to_save = ["classifier", "score"] + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) + if self.modules_to_save is None: + self.modules_to_save = {"classifier", "score"} + else: + self.modules_to_save.update({"classifier", "score"}) for name, _ in self.base_model.named_children(): if any(module_name in name for module_name in self.modules_to_save): @@ -353,7 +418,7 @@ class PeftModelForSequenceClassification(PeftModel): break # to make sure classifier layer is trainable - _set_trainable(self) + _set_trainable(self, adapter_name) def forward( self, @@ -510,8 +575,8 @@ class PeftModelForCausalLM(PeftModel): params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation def forward( @@ -647,8 +712,8 @@ class PeftModelForSeq2SeqLM(PeftModel): params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation self.base_model_prepare_encoder_decoder_kwargs_for_generation = ( self.base_model._prepare_encoder_decoder_kwargs_for_generation @@ -818,9 +883,12 @@ class PeftModelForTokenClassification(PeftModel): params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) - self.modules_to_save = ["classifier", "score"] + def __init__(self, model, peft_config: PeftConfig = None, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) + if self.modules_to_save is None: + self.modules_to_save = {"classifier", "score"} + else: + self.modules_to_save.update({"classifier", "score"}) for name, _ in self.base_model.named_children(): if any(module_name in name for module_name in self.modules_to_save): @@ -828,7 +896,7 @@ class PeftModelForTokenClassification(PeftModel): break # to make sure classifier layer is trainable - _set_trainable(self) + _set_trainable(self, adapter_name) def forward( self, diff --git a/src/peft/tuners/__init__.py b/src/peft/tuners/__init__.py index 38b7926..8f93079 100644 --- a/src/peft/tuners/__init__.py +++ b/src/peft/tuners/__init__.py @@ -17,7 +17,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .lora import LoraConfig, LoraModel 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/lora.py b/src/peft/tuners/lora.py index 0f65cbf..34cee7d 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -25,7 +25,13 @@ import torch.nn as nn import torch.nn.functional as F from transformers.pytorch_utils import Conv1D -from ..utils import PeftConfig, PeftType, transpose +from ..utils import ( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + PeftConfig, + PeftType, + _get_submodules, + transpose, +) def is_bnb_available(): @@ -48,8 +54,9 @@ class LoraConfig(PeftConfig): 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) - enable_lora ( `List[bool]`): Used with `lora.MergedLinear`. + 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`.: + enable_lora ( `List[bool]`): Used with `lora.MergedLinear`. Usually set to [True, False, True]. bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only' modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable and saved in the final checkpoint. @@ -72,7 +79,6 @@ class LoraConfig(PeftConfig): default=False, metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"}, ) - enable_lora: Optional[List[bool]] = field(default=None, metadata={"help": "Used with `lora.MergedLinear`."}) bias: str = field(default="none", metadata={"help": "Bias type for Lora. Can be 'none', 'all' or 'lora_only'"}) modules_to_save: Optional[List[str]] = field( default=None, @@ -88,38 +94,26 @@ class LoraConfig(PeftConfig): class LoraModel(torch.nn.Module): - """ - Creates Low Rank Adapter (Lora) model from a pretrained transformers model. - - Args: - model ([`transformers.PreTrainedModel`]): The model to be adapted. - config ([`LoraConfig`]): The configuration of the Lora model. - - Returns: - `torch.nn.Module`: The Lora model. - - Example:: - - >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import LoraModel, LoraConfig >>> - config = LoraConfig( - peft_type="LORA", 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") >>> lora_model = LoraModel(config, model) - - **Attributes**: - - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. - - **peft_config** ([`LoraConfig`]): The configuration of the Lora model. - """ - - def __init__(self, config, model): + def __init__(self, model, config, adapter_name): super().__init__() - self.peft_config = config self.model = model - self._find_and_replace() - mark_only_lora_as_trainable(self.model, self.peft_config.bias) self.forward = self.model.forward + self.config = config + self.add_adapter(adapter_name) - def _find_and_replace(self): + def add_adapter(self, adapter_name, config=None): + if config is not None: + config = self._prepare_lora_config(config, self.model.config.to_dict()) + self.config[adapter_name] = config + self._find_and_replace(adapter_name) + if len(self.config) > 1 and self.config[adapter_name].bias != "none": + 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.config[adapter_name].bias) + + def _find_and_replace(self, adapter_name): + lora_config = self.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( @@ -129,68 +123,72 @@ class LoraModel(torch.nn.Module): is_target_modules_in_base_model = False is_hf_device_map_available = hasattr(self.model, "hf_device_map") kwargs = { - "r": self.peft_config.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, + "merge_weights": (lora_config.merge_weights or lora_config.inference_mode) and not is_hf_device_map_available, } 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(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 self.peft_config.enable_lora is None: - new_module = Linear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) - else: - kwargs.update({"enable_lora": self.peft_config.enable_lora}) - new_module = MergedLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) - elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None: - new_module = Linear(target.in_features, target.out_features, bias=bias, **kwargs) - elif self.peft_config.enable_lora is not None: - kwargs.update({"enable_lora": self.peft_config.enable_lora}) - if isinstance(target, Conv1D): - in_features, out_features = ( - target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape + if isinstance(target, LoraLayer): + target.update_layer(adapter_name, lora_config.r, lora_config.lora_alpha, lora_config.lora_dropout) + 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 = Linear8bitLt( + adapter_name, target.in_features, target.out_features, bias=bias, **kwargs ) else: - 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 not a Conv1D. " - "Setting fan_in_fan_out to False." + 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 ) - kwargs["fan_in_fan_out"] = self.peft_config.fan_in_fan_out = False - new_module = MergedLinear(in_features, out_features, bias=bias, **kwargs) - self._replace_module(parent, target_name, new_module, target) + 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 = Linear(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." ) - def _get_submodules(self, key): - parent = self.model.get_submodule(".".join(key.split(".")[:-1])) - target_name = key.split(".")[-1] - target = self.model.get_submodule(key) - return parent, target, target_name - def _replace_module(self, parent_module, child_name, new_module, old_module): setattr(parent_module, child_name, new_module) new_module.weight = old_module.weight @@ -217,9 +215,12 @@ class LoraModel(torch.nn.Module): return None def get_peft_config_as_dict(self, inference: bool = False): - config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(self.peft_config).items()} - if inference: - config["inference_mode"] = True + config_dict = {} + for key, value in self.config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config return config def _set_adapter_layers(self, enabled=True): @@ -233,6 +234,34 @@ class LoraModel(torch.nn.Module): def disable_adapter_layers(self): self._set_adapter_layers(enabled=False) + def set_adapter(self, adapter_name): + for module in self.model.modules(): + if isinstance(module, LoraLayer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.active_adapter = adapter_name + + def merge_adapter(self): + for module in self.model.modules(): + if isinstance(module, LoraLayer): + module.merge() + + def unmerge_adapter(self): + for module in self.model.modules(): + if isinstance(module, LoraLayer): + module.unmerge() + + @staticmethod + def _prepare_lora_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] + if peft_config.inference_mode: + peft_config.merge_weights = True + return peft_config + # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py # and modified to work with PyTorch FSDP @@ -266,28 +295,53 @@ def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None: class LoraLayer: def __init__( self, - r: int, - lora_alpha: int, - lora_dropout: float, merge_weights: bool, + in_features: int, + out_features: int, ): - self.r = r - self.lora_alpha = lora_alpha - # Optional dropout - if lora_dropout > 0.0: - self.lora_dropout = nn.Dropout(p=lora_dropout) - else: - self.lora_dropout = lambda x: x + self.r = {} + self.lora_alpha = {} + self.scaling = {} + self.lora_dropout = nn.ModuleDict({}) + self.lora_A = nn.ModuleDict({}) + 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 + + def update_layer(self, adapter_name, r, lora_alpha, lora_dropout): + 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: + self.lora_A.update(nn.ModuleDict({nn.Linear(self.in_features, r, bias=False)})) + self.lora_B.update(nn.ModuleDict({nn.Linear(r, self.out_features, bias=False)})) + self.scaling[adapter_name] = lora_alpha / r + self.reset_lora_parameters(adapter_name) + + def reset_lora_parameters(self, adapter_name): + if adapter_name in self.lora_A.keys(): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A[adapter_name].weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B[adapter_name].weight) -class Linear(nn.Linear, LoraLayer): +class Linear(nn.Linear): # Lora implemented in a dense layer def __init__( self, + adapter_name: str, in_features: int, out_features: int, r: int = 0, @@ -298,185 +352,67 @@ class Linear(nn.Linear, LoraLayer): **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, merge_weights=merge_weights, 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: - self.lora_A = nn.Linear(in_features, r, bias=False) - self.lora_B = nn.Linear(r, out_features, bias=False) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.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"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) + self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.active_adapter = adapter_name - def train(self, mode: bool = True): - nn.Linear.train(self, mode) - self.lora_A.train(mode) - self.lora_B.train(mode) - if not mode and self.merge_weights and not self.merged: - # Merge the weights and mark it - if self.r > 0: - self.weight.data += ( - transpose(self.lora_B.weight @ self.lora_A.weight, self.fan_in_fan_out) * self.scaling + def merge(self): + 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 + if self.r[self.active_adapter] > 0: + self.weight.data += ( + transpose( + self.lora_B[self.active_adapter].weight @ self.lora_A[self.active_adapter].weight, + self.fan_in_fan_out, ) - self.merged = True - elif 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.weight @ self.lora_A.weight, self.fan_in_fan_out) * self.scaling - ) - self.merged = False - - def eval(self): - nn.Linear.eval(self) - self.lora_A.eval() - self.lora_B.eval() - - def forward(self, x: torch.Tensor): - if self.disable_adapters: - if self.r > 0 and self.merged: - self.weight.data -= ( - transpose(self.lora_B.weight @ self.lora_A.weight, self.fan_in_fan_out) * self.scaling - ) - self.merged = False - - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - elif 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_B(self.lora_A(self.lora_dropout(x))) * self.scaling - return result - else: - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - - -class MergedLinear(nn.Linear, LoraLayer): - # Lora implemented in a dense layer - def __init__( - self, - in_features: int, - out_features: int, - r: int = 0, - lora_alpha: int = 1, - lora_dropout: float = 0.0, - enable_lora: List[bool] = [False], - 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) - if out_features % len(enable_lora) != 0: - raise ValueError("The length of enable_lora must divide out_features") - self.enable_lora = enable_lora - self.fan_in_fan_out = fan_in_fan_out - # Actual trainable parameters - if r > 0 and any(enable_lora): - self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False) - self.lora_B = nn.Conv1d( - r * sum(enable_lora), - out_features // len(enable_lora) * sum(enable_lora), - kernel_size=1, - groups=2, - bias=False, + * self.scaling[self.active_adapter] ) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - # Compute the indices - self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1) - self.lora_ind[enable_lora, :] = True - self.lora_ind = self.lora_ind.view(-1) - 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"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) - - def zero_pad(self, x): - result = x.new_zeros((*x.shape[:-1], self.out_features)) - result = result.view(-1, self.out_features) - result[:, self.lora_ind] = x.reshape(-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora)) - return result.view((*x.shape[:-1], self.out_features)) - - def train(self, mode: bool = True): - nn.Linear.train(self, mode) - self.lora_A.train(mode) - self.lora_B.train(mode) - if not mode and self.merge_weights and not self.merged: - # Merge the weights and mark it - if self.r > 0 and any(self.enable_lora): - delta_w = ( - F.conv1d( - self.lora_A.weight.data.unsqueeze(0), - self.lora_B.weight.data, - groups=sum(self.enable_lora), - ) - .squeeze(0) - .transpose(-2, -1) - ) - self.weight.data += transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) self.merged = True - elif self.merge_weights and self.merged: - # Make sure that the weights are not merged - if self.r > 0 and any(self.enable_lora): - delta_w = ( - F.conv1d( - self.lora_A.weight.data.unsqueeze(0), - self.lora_B.weight.data, - groups=sum(self.enable_lora), - ) - .squeeze(0) - .transpose(-2, -1) - ) - self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) - self.merged = False - def eval(self): - nn.Linear.eval(self) - self.lora_A.eval() - self.lora_B.eval() + def unmerge(self): + 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 + if self.r[self.active_adapter] > 0: + self.weight.data -= ( + transpose( + self.lora_B[self.active_adapter].weight @ self.lora_A[self.active_adapter].weight, + self.fan_in_fan_out, + ) + * self.scaling[self.active_adapter] + ) + self.merged = False def forward(self, x: torch.Tensor): if self.disable_adapters: - if self.r > 0 and self.merged and any(self.enable_lora): - delta_w = ( - F.conv1d( - self.lora_A.weight.data.unsqueeze(0), - self.lora_B.weight.data, - groups=sum(self.enable_lora), - ) - .squeeze(0) - .transpose(-2, -1) - ) - self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) - self.merged = False + if self.r[self.active_adapter] > 0 and self.merged: + self.unmerge() return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - elif self.merged: - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - else: + 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) - if self.r > 0: - after_A = self.lora_A(self.lora_dropout(x)) - after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) - result += self.zero_pad(after_B) * self.scaling - return result + result += ( + self.lora_B[self.active_adapter]( + self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x)) + ) + * self.scaling[self.active_adapter] + ) + else: + return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if is_bnb_available(): @@ -485,6 +421,7 @@ if is_bnb_available(): # Lora implemented in a dense layer def __init__( self, + adapter_name, in_features, out_features, r: int = 0, @@ -502,115 +439,37 @@ 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: - self.lora_A = nn.Linear(in_features, r, bias=False) - self.lora_B = nn.Linear(r, out_features, bias=False) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - self.reset_parameters() + LoraLayer.__init__(self, merge_weights=False, in_features=in_features, out_features=out_features) - 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.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + + self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.active_adapter = adapter_name def forward(self, x: torch.Tensor): result = super().forward(x) if self.disable_adapters: 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_B(self.lora_A(self.lora_dropout(x))).to(expected_dtype) * self.scaling - result += output + output = ( + self.lora_B[self.active_adapter]( + self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x)) + ).to(expected_dtype) + * self.scaling[self.active_adapter] + ) else: - output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling - result += output - return result - - class MergedLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): - # Lora implemented in a dense layer - def __init__( - self, - in_features: int, - out_features: int, - r: int = 0, - lora_alpha: int = 1, - lora_dropout: float = 0.0, - enable_lora: List[bool] = [False], - **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) - if out_features % len(enable_lora) != 0: - raise ValueError("The length of enable_lora must divide out_features") - self.enable_lora = enable_lora - # Actual trainable parameters - if r > 0 and any(enable_lora): - self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False) - self.lora_B = nn.Conv1d( - r * sum(enable_lora), - out_features // len(enable_lora) * sum(enable_lora), - kernel_size=1, - groups=2, - bias=False, - ) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - # Compute the indices - self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1) - self.lora_ind[enable_lora, :] = True - self.lora_ind = self.lora_ind.view(-1) - 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.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) - - def zero_pad(self, x): - result = x.new_zeros((*x.shape[:-1], self.out_features)) - result = result.view(-1, self.out_features) - result[:, self.lora_ind] = x.reshape( - -1, self.out_features // len(self.enable_lora) * sum(self.enable_lora) - ) - return result.view((*x.shape[:-1], self.out_features)) - - 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() - after_A = self.lora_A(self.lora_dropout(x)) - after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) - output = self.zero_pad(after_B).to(expected_dtype) * self.scaling - result += output - else: - after_A = self.lora_A(self.lora_dropout(x)) - after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) - output = self.zero_pad(after_B) * self.scaling - result += output + output = ( + self.lora_B[self.active_adapter]( + self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x)) + ) + * self.scaling[self.active_adapter] + ) + result += output return result diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py index dd949c0..bfaabe8 100644 --- a/src/peft/utils/__init__.py +++ b/src/peft/utils/__init__.py @@ -17,14 +17,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .adapters_utils import CONFIG_NAME, WEIGHTS_NAME from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType from .other import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + CONFIG_NAME, + WEIGHTS_NAME, _set_trainable, bloom_model_postprocess_past_key_value, prepare_model_for_int8_training, shift_tokens_right, transpose, + _get_submodules, + _set_adapter, ) from .save_and_load import get_peft_model_state_dict, set_peft_model_state_dict diff --git a/src/peft/utils/adapters_utils.py b/src/peft/utils/adapters_utils.py deleted file mode 100644 index f2f8a95..0000000 --- a/src/peft/utils/adapters_utils.py +++ /dev/null @@ -1,18 +0,0 @@ -# coding=utf-8 -# Copyright 2023-present the HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -WEIGHTS_NAME = "adapter_model.bin" -CONFIG_NAME = "adapter_config.json" - -# TODO: add automapping and superclass here? diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index 3e2cf5b..544ec13 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -21,7 +21,7 @@ from typing import Optional, Union from huggingface_hub import hf_hub_download from transformers.utils import PushToHubMixin -from .adapters_utils import CONFIG_NAME +from .other import CONFIG_NAME class PeftType(str, enum.Enum): @@ -29,6 +29,7 @@ class PeftType(str, enum.Enum): P_TUNING = "P_TUNING" PREFIX_TUNING = "PREFIX_TUNING" LORA = "LORA" + MULTI_LORA = "MULTI_LORA" class TaskType(str, enum.Enum): @@ -82,7 +83,7 @@ class PeftConfigMixin(PushToHubMixin): writer.write(json.dumps(output_dict, indent=2, sort_keys=True)) @classmethod - def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None, **kwargs): r""" This method loads the configuration of your adapter model from a directory. @@ -92,11 +93,16 @@ class PeftConfigMixin(PushToHubMixin): **kwargs: Additional keyword arguments passed along to the child class initialization. """ - if os.path.isfile(os.path.join(pretrained_model_name_or_path, CONFIG_NAME)): - config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME) + path = ( + os.path.join(pretrained_model_name_or_path, subfolder) + if subfolder is not None + else pretrained_model_name_or_path + ) + if os.path.isfile(os.path.join(path, CONFIG_NAME)): + config_file = os.path.join(path, CONFIG_NAME) else: try: - config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME) + config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME, subfolder=subfolder) except Exception: raise ValueError(f"Can't find config.json at '{pretrained_model_name_or_path}'") diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 132b033..271bf89 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy + import torch @@ -86,11 +88,6 @@ def prepare_model_for_int8_training( return model -TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = { - "bloom": bloom_model_postprocess_past_key_value, -} - - # copied from transformers.models.bart.modeling_bart def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): """ @@ -113,11 +110,48 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start return shifted_input_ids -def _set_trainable(model): - if model.modules_to_save is not None: - for name, param in model.named_parameters(): - if any(module_name in name for module_name in model.modules_to_save): - param.requires_grad = True +class ModulesToSaveWrapper(torch.nn.Module): + def __init__(self, module_to_save, adapter_name): + super().__init__() + self.original_module = module_to_save + self.modules_to_save = torch.nn.ModuleDict({}) + self.update(adapter_name) + self.active_adapter = adapter_name + + def update(self, adapter_name): + self.modules_to_save.update(torch.nn.ModuleDict({adapter_name: copy.deepcopy(self.original_module)})) + + def forward(self, *args, **kwargs): + if self.active_adapter not in self.modules_to_save: + return self.original_module(*args, **kwargs) + return self.modules_to_save[self.active_adapter](*args, **kwargs) + + +def _get_submodules(model, key): + parent = model.get_submodule(".".join(key.split(".")[:-1])) + target_name = key.split(".")[-1] + target = model.get_submodule(key) + return parent, target, target_name + + +def _set_trainable(model, adapter_name): + key_list = [key for key, _ in model.named_modules()] + for key in key_list: + target_module_found = any(key.endswith(target_key) for target_key in model.modules_to_save) + if target_module_found: + parent, target, target_name = _get_submodules(key) + if isinstance(target, ModulesToSaveWrapper): + target.update(adapter_name) + else: + for param in target.parameters(): + param.requires_grad = True + setattr(parent, target_name, ModulesToSaveWrapper(target, adapter_name)) + + +def _set_adapter(model, adapter_name): + for module in model.modules(): + if isinstance(module, ModulesToSaveWrapper): + module.active_adapter = adapter_name def fsdp_auto_wrap_policy(model): @@ -157,3 +191,32 @@ def fsdp_auto_wrap_policy(model): def transpose(weight, fan_in_fan_out): return weight.T if fan_in_fan_out else weight + + +TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { + "t5": ["q", "v"], + "mt5": ["q", "v"], + "bart": ["q_proj", "v_proj"], + "gpt2": ["c_attn"], + "bloom": ["query_key_value"], + "opt": ["q_proj", "v_proj"], + "gptj": ["q_proj", "v_proj"], + "gpt_neox": ["query_key_value"], + "gpt_neo": ["q_proj", "v_proj"], + "bert": ["query", "value"], + "roberta": ["query", "value"], + "xlm-roberta": ["query", "value"], + "electra": ["query", "value"], + "deberta-v2": ["query_proj", "value_proj"], + "deberta": ["in_proj"], + "layoutlm": ["query", "value"], + "llama": ["q_proj", "v_proj"], + "chatglm": ["query_key_value"], +} + +TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = { + "bloom": bloom_model_postprocess_past_key_value, +} + +WEIGHTS_NAME = "adapter_model.bin" +CONFIG_NAME = "adapter_config.json" diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index c6596c7..43680fe 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -13,10 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .config import PeftType +from .config import PeftType, PromptLearningConfig -def get_peft_model_state_dict(model, state_dict=None): +def get_peft_model_state_dict(model, adapter_name, state_dict=None): """ Get the state dict of the Peft model. @@ -27,13 +27,14 @@ def get_peft_model_state_dict(model, state_dict=None): The state dict of the model. If not provided, the state dict of the model will be used. """ + config = model.peft_config[adapter_name] if state_dict is None: state_dict = model.state_dict() - if model.peft_config.peft_type == PeftType.LORA: + if config.peft_type == PeftType.LORA: # to_return = lora_state_dict(model, bias=model.peft_config.bias) # adapted from `https://github.com/microsoft/LoRA/blob/main/loralib/utils.py` - # to directly with the state dict which is necessary when using DeepSpeed or FSDP - bias = model.peft_config.bias + # to be used directly with the state dict which is necessary when using DeepSpeed or FSDP + bias = config.bias if bias == "none": to_return = {k: state_dict[k] for k in state_dict if "lora_" in k} elif bias == "all": @@ -48,21 +49,26 @@ def get_peft_model_state_dict(model, state_dict=None): to_return[bias_name] = state_dict[bias_name] else: raise NotImplementedError - else: + 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 = {} - if model.peft_config.inference_mode: + if config.inference_mode: prompt_embeddings = model.prompt_encoder.embedding.weight else: - prompt_embeddings = model.get_prompt_embedding_to_save() + prompt_embeddings = model.get_prompt_embedding_to_save(adapter_name) to_return["prompt_embeddings"] = prompt_embeddings + else: + raise NotImplementedError if model.modules_to_save is not None: for key, value in state_dict.items(): - if any(module_name in key for module_name in model.modules_to_save): - to_return[key] = value + if any(f"{module_name}.modules_to_save.{adapter_name}" in key for module_name in model.modules_to_save): + to_return[key.replace("modules_to_save.", "")] = value + + to_return = {k.replace(f"{adapter_name}.", ""): v for k, v in to_return.items()} return to_return -def set_peft_model_state_dict(model, peft_model_state_dict): +def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): """ Set the state dict of the Peft model. @@ -70,10 +76,33 @@ def set_peft_model_state_dict(model, peft_model_state_dict): model ([`PeftModel`]): The Peft model. peft_model_state_dict (`dict`): The state dict of the Peft model. """ + config = model.peft_config[adapter_name] + state_dict = {} + if model.modules_to_save is not None: + for key, value in peft_model_state_dict.items(): + if any(module_name in key for module_name in model.modules_to_save): + for module_name in model.modules_to_save: + if module_name in key: + key = key.replace(module_name, f"{module_name}.modules_to_save.{adapter_name}") + break + state_dict[key] = value + + if config.peft_type == PeftType.LORA: + peft_model_state_dict = {} + for k, v in state_dict.items(): + if "lora_" in k: + suffix_to_replace = ".".join(k.split("lora_")[1].split(".")[1:]) + k = k.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + peft_model_state_dict[k] = v + else: + peft_model_state_dict[k] = v + elif isinstance(config, PromptLearningConfig): + peft_model_state_dict = state_dict + else: + raise NotImplementedError model.load_state_dict(peft_model_state_dict, strict=False) - if model.peft_config.peft_type != PeftType.LORA: - model.prompt_encoder.embedding.load_state_dict( + if isinstance(config, PromptLearningConfig): + model.prompt_encoder[adapter_name].embedding.load_state_dict( {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True ) - return model From af252b709bd8f69dc47c10efb0c9798ec7f639de Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:29:24 +0530 Subject: [PATCH 13/76] Update peft_model.py --- src/peft/peft_model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index cbc8a89..b54a148 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -86,9 +86,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.active_adapter = adapter_name if not isinstance(peft_config, PromptLearningConfig): self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( - self.base_model, peft_config, adapter_name + self.base_model, self.peft_config, adapter_name ) - self.add_adapter(adapter_name, peft_config) + else: + self.add_adapter(adapter_name, peft_config) def save_pretrained(self, save_directory, **kwargs): r""" From 7d7c598647a68418705231286d1e6f90eb44646f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:32:21 +0530 Subject: [PATCH 14/76] Update peft_model.py --- src/peft/peft_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index b54a148..71b9ef2 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -85,6 +85,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.peft_config = {} self.active_adapter = adapter_name if not isinstance(peft_config, PromptLearningConfig): + self.peft_config[adapter_name] = peft_config self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( self.base_model, self.peft_config, adapter_name ) From 64cae2aab2174214987758f0486c3d6ae8a96563 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:34:04 +0530 Subject: [PATCH 15/76] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 34cee7d..5459381 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -139,7 +139,7 @@ class LoraModel(torch.nn.Module): if target_module_found: if not is_target_modules_in_base_model: is_target_modules_in_base_model = True - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(self.model, key) bias = target.bias is not None if isinstance(target, LoraLayer): target.update_layer(adapter_name, lora_config.r, lora_config.lora_alpha, lora_config.lora_dropout) From e9d45da4c5868c4a5b0624101a39bdebc2d4d138 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:35:49 +0530 Subject: [PATCH 16/76] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 5459381..215a315 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -361,7 +361,7 @@ class Linear(nn.Linear): self.weight.data = self.weight.data.T nn.Linear.reset_parameters(self) - self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.update_layer(adapter_name, r, lora_alpha, lora_dropout) self.active_adapter = adapter_name def merge(self): From 8ec7cb84350239d6ca51aa1223c91d40714f6a1b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:36:41 +0530 Subject: [PATCH 17/76] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 215a315..058a3f3 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -444,7 +444,7 @@ if is_bnb_available(): # Freezing the pre-trained weight matrix self.weight.requires_grad = False - self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.update_layer(adapter_name, r, lora_alpha, lora_dropout) self.active_adapter = adapter_name def forward(self, x: torch.Tensor): From 090d0743992c5ac43ae931b8e02d109b1b27d900 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:38:22 +0530 Subject: [PATCH 18/76] Update lora.py --- src/peft/tuners/lora.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 058a3f3..45d88ec 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -325,8 +325,8 @@ class LoraLayer: self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer})) # Actual trainable parameters if r > 0: - self.lora_A.update(nn.ModuleDict({nn.Linear(self.in_features, r, bias=False)})) - self.lora_B.update(nn.ModuleDict({nn.Linear(r, self.out_features, bias=False)})) + self.lora_A.update(nn.ModuleDict({adapter_name: nn.Linear(self.in_features, r, bias=False)})) + self.lora_B.update(nn.ModuleDict({adapter_name: nn.Linear(r, self.out_features, bias=False)})) self.scaling[adapter_name] = lora_alpha / r self.reset_lora_parameters(adapter_name) From 7c8ee5814a1a4ced386fc39cc13b3caf83a30930 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:40:06 +0530 Subject: [PATCH 19/76] Update peft_model.py --- src/peft/peft_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 71b9ef2..f2274fe 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -324,8 +324,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.add_adapter(adapter_name, peft_config) # load weights if any - if kwargs.get("subfolder", None) is not None: - path = os.path.join(model_id, kwargs["subfolder"]) + path = os.path.join(model_id, kwargs["subfolder"]) if kwargs.get("subfolder", None) is not None else model_id + if os.path.exists(os.path.join(path, WEIGHTS_NAME)): filename = os.path.join(path, WEIGHTS_NAME) else: From 002da1b450a1c3d4e206f9bd1556a45b7e90e8e6 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 20:19:06 +0530 Subject: [PATCH 20/76] fix bugs --- src/peft/tuners/lora.py | 1 + src/peft/utils/save_and_load.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 45d88ec..4739dce 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -329,6 +329,7 @@ class LoraLayer: self.lora_B.update(nn.ModuleDict({adapter_name: nn.Linear(r, self.out_features, bias=False)})) self.scaling[adapter_name] = lora_alpha / r 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(): diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 43680fe..fb4b252 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -86,6 +86,8 @@ def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): key = key.replace(module_name, f"{module_name}.modules_to_save.{adapter_name}") break state_dict[key] = value + else: + state_dict = peft_model_state_dict if config.peft_type == PeftType.LORA: peft_model_state_dict = {} @@ -100,7 +102,6 @@ def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): 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( From 1141b125d0e7f4590c2bb53b03faf5b5888b3399 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 29 Mar 2023 19:56:23 -0400 Subject: [PATCH 21/76] Impelment the budget finalization --- src/peft/tuners/adalora.py | 74 ++++++++++++++++++++++++++++++++------ 1 file changed, 64 insertions(+), 10 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 61bd512..796ae32 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -185,8 +185,63 @@ class AdaLoraModel(LoraModel): outputs.loss += orth_reg_weight * regu_loss return outputs + + def _prepare_new_module(self, target, rank_idx): + rank = rank_idx.sum().item() + 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, + } + loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) + if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt) and self.peft_config.enable_lora is None: + 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) and self.peft_config.enable_lora is None: + new_module = SVDLinear(target.in_features, target.out_features, bias=bias, **kwargs) + + with torch.no_grad(): + new_module = new_module.to(target.weight.device) + new_module.weight.copy_(target.weight) + if bias: + new_module.bias.copy_(target.bias) + if rank > 0: + rank_idx = rank_idx.view(-1) + 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 update_and_allocate(self, global_step): - self.rankallocator.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) + # 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) + + for name,rank_idx in rank_pattern.items(): + key = ".".join(name.split(".")[1:-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) + # Pass the function to do forward propagation + else: + return None @@ -463,25 +518,24 @@ class RankAllocator(object): 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.data.masked_fill_(triplet_ipt[n]<=mask_threshold, 0.0) - return mask_threshold + p.masked_fill_(triplet_ipt[n]<=mask_threshold, 0.0) + rank_pattern[n] = (~(triplet_ipt[n]<=mask_threshold)).to(p.device) + return rank_pattern def update_and_allocate(self, model, 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: self.update_ipt(model) - # TODO: Finalize the budget distribution by replacing with new Linear. budget, mask_ind = self.budget_schedule(global_step) - print("budget:", budget) if mask_ind: - mask_threshold = self.mask_to_budget(model, budget) - print("mask threshold:", mask_threshold) + rank_pattern = self.mask_to_budget(model, budget) else: - mask_threshold = None - return budget, mask_threshold + rank_pattern = None + return budget, rank_pattern From d6ae6650b2c909aad57fa0f7814e5d1c24691373 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Thu, 30 Mar 2023 00:45:33 +0000 Subject: [PATCH 22/76] Finish the test for rank finalization --- .../peft_adalora_seq2seq.py | 4 +- src/peft/tuners/adalora.py | 37 +++++++++++++------ 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index 49fa497..55c2161 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -20,7 +20,7 @@ text_column = "sentence" label_column = "text_label" max_length = 128 lr = 1e-3 -num_epochs = 8 +num_epochs = 2 batch_size = 8 @@ -28,7 +28,7 @@ batch_size = 8 peft_config = AdaLoraConfig( init_r=12, target_r=8, beta1=0.85, beta2=0.85, - tinit=200, tfinal=1000, deltaT=10, + tinit=2, tfinal=300, deltaT=10, lora_alpha=32, lora_dropout=0.1, task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 796ae32..648c05a 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -195,6 +195,7 @@ class AdaLoraModel(LoraModel): "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) and self.peft_config.enable_lora is None: kwargs.update( @@ -208,9 +209,9 @@ class AdaLoraModel(LoraModel): new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None: 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 = new_module.to(target.weight.device) new_module.weight.copy_(target.weight) if bias: new_module.bias.copy_(target.bias) @@ -221,7 +222,6 @@ class AdaLoraModel(LoraModel): 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 @@ -231,7 +231,7 @@ class AdaLoraModel(LoraModel): budget, rank_pattern = self.rankallocator.update_and_allocate(self, global_step) # 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) + budget, rank_pattern = self.rankallocator.update_and_allocate(self, global_step, force_mask=True) for name,rank_idx in rank_pattern.items(): key = ".".join(name.split(".")[1:-1]) @@ -239,7 +239,9 @@ class AdaLoraModel(LoraModel): new_module = self._prepare_new_module(target, rank_idx) self._replace_module(parent, target_name, new_module, target) - # Pass the function to do forward propagation + print("Finalize the rank pattern.") + self.rankallocator.reset_ipt() + # Pass the function and do forward propagation else: return None @@ -393,19 +395,23 @@ if is_bnb_available(): class RankAllocator(object): + """ + The RankAllocator for AdaLoraModel. + Paper: https://openreview.net/pdf?id=lq62uWRJjiY + + Args: + config ([`AdaLoraConfig`]): The configuration of the AdaLora model. + param_iterator: the parameter iterator to initalize the rankallocator. + + """ def __init__(self, peft_config, param_iterator): self.peft_config = peft_config - - self.ipt = {} - self.exp_avg_ipt = {} - self.exp_avg_unc = {} - self.cat_ipt = {} - 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(param_iterator) @@ -413,6 +419,12 @@ class RankAllocator(object): 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): self.init_bgt = 0 self.name_set = set() @@ -527,12 +539,13 @@ class RankAllocator(object): rank_pattern[n] = (~(triplet_ipt[n]<=mask_threshold)).to(p.device) return rank_pattern - def update_and_allocate(self, model, global_step): + 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) - if mask_ind: + print("budget:", budget) + if mask_ind or force_mask: rank_pattern = self.mask_to_budget(model, budget) else: rank_pattern = None From ce61e2452ae623f6cf41c755b51f35f1aff30dc4 Mon Sep 17 00:00:00 2001 From: Zhang Date: Wed, 29 Mar 2023 21:03:48 -0400 Subject: [PATCH 23/76] define the resize function --- src/peft/tuners/adalora.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 648c05a..ae5fe20 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -225,6 +225,15 @@ class AdaLoraModel(LoraModel): return new_module + def resize_modules_by_rank_pattern(self, rank_pattern): + for name,rank_idx in rank_pattern.items(): + key = ".".join(name.split(".")[1:-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: @@ -233,14 +242,17 @@ class AdaLoraModel(LoraModel): 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) - for name,rank_idx in rank_pattern.items(): - key = ".".join(name.split(".")[1:-1]) - parent, target, target_name = self._get_submodules(key) + self.resize_modules_by_rank_pattern(rank_pattern) - new_module = self._prepare_new_module(target, rank_idx) - self._replace_module(parent, target_name, new_module, target) + # for name,rank_idx in rank_pattern.items(): + # key = ".".join(name.split(".")[1:-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) print("Finalize the rank pattern.") self.rankallocator.reset_ipt() + self.rank_pattern = rank_pattern # Pass the function and do forward propagation else: return None From d3a48a891ed37ab961646d473183f04b6168d548 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Thu, 30 Mar 2023 01:04:11 +0000 Subject: [PATCH 24/76] save rank pattern --- src/peft/tuners/adalora.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 648c05a..0a323fc 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -55,6 +55,10 @@ class AdaLoraConfig(LoraConfig): 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 @@ -229,6 +233,8 @@ class AdaLoraModel(LoraModel): # 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) + 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) From 300abd1439dd14903ae3e55eae948372b6a456e7 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Thu, 30 Mar 2023 05:53:04 +0000 Subject: [PATCH 25/76] refine the key of rank pattern --- src/peft/tuners/adalora.py | 28 +++++++++++----------------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index d6343ff..f3edeba 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -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 = {} From e3b4cd46717d57145fa8a992af823199b962f34c Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Thu, 30 Mar 2023 01:59:26 -0400 Subject: [PATCH 26/76] Implement the save_pretrained for AdaLoRA --- src/peft/utils/save_and_load.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 86e388b..8be9314 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -70,8 +70,12 @@ def set_peft_model_state_dict(model, peft_model_state_dict): model ([`PeftModel`]): The Peft model. peft_model_state_dict (`dict`): The state dict of the Peft model. """ - + if model.peft_config.peft_type is PeftType.ADALORA: + rank_pattern = model.peft_config.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 From d4292300a0cd5ae7782134122a686ecf130f6b03 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Thu, 30 Mar 2023 06:19:45 +0000 Subject: [PATCH 27/76] Finish the test for model load and save --- .../conditional_generation/peft_adalora_seq2seq.py | 7 ++----- src/peft/tuners/adalora.py | 12 ++++-------- src/peft/utils/save_and_load.py | 2 +- 3 files changed, 7 insertions(+), 14 deletions(-) diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index 55c2161..391be4a 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -20,7 +20,7 @@ text_column = "sentence" label_column = "text_label" max_length = 128 lr = 1e-3 -num_epochs = 2 +num_epochs = 8 batch_size = 8 @@ -28,7 +28,7 @@ batch_size = 8 peft_config = AdaLoraConfig( init_r=12, target_r=8, beta1=0.85, beta2=0.85, - tinit=2, tfinal=300, deltaT=10, + tinit=200, tfinal=1000, deltaT=10, lora_alpha=32, lora_dropout=0.1, task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False @@ -37,7 +37,6 @@ peft_config = AdaLoraConfig( model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) model = get_peft_model(model, peft_config) model.print_trainable_parameters() -model # loading dataset @@ -53,8 +52,6 @@ dataset = dataset.map( num_proc=1, ) -dataset["train"][0] - # data preprocessing tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index f3edeba..5e032a6 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -39,6 +39,7 @@ class AdaLoraConfig(LoraConfig): 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."}) @@ -159,7 +160,6 @@ class AdaLoraModel(LoraModel): f"Please check the target modules and try again." ) - def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" try: @@ -167,7 +167,6 @@ class AdaLoraModel(LoraModel): except AttributeError: return getattr(self.model, name) - def forward(self, *args, **kwargs): outputs = self.model.forward(*args, **kwargs) @@ -197,7 +196,7 @@ class AdaLoraModel(LoraModel): rank_idx = rank_idx.view(-1) rank = rank_idx.sum().item() else: - raise ValueError("Unexcepted type of rank_idx") + raise ValueError(f"Unexcepted type of rank_idx") kwargs = { "r": rank, "lora_alpha": self.peft_config.lora_alpha, @@ -233,7 +232,6 @@ class AdaLoraModel(LoraModel): 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]) @@ -241,7 +239,6 @@ class AdaLoraModel(LoraModel): 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: @@ -256,7 +253,6 @@ class AdaLoraModel(LoraModel): self.resize_modules_by_rank_pattern(rank_pattern) self.peft_config.rank_pattern = rank_pattern self.rankallocator.reset_ipt() - print("Finalize the rank pattern.") # Pass the function and do forward propagation else: return None @@ -279,7 +275,7 @@ class SVDLinear(nn.Linear, LoraLayer): 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: @@ -552,7 +548,7 @@ class RankAllocator(object): if global_step < self.peft_config.total_step - self.peft_config.tfinal: self.update_ipt(model) budget, mask_ind = self.budget_schedule(global_step) - print("budget:", budget) + # Allocate the budget according to importance scores if mask_ind or force_mask: rank_pattern = self.mask_to_budget(model, budget) else: diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 8be9314..cf4c813 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -70,7 +70,7 @@ def set_peft_model_state_dict(model, peft_model_state_dict): model ([`PeftModel`]): The Peft model. peft_model_state_dict (`dict`): The state dict of the Peft model. """ - if model.peft_config.peft_type is PeftType.ADALORA: + if model.peft_config.peft_type == PeftType.ADALORA: rank_pattern = model.peft_config.rank_pattern if rank_pattern: model.base_model.resize_modules_by_rank_pattern(rank_pattern) From 7ef47be5f5f8d608773312d2f3e037f073f27e3b Mon Sep 17 00:00:00 2001 From: tpoisonooo Date: Mon, 3 Apr 2023 14:02:13 +0800 Subject: [PATCH 28/76] Update other.py typo --- src/peft/utils/other.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 132b033..585da64 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -34,7 +34,7 @@ def prepare_model_for_int8_training( model, output_embedding_layer_name="lm_head", use_gradient_checkpointing=True, layer_norm_names=["layer_norm"] ): r""" - This method wrapps the entire protocol for preparing a model before running a training. This includes: + This method wraps the entire protocol for preparing a model before running a training. This includes: 1- Cast the layernorm in fp32 2- making output embedding layer require grads 3- Add the upcasting of the lm head to fp32 From 39cbd7d8ed6b2fc56442ada66fba32898cfd00aa Mon Sep 17 00:00:00 2001 From: Guspan Tanadi <36249910+guspan-tanadi@users.noreply.github.com> Date: Mon, 3 Apr 2023 16:13:33 +0700 Subject: [PATCH 29/76] docs: have fix bit typo README Improve readability --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 5653c1a..0dc7c8a 100644 --- a/README.md +++ b/README.md @@ -64,7 +64,7 @@ Hardware: Single A100 80GB GPU with CPU RAM above 64GB | bigscience/bloomz-7b1 (7B params) | OOM GPU | 32GB GPU / 3.8GB CPU | 18.1GB GPU / 35GB CPU | Performance of PEFT-LoRA tuned [`bigscience/T0_3B`](https://huggingface.co/bigscience/T0_3B) on [`ought/raft/twitter_complaints`](https://huggingface.co/datasets/ought/raft/viewer/twitter_complaints) leaderboard. -A point to note is that we didn't try to sequeeze performance by playing around with input instruction templates, LoRA hyperparams and other training related hyperparams. Also, we didn't use the larger 13B [mt0-xxl](https://huggingface.co/bigscience/mt0-xxl) model. +A point to note is that we didn't try to squeeze performance by playing around with input instruction templates, LoRA hyperparams and other training related hyperparams. Also, we didn't use the larger 13B [mt0-xxl](https://huggingface.co/bigscience/mt0-xxl) model. So, we are already seeing comparable performance to SoTA with parameter efficient tuning. Also, the final checkpoint size is just `19MB` in comparison to `11GB` size of the backbone [`bigscience/T0_3B`](https://huggingface.co/bigscience/T0_3B) model. | Submission Name | Accuracy | @@ -81,7 +81,7 @@ GPU memory required by different settings during training is given below. The fi Hardware: Single A100 80GB GPU with CPU RAM above 64GB -| Model | Full Finetuning | PEFT-LoRA | PEFT-LoRA with Gradient Checkpoitning | +| Model | Full Finetuning | PEFT-LoRA | PEFT-LoRA with Gradient Checkpointing | | --------- | ---- | ---- | ---- | | CompVis/stable-diffusion-v1-4 | 27.5GB GPU / 3.97GB CPU | 15.5GB GPU / 3.84GB CPU | 8.12GB GPU / 3.77GB CPU | @@ -148,7 +148,7 @@ Another example is fine-tuning [`roberta-large`](https://huggingface.co/roberta- ## PEFT + 🤗 Accelerate -PEFT models work with 🤗 Accelerate out of the box. Use 🤗 Accelerate for Distributed training on various hardware such as GPUs, Apple Silicon devices etc during training. +PEFT models work with 🤗 Accelerate out of the box. Use 🤗 Accelerate for Distributed training on various hardware such as GPUs, Apple Silicon devices, etc during training. Use 🤗 Accelerate for inferencing on consumer hardware with small resources. ### Example of PEFT model training using 🤗 Accelerate's DeepSpeed integration From 4ddb85ce1e2a25d11e5c32e485e2348df792da2c Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Mon, 3 Apr 2023 17:08:42 +0200 Subject: [PATCH 30/76] Update stale.py --- scripts/stale.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/stale.py b/scripts/stale.py index a0bd10a..e910135 100644 --- a/scripts/stale.py +++ b/scripts/stale.py @@ -28,6 +28,7 @@ LABELS_TO_EXEMPT = [ "feature request", "new model", "wip", + "PRs welcome to address this", ] @@ -59,4 +60,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() From 45d7aab39a0a580201209709cbc38d248cec0193 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Mon, 3 Apr 2023 08:51:01 -0700 Subject: [PATCH 31/76] Update README.md --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 5653c1a..dffb656 100644 --- a/README.md +++ b/README.md @@ -25,10 +25,10 @@ Seamlessly integrated with 🤗 Accelerate for large scale models leveraging Dee Supported methods: -1. LoRA: [LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS](https://arxiv.org/pdf/2106.09685.pdf) +1. LoRA: [LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS](https://arxiv.org/abs/2106.09685) 2. Prefix Tuning: [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://aclanthology.org/2021.acl-long.353/), [P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks](https://arxiv.org/pdf/2110.07602.pdf) -3. P-Tuning: [GPT Understands, Too](https://arxiv.org/pdf/2103.10385.pdf) -4. Prompt Tuning: [The Power of Scale for Parameter-Efficient Prompt Tuning](https://arxiv.org/pdf/2104.08691.pdf) +3. P-Tuning: [GPT Understands, Too](https://arxiv.org/abs/2103.10385) +4. Prompt Tuning: [The Power of Scale for Parameter-Efficient Prompt Tuning](https://arxiv.org/abs/2104.08691) ## Getting started From ff9a1edbfd2d405b86d50a2e5299cc1bbd49d887 Mon Sep 17 00:00:00 2001 From: toncho11 Date: Mon, 3 Apr 2023 18:28:11 +0200 Subject: [PATCH 32/76] Fixing a bug where a wrong parameter name is used. --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index f9573bb..85757b7 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -164,7 +164,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if getattr(model, "hf_device_map", None) is not None: device_map = kwargs.get("device_map", "auto") max_memory = kwargs.get("max_memory", None) - offload_dir = kwargs.get("offload_dir", None) + offload_dir = kwargs.get("offload_folder", None) offload_index = kwargs.get("offload_index", None) dispatch_model_kwargs = {} From c2ef46f1454987f2fed2b7d45b655c75498022a5 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 07:58:48 +0000 Subject: [PATCH 33/76] v1 --- src/peft/peft_model.py | 69 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index f9573bb..fbb4867 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -1034,3 +1034,72 @@ class PeftModelForTokenClassification(PeftModel): hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) + + +class PeftModelForVision2Seq(PeftModel): + """ + Peft model for vision to text models. + + Args: + model ([`~transformers.PreTrainedModel`]): Base transformer model. + peft_config ([`PeftConfig`]): Peft config. + + + Example: + + ```py + >>> from transformers import AutoModelForVision2Seq + >>> from peft import PeftModelForVision2Seq, get_peft_config + + >>> config = { + ... "peft_type": "LORA", + ... "task_type": "VISION_2_SEQ", + ... "inference_mode": False, + ... "r": 8, + ... "target_modules": ["q", "v"], + ... "lora_alpha": 32, + ... "lora_dropout": 0.1, + ... "merge_weights": False, + ... "fan_in_fan_out": False, + ... "enable_lora": None, + ... "bias": "none", + ... } + + >>> peft_config = get_peft_config(config) + >>> model = AutoModelForCausalLM.from_pretrained("Salesforce/blip2-flan-t5-xl") + >>> peft_model = PeftModelForVision2Seq(model, peft_config) + >>> peft_model.print_trainable_parameters() + trainable params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544 + ``` + """ + + def __init__(self, model, peft_config: PeftConfig): + super().__init__(model, peft_config) + self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation + + def forward( + self, + pixel_values=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + r""" + A simple wrapper around the base model's forward method. + """ + return self.base_model( + pixel_values=pixel_values, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + labels=labels, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) \ No newline at end of file From c7e22ccd757c7f5ba5e459bcf86416fd803c3afa Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 07:59:03 +0000 Subject: [PATCH 34/76] v1 --- src/peft/mapping.py | 8 +++++++- src/peft/tuners/lora.py | 11 +++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/peft/mapping.py b/src/peft/mapping.py index dbb9f36..1e98edb 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -19,6 +19,7 @@ from .peft_model import ( PeftModelForSeq2SeqLM, PeftModelForSequenceClassification, PeftModelForTokenClassification, + PeftModelForVision2Seq, ) from .tuners import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from .utils import PromptLearningConfig @@ -29,6 +30,7 @@ MODEL_TYPE_TO_PEFT_MODEL_MAPPING = { "SEQ_2_SEQ_LM": PeftModelForSeq2SeqLM, "CAUSAL_LM": PeftModelForCausalLM, "TOKEN_CLS": PeftModelForTokenClassification, + "VISION_2_SEQ": PeftModelForVision2Seq, } PEFT_TYPE_TO_CONFIG_MAPPING = { @@ -44,6 +46,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "bart": ["q_proj", "v_proj"], "gpt2": ["c_attn"], "bloom": ["query_key_value"], + "blip2": ["q", "v", "q_proj", "v_proj"], "opt": ["q_proj", "v_proj"], "gptj": ["q_proj", "v_proj"], "gpt_neox": ["query_key_value"], @@ -134,9 +137,12 @@ def get_peft_model(model, peft_config): model ([`transformers.PreTrainedModel`]): Model to be wrapped. peft_config ([`PeftConfig`]): Configuration object containing the parameters of the Peft model. """ - model_config = model.config.to_dict() peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None) + + if peft_config.task_type == "VISION_2_SEQ" and not isinstance(peft_config, LoraConfig): + raise ValueError("Vision2Seq task type is only supported with LORA") + if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): peft_config = _prepare_lora_config(peft_config, model_config) return PeftModel(model, peft_config) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 51cd56f..d4d17ac 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -394,6 +394,7 @@ class Linear(nn.Linear, LoraLayer): self.lora_B.eval() def forward(self, x: torch.Tensor): + if self.disable_adapters: if self.r > 0 and self.merged: self.weight.data -= ( @@ -401,14 +402,20 @@ class Linear(nn.Linear, LoraLayer): ) self.merged = False - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + + return result elif 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: + x = x.to(self.lora_A.weight.dtype) + result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling return result else: - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + + return result class MergedLinear(nn.Linear, LoraLayer): From af6794e424facafe2e390339fd7fce791f84ee59 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 08:18:47 +0000 Subject: [PATCH 35/76] add blip2 --- README.md | 6 ++ .../int8_training/fine_tune_blip2_int8.py | 88 +++++++++++++++++++ src/peft/peft_model.py | 2 +- src/peft/tuners/lora.py | 1 - 4 files changed, 95 insertions(+), 2 deletions(-) create mode 100644 examples/int8_training/fine_tune_blip2_int8.py diff --git a/README.md b/README.md index ccdcd55..af84399 100644 --- a/README.md +++ b/README.md @@ -274,6 +274,12 @@ An example is provided in `~examples/causal_language_modeling/peft_lora_clm_acce | ViT | ✅ | | | | | Swin | ✅ | | | | +### Image to text (Multi-modal models) + +| Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning | +| --------- | ---- | ---- | ---- | ---- | +| Blip-2 | ✅ | | | | + ___Note that we have tested LoRA for [ViT](https://huggingface.co/docs/transformers/model_doc/vit) and [Swin](https://huggingface.co/docs/transformers/model_doc/swin) for fine-tuning on image classification. However, it should be possible to use LoRA for any compatible model [provided](https://huggingface.co/models?pipeline_tag=image-classification&sort=downloads&search=vit) by 🤗 Transformers. Check out the respective examples to learn more. If you run into problems, please open an issue.___ diff --git a/examples/int8_training/fine_tune_blip2_int8.py b/examples/int8_training/fine_tune_blip2_int8.py new file mode 100644 index 0000000..526336a --- /dev/null +++ b/examples/int8_training/fine_tune_blip2_int8.py @@ -0,0 +1,88 @@ +import torch +from datasets import load_dataset +from torch.utils.data import DataLoader, Dataset +from transformers import AutoModelForVision2Seq, AutoProcessor + +from peft import LoraConfig, get_peft_model + + +config = LoraConfig( + r=16, + lora_alpha=32, + target_modules=["q_proj", "v_proj"], + lora_dropout=0.05, + bias="none", + task_type="VISION_2_SEQ", +) + +model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True, device_map={"": 0}) +processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b") +model = get_peft_model(model, config) + +model.print_trainable_parameters() + +dataset = load_dataset("ybelkada/football-dataset", split="train") + + +class ImageCaptioningDataset(Dataset): + def __init__(self, dataset, processor): + self.dataset = dataset + self.processor = processor + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, idx): + item = self.dataset[idx] + encoding = self.processor(images=item["image"], padding="max_length", return_tensors="pt") + # remove batch dimension + encoding = {k: v.squeeze() for k, v in encoding.items()} + encoding["text"] = item["text"] + return encoding + + +def collator(batch): + # pad the input_ids and attention_mask + processed_batch = {} + for key in batch[0].keys(): + if key != "text": + processed_batch[key] = torch.stack([example[key] for example in batch]) + else: + text_inputs = processor.tokenizer( + [example["text"] for example in batch], padding=True, return_tensors="pt" + ) + processed_batch["input_ids"] = text_inputs["input_ids"] + processed_batch["attention_mask"] = text_inputs["attention_mask"] + return processed_batch + + +train_dataset = ImageCaptioningDataset(dataset, processor) +train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=2, collate_fn=collator) + +optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5) + +device = "cuda" if torch.cuda.is_available() else "cpu" +model.to(device) + +model.train() + +for epoch in range(50): + print("Epoch:", epoch) + for idx, batch in enumerate(train_dataloader): + input_ids = batch.pop("input_ids").to(device) + pixel_values = batch.pop("pixel_values").to(device, torch.float16) + + outputs = model(input_ids=input_ids, pixel_values=pixel_values, labels=input_ids) + + loss = outputs.loss + + print("Loss:", loss.item()) + + loss.backward() + + optimizer.step() + optimizer.zero_grad() + + if idx % 10 == 0: + generated_output = model.generate(pixel_values=pixel_values) + print(processor.batch_decode(generated_output, skip_special_tokens=True)) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 79d7464..0305b79 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -1102,4 +1102,4 @@ class PeftModelForVision2Seq(PeftModel): output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs, - ) \ No newline at end of file + ) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index d4d17ac..1674754 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -394,7 +394,6 @@ class Linear(nn.Linear, LoraLayer): self.lora_B.eval() def forward(self, x: torch.Tensor): - if self.disable_adapters: if self.r > 0 and self.merged: self.weight.data -= ( From f569bc682bb1998598f645ca9b5c98f847f5e90c Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Tue, 4 Apr 2023 10:21:38 +0200 Subject: [PATCH 36/76] Update src/peft/peft_model.py --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 0305b79..6706b23 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -1066,7 +1066,7 @@ class PeftModelForVision2Seq(PeftModel): ... } >>> peft_config = get_peft_config(config) - >>> model = AutoModelForCausalLM.from_pretrained("Salesforce/blip2-flan-t5-xl") + >>> model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip2-flan-t5-xl") >>> peft_model = PeftModelForVision2Seq(model, peft_config) >>> peft_model.print_trainable_parameters() trainable params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544 From 46ab59628cd75158ef232e9a63999626a9d9d949 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 08:23:47 +0000 Subject: [PATCH 37/76] revert --- src/peft/tuners/lora.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 1674754..0553f18 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -412,9 +412,8 @@ class Linear(nn.Linear, LoraLayer): result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling return result else: - result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - return result class MergedLinear(nn.Linear, LoraLayer): From 96cd0390367f73ba494e23e6cd51b75b3df17a16 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 08:29:51 +0000 Subject: [PATCH 38/76] fix --- examples/int8_training/fine_tune_blip2_int8.py | 1 - src/peft/mapping.py | 2 +- src/peft/tuners/lora.py | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/int8_training/fine_tune_blip2_int8.py b/examples/int8_training/fine_tune_blip2_int8.py index 526336a..1e73ffb 100644 --- a/examples/int8_training/fine_tune_blip2_int8.py +++ b/examples/int8_training/fine_tune_blip2_int8.py @@ -9,7 +9,6 @@ from peft import LoraConfig, get_peft_model config = LoraConfig( r=16, lora_alpha=32, - target_modules=["q_proj", "v_proj"], lora_dropout=0.05, bias="none", task_type="VISION_2_SEQ", diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 1e98edb..4abbd5a 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -46,7 +46,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "bart": ["q_proj", "v_proj"], "gpt2": ["c_attn"], "bloom": ["query_key_value"], - "blip2": ["q", "v", "q_proj", "v_proj"], + "blip-2": ["q", "v", "q_proj", "v_proj"], "opt": ["q_proj", "v_proj"], "gptj": ["q_proj", "v_proj"], "gpt_neox": ["query_key_value"], diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 0553f18..6fae36e 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -415,7 +415,6 @@ class Linear(nn.Linear, LoraLayer): return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - class MergedLinear(nn.Linear, LoraLayer): # Lora implemented in a dense layer def __init__( From 4cbd6cfd43c76d4762dcf93ed97cfa803b61a84f Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 08:31:37 +0000 Subject: [PATCH 39/76] revert --- src/peft/tuners/lora.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 6fae36e..9745467 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -401,9 +401,7 @@ class Linear(nn.Linear, LoraLayer): ) self.merged = False - result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - - return result + return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) elif 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: From 8c83386ef413174bb2ee7167d1d13fa0fe00516c Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 08:37:32 +0000 Subject: [PATCH 40/76] few fixes --- .../int8_training/fine_tune_blip2_int8.py | 21 +++++++++++++++++-- src/peft/tuners/lora.py | 2 -- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/examples/int8_training/fine_tune_blip2_int8.py b/examples/int8_training/fine_tune_blip2_int8.py index 1e73ffb..25121f6 100644 --- a/examples/int8_training/fine_tune_blip2_int8.py +++ b/examples/int8_training/fine_tune_blip2_int8.py @@ -1,3 +1,17 @@ +# coding=utf-8 +# Copyright 2023-present the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import torch from datasets import load_dataset from torch.utils.data import DataLoader, Dataset @@ -6,6 +20,7 @@ from transformers import AutoModelForVision2Seq, AutoProcessor from peft import LoraConfig, get_peft_model +# Let's define the LoraConfig config = LoraConfig( r=16, lora_alpha=32, @@ -14,12 +29,15 @@ config = LoraConfig( task_type="VISION_2_SEQ", ) +# We load our model and processor using `transformers` model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True, device_map={"": 0}) processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b") -model = get_peft_model(model, config) +# Get our peft model and print the number of trainable parameters +model = get_peft_model(model, config) model.print_trainable_parameters() +# Let's load the dataset here! dataset = load_dataset("ybelkada/football-dataset", split="train") @@ -61,7 +79,6 @@ train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=2, collate optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5) device = "cuda" if torch.cuda.is_available() else "cpu" -model.to(device) model.train() diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 9745467..51cd56f 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -405,8 +405,6 @@ class Linear(nn.Linear, LoraLayer): elif 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: - x = x.to(self.lora_A.weight.dtype) - result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling return result else: From 7ed9ad04bfa775397e9cb4f112f8d6f04df6c4e0 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Tue, 4 Apr 2023 11:01:07 +0000 Subject: [PATCH 41/76] revert changes --- .../int8_training/fine_tune_blip2_int8.py | 1 - src/peft/mapping.py | 5 -- src/peft/peft_model.py | 69 ------------------- 3 files changed, 75 deletions(-) diff --git a/examples/int8_training/fine_tune_blip2_int8.py b/examples/int8_training/fine_tune_blip2_int8.py index 25121f6..ca6ba40 100644 --- a/examples/int8_training/fine_tune_blip2_int8.py +++ b/examples/int8_training/fine_tune_blip2_int8.py @@ -26,7 +26,6 @@ config = LoraConfig( lora_alpha=32, lora_dropout=0.05, bias="none", - task_type="VISION_2_SEQ", ) # We load our model and processor using `transformers` diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 4abbd5a..35a6901 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -19,7 +19,6 @@ from .peft_model import ( PeftModelForSeq2SeqLM, PeftModelForSequenceClassification, PeftModelForTokenClassification, - PeftModelForVision2Seq, ) from .tuners import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from .utils import PromptLearningConfig @@ -30,7 +29,6 @@ MODEL_TYPE_TO_PEFT_MODEL_MAPPING = { "SEQ_2_SEQ_LM": PeftModelForSeq2SeqLM, "CAUSAL_LM": PeftModelForCausalLM, "TOKEN_CLS": PeftModelForTokenClassification, - "VISION_2_SEQ": PeftModelForVision2Seq, } PEFT_TYPE_TO_CONFIG_MAPPING = { @@ -140,9 +138,6 @@ def get_peft_model(model, peft_config): model_config = model.config.to_dict() peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None) - if peft_config.task_type == "VISION_2_SEQ" and not isinstance(peft_config, LoraConfig): - raise ValueError("Vision2Seq task type is only supported with LORA") - if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): peft_config = _prepare_lora_config(peft_config, model_config) return PeftModel(model, peft_config) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 6706b23..85757b7 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -1034,72 +1034,3 @@ class PeftModelForTokenClassification(PeftModel): hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) - - -class PeftModelForVision2Seq(PeftModel): - """ - Peft model for vision to text models. - - Args: - model ([`~transformers.PreTrainedModel`]): Base transformer model. - peft_config ([`PeftConfig`]): Peft config. - - - Example: - - ```py - >>> from transformers import AutoModelForVision2Seq - >>> from peft import PeftModelForVision2Seq, get_peft_config - - >>> config = { - ... "peft_type": "LORA", - ... "task_type": "VISION_2_SEQ", - ... "inference_mode": False, - ... "r": 8, - ... "target_modules": ["q", "v"], - ... "lora_alpha": 32, - ... "lora_dropout": 0.1, - ... "merge_weights": False, - ... "fan_in_fan_out": False, - ... "enable_lora": None, - ... "bias": "none", - ... } - - >>> peft_config = get_peft_config(config) - >>> model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip2-flan-t5-xl") - >>> peft_model = PeftModelForVision2Seq(model, peft_config) - >>> peft_model.print_trainable_parameters() - trainable params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544 - ``` - """ - - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) - self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation - - def forward( - self, - pixel_values=None, - attention_mask=None, - decoder_input_ids=None, - decoder_attention_mask=None, - labels=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - **kwargs, - ): - r""" - A simple wrapper around the base model's forward method. - """ - return self.base_model( - pixel_values=pixel_values, - attention_mask=attention_mask, - decoder_input_ids=decoder_input_ids, - decoder_attention_mask=decoder_attention_mask, - labels=labels, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - **kwargs, - ) From bd80d61b2a77f439c2e3428151c164840a5c020d Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 17:44:24 +0530 Subject: [PATCH 42/76] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 126 ++++++++++++++++++++++------------------ src/peft/tuners/lora.py | 12 +--- 2 files changed, 70 insertions(+), 68 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 4df5cd9..b98fc9f 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -84,6 +84,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): 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 if not isinstance(peft_config, PromptLearningConfig): self.peft_config[adapter_name] = peft_config self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( @@ -120,10 +121,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if peft_config.base_model_name_or_path is None: peft_config.base_model_name_or_path = ( self.base_model.__dict__.get("name_or_path", None) - if isinstance(self.peft_config, PromptLearningConfig) + if isinstance(peft_config, PromptLearningConfig) else self.base_model.model.__dict__.get("name_or_path", None) ) - inference_mode = self.peft_config.inference_mode + inference_mode = peft_config.inference_mode peft_config.inference_mode = True peft_config.save_pretrained(output_dir) peft_config.inference_mode = inference_mode @@ -213,32 +214,33 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Returns the virtual prompts to use for Peft. Only applicable when `peft_config.peft_type != PeftType.LORA`. """ + peft_config = self.active_peft_config prompt_encoder = self.prompt_encoder[self.active_adapter] prompt_tokens = self.prompt_tokens[self.active_adapter].unsqueeze(0).expand(batch_size, -1).to(self.device) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: - prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens] - if self.peft_config.inference_mode: + if peft_config.peft_type == PeftType.PREFIX_TUNING: + prompt_tokens = prompt_tokens[:, : peft_config.num_virtual_tokens] + if peft_config.inference_mode: past_key_values = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: past_key_values = prompt_encoder(prompt_tokens) past_key_values = past_key_values.view( batch_size, - self.peft_config.num_virtual_tokens, - self.peft_config.num_layers * 2, - self.peft_config.num_attention_heads, - self.peft_config.token_dim // self.peft_config.num_attention_heads, + peft_config.num_virtual_tokens, + peft_config.num_layers * 2, + peft_config.num_attention_heads, + peft_config.token_dim // peft_config.num_attention_heads, ) - if self.peft_config.num_transformer_submodules == 2: + if peft_config.num_transformer_submodules == 2: past_key_values = torch.cat([past_key_values, past_key_values], dim=2) past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split( - self.peft_config.num_transformer_submodules * 2 + peft_config.num_transformer_submodules * 2 ) if TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING.get(self.config.model_type, None) is not None: post_process_fn = TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING[self.config.model_type] past_key_values = post_process_fn(past_key_values) return past_key_values else: - if self.peft_config.inference_mode: + if peft_config.inference_mode: prompts = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: prompts = prompt_encoder(prompt_tokens) @@ -281,13 +283,13 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Disables the adapter module. """ - if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): + if isinstance(self.active_peft_config, PromptLearningConfig): old_forward = self.forward self.forward = self.base_model.forward else: self.base_model.disable_adapter_layers() yield - if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): + if isinstance(self.active_peft_config, PromptLearningConfig): self.forward = old_forward else: self.base_model.enable_adapter_layers() @@ -296,13 +298,14 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Returns the base model. """ - return ( - self.base_model - if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig) - else self.base_model.model - ) + return self.base_model if isinstance(self.active_peft_config, PromptLearningConfig) else self.base_model.model def add_adapter(self, adapter_name, peft_config): + if peft_config.peft_type != self.peft_type: + raise ValueError( + f"Cannot combine adapters with different peft types. " + f"Found {self.peft_type} and {peft_config.peft_type}." + ) self.peft_config[adapter_name] = peft_config if isinstance(peft_config, PromptLearningConfig): self._setup_prompt_encoder(adapter_name) @@ -380,11 +383,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): **dispatch_model_kwargs, ) hook = AlignDevicesHook(io_same_device=True) - if not isinstance(self.peft_config[adapter_name]) == PeftType.LORA: - add_hook_to_module(self.base_model.model, hook) - else: + if isinstance(self.peft_config[adapter_name], PromptLearningConfig): remove_hook_from_submodules(self.prompt_encoder) - add_hook_to_module(self.base_model, hook) + add_hook_to_module(self.get_base_model(), hook) def set_adapter(self, adapter_name): """ @@ -397,6 +398,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.base_model.set_adapter(adapter_name) _set_adapter(self, adapter_name) + @property + def active_peft_config(self): + return self.peft_config[self.active_adapter] + class PeftModelForSequenceClassification(PeftModel): """ @@ -465,8 +470,8 @@ class PeftModelForSequenceClassification(PeftModel): **kwargs, ): return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if not isinstance(self.peft_config, PromptLearningConfig): + peft_config = self.active_peft_config + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -481,7 +486,7 @@ class PeftModelForSequenceClassification(PeftModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") @@ -496,13 +501,13 @@ class PeftModelForSequenceClassification(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: return self._prefix_tuning_forward(input_ids=input_ids, **kwargs) else: if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( - torch.zeros(batch_size, self.peft_config.num_virtual_tokens).to(self.device), + torch.zeros(batch_size, peft_config.num_virtual_tokens).to(self.device), kwargs["token_type_ids"], ), dim=1, @@ -638,7 +643,8 @@ class PeftModelForCausalLM(PeftModel): return_dict=None, **kwargs, ): - if not isinstance(self.peft_config, PromptLearningConfig): + peft_config = self.active_peft_config + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -653,7 +659,7 @@ class PeftModelForCausalLM(PeftModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -672,7 +678,7 @@ class PeftModelForCausalLM(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.base_model(input_ids=input_ids, past_key_values=past_key_values, **kwargs) else: @@ -680,7 +686,7 @@ class PeftModelForCausalLM(PeftModel): inputs_embeds = self.word_embeddings(input_ids) # concat prompt labels if labels is not None: - prefix_labels = torch.full((batch_size, self.peft_config.num_virtual_tokens), -100).to(self.device) + prefix_labels = torch.full((batch_size, peft_config.num_virtual_tokens), -100).to(self.device) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) prompts = prompts.to(inputs_embeds.dtype) @@ -688,9 +694,10 @@ class PeftModelForCausalLM(PeftModel): return self.base_model(inputs_embeds=inputs_embeds, **kwargs) def generate(self, **kwargs): + peft_config = self.active_peft_config self.base_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation try: - if not isinstance(self.peft_config, PromptLearningConfig): + if not isinstance(peft_config, PromptLearningConfig): outputs = self.base_model.generate(**kwargs) else: if "input_ids" not in kwargs: @@ -698,13 +705,13 @@ class PeftModelForCausalLM(PeftModel): # For gpt2 models, we construct postion_ids on the fly by using attention mask, and position ids need to match input_shape. # for prefix tuning, input shape is determined using `input_ids`. Thus we should not expand 'attention_mask' here # for prompt tuning input_ids is not passed but a concatenated input_embeds is passed. Thus attention_mask needs to be of same size of num_virtual_tokens + input_ids - if kwargs.get("attention_mask", None) is not None and self.peft_config.peft_type in [ + if kwargs.get("attention_mask", None) is not None and peft_config.peft_type in [ PeftType.PROMPT_TUNING, PeftType.P_TUNING, ]: # concat prompt attention mask prefix_attention_mask = torch.ones( - kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens + kwargs["input_ids"].shape[0], peft_config.num_virtual_tokens ).to(kwargs["input_ids"].device) kwargs["attention_mask"] = torch.cat((prefix_attention_mask, kwargs["attention_mask"]), dim=1) @@ -728,17 +735,18 @@ class PeftModelForCausalLM(PeftModel): return outputs def prepare_inputs_for_generation(self, *args, **kwargs): + peft_config = self.active_peft_config model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs) - if isinstance(self.peft_config, PromptLearningConfig): - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if isinstance(peft_config, PromptLearningConfig): + if peft_config.peft_type == PeftType.PREFIX_TUNING: prefix_attention_mask = torch.ones( - model_kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens + model_kwargs["input_ids"].shape[0], peft_config.num_virtual_tokens ).to(model_kwargs["input_ids"].device) model_kwargs["attention_mask"] = torch.cat( (prefix_attention_mask, model_kwargs["attention_mask"]), dim=1 ) - if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if model_kwargs["past_key_values"] is None and peft_config.peft_type == PeftType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0]) model_kwargs["past_key_values"] = past_key_values else: @@ -810,7 +818,8 @@ class PeftModelForSeq2SeqLM(PeftModel): return_dict=None, **kwargs, ): - if not isinstance(self.peft_config, PromptLearningConfig): + peft_config = self.active_peft_config + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -828,7 +837,7 @@ class PeftModelForSeq2SeqLM(PeftModel): batch_size = input_ids.shape[0] if decoder_attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -848,7 +857,7 @@ class PeftModelForSeq2SeqLM(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.base_model( input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs @@ -864,35 +873,36 @@ class PeftModelForSeq2SeqLM(PeftModel): if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1) # concat prompt labels if labels is not None: - if self.peft_config.num_transformer_submodules == 1: + if peft_config.num_transformer_submodules == 1: kwargs["labels"] = labels - elif self.peft_config.num_transformer_submodules == 2: - prefix_labels = torch.full((batch_size, self.peft_config.num_virtual_tokens), -100).to(self.device) + elif peft_config.num_transformer_submodules == 2: + prefix_labels = torch.full((batch_size, peft_config.num_virtual_tokens), -100).to(self.device) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) prompts = prompts.to(inputs_embeds.dtype) - inputs_embeds = torch.cat((prompts[:, : self.peft_config.num_virtual_tokens], inputs_embeds), dim=1) - if self.peft_config.num_transformer_submodules == 1: + inputs_embeds = torch.cat((prompts[:, : peft_config.num_virtual_tokens], inputs_embeds), dim=1) + if peft_config.num_transformer_submodules == 1: return self.base_model(inputs_embeds=inputs_embeds, **kwargs) - elif self.peft_config.num_transformer_submodules == 2: + elif peft_config.num_transformer_submodules == 2: decoder_inputs_embeds = torch.cat( - (prompts[:, self.peft_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1 + (prompts[:, peft_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1 ) return self.base_model( inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs ) def generate(self, **kwargs): + peft_config = self.active_peft_config self.base_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation self.base_model._prepare_encoder_decoder_kwargs_for_generation = ( self._prepare_encoder_decoder_kwargs_for_generation ) try: - if not isinstance(self.peft_config, PromptLearningConfig): + if not isinstance(peft_config, PromptLearningConfig): outputs = self.base_model.generate(**kwargs) else: if "input_ids" not in kwargs: @@ -908,7 +918,7 @@ class PeftModelForSeq2SeqLM(PeftModel): ) kwargs["token_type_ids"] = None - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: outputs = self.base_model.generate(**kwargs) else: raise NotImplementedError @@ -926,8 +936,9 @@ class PeftModelForSeq2SeqLM(PeftModel): return outputs def prepare_inputs_for_generation(self, *args, **kwargs): + peft_config = self.active_peft_config model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs) - if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if model_kwargs["past_key_values"] is None and peft_config.peft_type == PeftType.PREFIX_TUNING: batch_size = model_kwargs["decoder_input_ids"].shape[0] past_key_values = self.get_prompt(batch_size) model_kwargs["past_key_values"] = past_key_values @@ -1000,9 +1011,10 @@ class PeftModelForTokenClassification(PeftModel): return_dict=None, **kwargs, ): + peft_config = self.active_peft_config return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if not isinstance(self.peft_config, PromptLearningConfig): + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -1017,7 +1029,7 @@ class PeftModelForTokenClassification(PeftModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") @@ -1032,13 +1044,13 @@ class PeftModelForTokenClassification(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: return self._prefix_tuning_forward(input_ids=input_ids, **kwargs) else: if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( - torch.zeros(batch_size, self.peft_config.num_virtual_tokens).to(self.device), + torch.zeros(batch_size, peft_config.num_virtual_tokens).to(self.device), kwargs["token_type_ids"], ), dim=1, diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 977a3b7..0b96a77 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -323,17 +323,7 @@ class LoraModel(torch.nn.Module): if isinstance(target, LoraLayer): bias = target.bias is not None new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias) - - # manually merge if not merged - if not target.merged: - # merge weights per: https://arxiv.org/pdf/2106.09685.pdf / page 4 - if target.r > 0: - target.weight.data += ( - transpose(target.lora_B.weight @ target.lora_A.weight, target.fan_in_fan_out) - * target.scaling - ).to(target.weight.dtype) - target.merged = True - + target.merge() self._replace_module(parent, target_name, new_module, target) return self.model From 96ca100e34e1d0207a517d34ccc9a73fb4abf69e Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 18:03:32 +0530 Subject: [PATCH 43/76] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 0b96a77..aab5cdf 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -166,7 +166,7 @@ class LoraModel(torch.nn.Module): "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": self.peft_config.init_lora_weights, + "init_lora_weights": lora_config.init_lora_weights, } key_list = [key for key, _ in self.model.named_modules()] for key in key_list: From d4b64c82801b9a939bd909d94d147307c07c7926 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 18:27:23 +0530 Subject: [PATCH 44/76] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index aab5cdf..417bb7a 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -434,6 +434,8 @@ class Linear(nn.Linear): self.active_adapter = adapter_name 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 @@ -451,6 +453,8 @@ class Linear(nn.Linear): self.merged = True 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 @@ -468,6 +472,8 @@ class Linear(nn.Linear): 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() @@ -520,7 +526,7 @@ if is_bnb_available(): 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[self.active_adapter] > 0: if not torch.is_autocast_enabled(): From 18ccde8e86a1dcebc1d2ddf85350d09a341342d2 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 19:44:58 +0530 Subject: [PATCH 45/76] =?UTF-8?q?fixing=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 6 ++++-- src/peft/tuners/lora.py | 2 +- src/peft/utils/save_and_load.py | 4 ++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index b98fc9f..1c3e7e1 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -112,7 +112,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): for adapter_name, peft_config in self.peft_config.items(): # save only the trainable weights - output_state_dict = get_peft_model_state_dict(self, adapter_name, kwargs.get("state_dict", None)) + output_state_dict = get_peft_model_state_dict( + self, state_dict=kwargs.get("state_dict", None), adapter_name=adapter_name + ) output_dir = os.path.join(save_directory, adapter_name) if adapter_name != "default" else save_directory os.makedirs(output_dir, exist_ok=True) torch.save(output_state_dict, os.path.join(output_dir, WEIGHTS_NAME)) @@ -346,7 +348,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu") ) # load the weights into the model - set_peft_model_state_dict(self, adapter_name, adapters_weights) + set_peft_model_state_dict(self, adapters_weights, adapter_name=adapter_name) if ( (getattr(self, "hf_device_map", None) is not None) and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 417bb7a..0023717 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -136,7 +136,7 @@ class LoraModel(torch.nn.Module): self.model = model self.forward = self.model.forward self.config = config - self.add_adapter(adapter_name) + self.add_adapter(adapter_name, self.config[adapter_name]) def add_adapter(self, adapter_name, config=None): if config is not None: diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index fb4b252..7ebdfaf 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -16,7 +16,7 @@ from .config import PeftType, PromptLearningConfig -def get_peft_model_state_dict(model, adapter_name, state_dict=None): +def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): """ Get the state dict of the Peft model. @@ -68,7 +68,7 @@ def get_peft_model_state_dict(model, adapter_name, state_dict=None): return to_return -def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): +def set_peft_model_state_dict(model, peft_model_state_dict, adapter_name="default"): """ Set the state dict of the Peft model. From 122f708ae8551f5b88cfbe3b9579fcdb3cb1f7ec Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 20:05:59 +0530 Subject: [PATCH 46/76] =?UTF-8?q?=F0=9F=98=85.=20Fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 2 +- src/peft/utils/save_and_load.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 0023717..63201cc 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -404,7 +404,7 @@ class LoraLayer: nn.init.zeros_(self.lora_B[adapter_name].weight) -class Linear(nn.Linear): +class Linear(nn.Linear, LoraLayer): # Lora implemented in a dense layer def __init__( self, diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 7ebdfaf..a258c32 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -53,7 +53,7 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): elif isinstance(config, PromptLearningConfig): to_return = {} if config.inference_mode: - prompt_embeddings = model.prompt_encoder.embedding.weight + prompt_embeddings = model.prompt_encoder[adapter_name].embedding.weight else: prompt_embeddings = model.get_prompt_embedding_to_save(adapter_name) to_return["prompt_embeddings"] = prompt_embeddings From d4c2bc60e4bcbe47d0ba25cae23783b1503c5b49 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:04:42 +0530 Subject: [PATCH 47/76] =?UTF-8?q?fix=20more=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 63201cc..13ae8f5 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -135,22 +135,22 @@ class LoraModel(torch.nn.Module): super().__init__() self.model = model self.forward = self.model.forward - self.config = config - self.add_adapter(adapter_name, self.config[adapter_name]) + self.peft_config = config + self.add_adapter(adapter_name, self.peft_config[adapter_name]) def add_adapter(self, adapter_name, config=None): if config is not None: config = self._prepare_lora_config(config, self.model.config.to_dict()) - self.config[adapter_name] = config + self.peft_config[adapter_name] = config self._find_and_replace(adapter_name) - if len(self.config) > 1 and self.config[adapter_name].bias != "none": + if len(self.peft_config) > 1 and self.peft_config[adapter_name].bias != "none": 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.config[adapter_name].bias) + mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias) def _find_and_replace(self, adapter_name): - lora_config = self.config[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( @@ -260,7 +260,7 @@ class LoraModel(torch.nn.Module): def get_peft_config_as_dict(self, inference: bool = False): config_dict = {} - for key, value in self.config.items(): + for key, value in self.peft_config.items(): config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} if inference: config["inference_mode"] = True From 41b2fd770f50aa1c632068ec865a34ab2379dc34 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:13:41 +0530 Subject: [PATCH 48/76] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 13ae8f5..35f355e 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -477,7 +477,7 @@ class Linear(nn.Linear, LoraLayer): if self.disable_adapters: if self.r[self.active_adapter] > 0 and self.merged: self.unmerge() - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + 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 += ( @@ -487,7 +487,8 @@ class Linear(nn.Linear, LoraLayer): * self.scaling[self.active_adapter] ) else: - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + return result if is_bnb_available(): From dbdb8f375708862d9e286a2962613025eb24fbae Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:19:19 +0530 Subject: [PATCH 49/76] =?UTF-8?q?fix=20more=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 35f355e..53cbba7 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -319,7 +319,7 @@ class LoraModel(torch.nn.Module): key_list = [key for key, _ in self.model.named_modules() if "lora" not in key] for key in key_list: - parent, target, target_name = self._get_submodules(key) + parent, target, target_name = _get_submodules(key) if isinstance(target, LoraLayer): bias = target.bias is not None new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias) From 6f1f26f426861e402ddcf67ff23a26b66f62e877 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:30:07 +0530 Subject: [PATCH 50/76] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 53cbba7..74ed008 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -319,7 +319,7 @@ class LoraModel(torch.nn.Module): key_list = [key for key, _ in self.model.named_modules() if "lora" not in key] for key in key_list: - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(self.model, key) if isinstance(target, LoraLayer): bias = target.bias is not None new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias) From b9433a82082f7bb56504d5e0d04558f69d5e929f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:39:51 +0530 Subject: [PATCH 51/76] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/utils/other.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 7cd285f..53e2ee8 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -139,7 +139,7 @@ def _set_trainable(model, adapter_name): for key in key_list: target_module_found = any(key.endswith(target_key) for target_key in model.modules_to_save) if target_module_found: - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(model, key) if isinstance(target, ModulesToSaveWrapper): target.update(adapter_name) else: From 75131959d1d5b9c00acd7a59aefcc196b54d5ef9 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 02:03:31 +0530 Subject: [PATCH 52/76] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 1c3e7e1..fbe98e4 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -352,7 +352,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if ( (getattr(self, "hf_device_map", None) is not None) and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0) - and len(self.peft_config == 1) + and len(self.peft_config) == 1 ): device_map = kwargs.get("device_map", "auto") max_memory = kwargs.get("max_memory", None) From 44f3e86b6285901f4adea14c88642f0acbc1c70c Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 03:35:08 +0530 Subject: [PATCH 53/76] Update config.py --- src/peft/utils/config.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index 927daf0..34e98e9 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -29,7 +29,6 @@ class PeftType(str, enum.Enum): P_TUNING = "P_TUNING" PREFIX_TUNING = "PREFIX_TUNING" LORA = "LORA" - MULTI_LORA = "MULTI_LORA" class TaskType(str, enum.Enum): From 405f68f54abae0ffb281c16b1abbdb86bd79001f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 12:40:02 +0530 Subject: [PATCH 54/76] fix doc failure --- docs/source/package_reference/tuners.mdx | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/source/package_reference/tuners.mdx b/docs/source/package_reference/tuners.mdx index 2ec0824..9404266 100644 --- a/docs/source/package_reference/tuners.mdx +++ b/docs/source/package_reference/tuners.mdx @@ -14,8 +14,6 @@ For finetuning a model with LoRA. [[autodoc]] tuners.lora.Linear -[[autodoc]] tuners.lora.MergedLinear - ## P-tuning [[autodoc]] tuners.p_tuning.PromptEncoderConfig From deff03f2c251534fffd2511fc2d440e84cc54b1b Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Wed, 5 Apr 2023 10:17:26 +0200 Subject: [PATCH 55/76] [`tests`] Adds more tests + fix failing tests (#238) * adds more tests - refactor tests - add enc-dec tests - skips generate tests for non-lora adapters * rm unneeded file * fix tests * fix * more checks * fix issue --- src/peft/tuners/lora.py | 25 ++- tests/test_decoder_models.py | 85 ++++++++++ tests/test_encoder_decoder_models.py | 88 ++++++++++ tests/test_peft_model.py | 238 --------------------------- tests/testing_common.py | 194 +++++++++++++++++++++- 5 files changed, 381 insertions(+), 249 deletions(-) create mode 100644 tests/test_decoder_models.py create mode 100644 tests/test_encoder_decoder_models.py delete mode 100644 tests/test_peft_model.py diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 51cd56f..ab5259d 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -270,11 +270,26 @@ class LoraModel(torch.nn.Module): # manually merge if not merged if not target.merged: # merge weights per: https://arxiv.org/pdf/2106.09685.pdf / page 4 - if target.r > 0: - target.weight.data += ( - transpose(target.lora_B.weight @ target.lora_A.weight, target.fan_in_fan_out) - * target.scaling - ).to(target.weight.dtype) + if isinstance(target, Linear): + if target.r > 0: + target.weight.data += ( + transpose(target.lora_B.weight @ target.lora_A.weight, target.fan_in_fan_out) + * target.scaling + ).to(target.weight.dtype) + else: + if target.r > 0: + delta_w = ( + F.conv1d( + target.lora_A.weight.data.unsqueeze(0), + target.lora_B.weight.data, + groups=sum(target.enable_lora), + ) + .squeeze(0) + .transpose(-2, -1) + ) + target.weight.data += transpose( + target.zero_pad(delta_w * target.scaling), not target.fan_in_fan_out + ).to(target.weight.dtype) target.merged = True self._replace_module(parent, target_name, new_module, target) diff --git a/tests/test_decoder_models.py b/tests/test_decoder_models.py new file mode 100644 index 0000000..f0a575e --- /dev/null +++ b/tests/test_decoder_models.py @@ -0,0 +1,85 @@ +# coding=utf-8 +# Copyright 2023-present the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import unittest + +import torch +from parameterized import parameterized +from transformers import AutoModelForCausalLM + +from .testing_common import PeftCommonTester, PeftTestConfigManager + + +PEFT_DECODER_MODELS_TO_TEST = [ + "hf-internal-testing/tiny-random-OPTForCausalLM", + "hf-internal-testing/tiny-random-GPTNeoXForCausalLM", + "hf-internal-testing/tiny-random-GPT2LMHeadModel", + "hf-internal-testing/tiny-random-BloomForCausalLM", + "hf-internal-testing/tiny-random-gpt_neo", + "hf-internal-testing/tiny-random-GPTJForCausalLM", +] + +FULL_GRID = { + "model_ids": PEFT_DECODER_MODELS_TO_TEST, + "task_type": "CAUSAL_LM", +} + + +class PeftDecoderModelTester(unittest.TestCase, PeftCommonTester): + r""" + Test if the PeftModel behaves as expected. This includes: + - test if the model has the expected methods + + We use parametrized.expand for debugging purposes to test each model individually. + """ + transformers_class = AutoModelForCausalLM + + def prepare_inputs_for_testing(self): + input_ids = torch.tensor([[1, 1, 1], [1, 2, 1]]).to(self.torch_device) + attention_mask = torch.tensor([[1, 1, 1], [1, 0, 1]]).to(self.torch_device) + + input_dict = { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + return input_dict + + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) + def test_attributes_parametrized(self, test_name, model_id, config_cls, config_kwargs): + self._test_model_attr(model_id, config_cls, config_kwargs) + + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) + def test_prepare_for_training_parametrized(self, test_name, model_id, config_cls, config_kwargs): + self._test_prepare_for_training(model_id, config_cls, config_kwargs) + + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) + def test_save_pretrained(self, test_name, model_id, config_cls, config_kwargs): + self._test_save_pretrained(model_id, config_cls, config_kwargs) + + @parameterized.expand( + PeftTestConfigManager.get_grid_parameters( + { + "model_ids": PEFT_DECODER_MODELS_TO_TEST, + "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]}, + "task_type": "CAUSAL_LM", + }, + ) + ) + def test_merge_layers(self, test_name, model_id, config_cls, config_kwargs): + self._test_merge_layers(model_id, config_cls, config_kwargs) + + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) + def test_generate(self, test_name, model_id, config_cls, config_kwargs): + self._test_generate(model_id, config_cls, config_kwargs) diff --git a/tests/test_encoder_decoder_models.py b/tests/test_encoder_decoder_models.py new file mode 100644 index 0000000..6cb5a36 --- /dev/null +++ b/tests/test_encoder_decoder_models.py @@ -0,0 +1,88 @@ +# coding=utf-8 +# Copyright 2023-present the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import unittest + +import torch +from parameterized import parameterized +from transformers import AutoModelForSeq2SeqLM + +from .testing_common import PeftCommonTester, PeftTestConfigManager + + +PEFT_ENCODER_DECODER_MODELS_TO_TEST = [ + "hf-internal-testing/tiny-random-T5ForConditionalGeneration", + "hf-internal-testing/tiny-random-BartForConditionalGeneration", +] + +FULL_GRID = {"model_ids": PEFT_ENCODER_DECODER_MODELS_TO_TEST, "task_type": "SEQ_2_SEQ_LM"} + + +def skip_non_lora_or_pt(test_list): + r""" + Skip tests that are not lora or prefix tuning + """ + return [test for test in test_list if ("lora" in test[0] or "prefix_tuning" in test[0])] + + +class PeftEncoderDecoderModelTester(unittest.TestCase, PeftCommonTester): + r""" + Test if the PeftModel behaves as expected. This includes: + - test if the model has the expected methods + + We use parametrized.expand for debugging purposes to test each model individually. + """ + transformers_class = AutoModelForSeq2SeqLM + + def prepare_inputs_for_testing(self): + input_ids = torch.tensor([[1, 1, 1], [1, 2, 1]]).to(self.torch_device) + decoder_input_ids = torch.tensor([[1, 1, 1], [1, 2, 1]]).to(self.torch_device) + attention_mask = torch.tensor([[1, 1, 1], [1, 0, 1]]).to(self.torch_device) + + input_dict = { + "input_ids": input_ids, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + } + + return input_dict + + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) + def test_attributes_parametrized(self, test_name, model_id, config_cls, config_kwargs): + self._test_model_attr(model_id, config_cls, config_kwargs) + + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) + def test_prepare_for_training_parametrized(self, test_name, model_id, config_cls, config_kwargs): + self._test_prepare_for_training(model_id, config_cls, config_kwargs) + + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) + def test_save_pretrained(self, test_name, model_id, config_cls, config_kwargs): + self._test_save_pretrained(model_id, config_cls, config_kwargs) + + @parameterized.expand( + PeftTestConfigManager.get_grid_parameters( + { + "model_ids": PEFT_ENCODER_DECODER_MODELS_TO_TEST, + "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [True, False]}, + "task_type": "SEQ_2_SEQ_LM", + }, + ) + ) + def test_merge_layers(self, test_name, model_id, config_cls, config_kwargs): + self._test_merge_layers(model_id, config_cls, config_kwargs) + + # skip non lora models - generate does not work for prefix tuning, prompt tuning + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID, filter_params_func=skip_non_lora_or_pt)) + def test_generate(self, test_name, model_id, config_cls, config_kwargs): + self._test_generate(model_id, config_cls, config_kwargs) diff --git a/tests/test_peft_model.py b/tests/test_peft_model.py deleted file mode 100644 index 4280ff3..0000000 --- a/tests/test_peft_model.py +++ /dev/null @@ -1,238 +0,0 @@ -# coding=utf-8 -# Copyright 2023-present the HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import os -import tempfile -import unittest - -import torch -from parameterized import parameterized -from transformers import AutoModelForCausalLM - -from peft import ( - PeftModel, - get_peft_model, - get_peft_model_state_dict, - prepare_model_for_int8_training, -) - -from .testing_common import PeftTestConfigManager - - -PEFT_DECODER_MODELS_TO_TEST = [ - "hf-internal-testing/tiny-random-OPTForCausalLM", - "hf-internal-testing/tiny-random-GPTNeoXForCausalLM", - "hf-internal-testing/tiny-random-GPT2LMHeadModel", - "hf-internal-testing/tiny-random-BloomForCausalLM", - "hf-internal-testing/tiny-random-gpt_neo", - "hf-internal-testing/tiny-random-GPTJForCausalLM", -] - -FULL_GRID = { - "model_ids": PEFT_DECODER_MODELS_TO_TEST, -} - - -class PeftTestMixin: - torch_device = "cuda" if torch.cuda.is_available() else "cpu" - - -class PeftModelTester(unittest.TestCase, PeftTestMixin): - r""" - Test if the PeftModel behaves as expected. This includes: - - test if the model has the expected methods - - We use parametrized.expand for debugging purposes to test each model individually. - """ - - def _test_model_attr(self, model_id, config_cls, config_kwargs): - model = AutoModelForCausalLM.from_pretrained(model_id) - config = config_cls( - base_model_name_or_path=model_id, - **config_kwargs, - ) - model = get_peft_model(model, config) - - self.assertTrue(hasattr(model, "save_pretrained")) - self.assertTrue(hasattr(model, "from_pretrained")) - self.assertTrue(hasattr(model, "push_to_hub")) - - @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) - def test_attributes_parametrized(self, test_name, model_id, config_cls, config_kwargs): - self._test_model_attr(model_id, config_cls, config_kwargs) - - def _test_prepare_for_training(self, model_id, config_cls, config_kwargs): - model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device) - config = config_cls( - base_model_name_or_path=model_id, - **config_kwargs, - ) - model = get_peft_model(model, config) - - dummy_input = torch.LongTensor([[1, 1, 1]]).to(self.torch_device) - dummy_output = model.get_input_embeddings()(dummy_input) - - self.assertTrue(not dummy_output.requires_grad) - - # load with `prepare_model_for_int8_training` - model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device) - model = prepare_model_for_int8_training(model) - - for param in model.parameters(): - self.assertTrue(not param.requires_grad) - - config = config_cls( - base_model_name_or_path=model_id, - **config_kwargs, - ) - model = get_peft_model(model, config) - - # For backward compatibility - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) - - dummy_input = torch.LongTensor([[1, 1, 1]]).to(self.torch_device) - dummy_output = model.get_input_embeddings()(dummy_input) - - self.assertTrue(dummy_output.requires_grad) - - @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) - def test_prepare_for_training_parametrized(self, test_name, model_id, config_cls, config_kwargs): - self._test_prepare_for_training(model_id, config_cls, config_kwargs) - - def _test_save_pretrained(self, model_id, config_cls, config_kwargs): - model = AutoModelForCausalLM.from_pretrained(model_id) - config = config_cls( - base_model_name_or_path=model_id, - **config_kwargs, - ) - model = get_peft_model(model, config) - model = model.to(self.torch_device) - - with tempfile.TemporaryDirectory() as tmp_dirname: - model.save_pretrained(tmp_dirname) - - model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id) - model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname) - - # check if the state dicts are equal - state_dict = get_peft_model_state_dict(model) - state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained) - - # check if same keys - self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys()) - - # check if tensors equal - for key in state_dict.keys(): - self.assertTrue( - torch.allclose( - state_dict[key].to(self.torch_device), state_dict_from_pretrained[key].to(self.torch_device) - ) - ) - - # check if `adapter_model.bin` is present - self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin"))) - - # check if `adapter_config.json` is present - self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json"))) - - # check if `pytorch_model.bin` is not present - self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin"))) - - # check if `config.json` is not present - self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json"))) - - @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) - def test_save_pretrained(self, test_name, model_id, config_cls, config_kwargs): - self._test_save_pretrained(model_id, config_cls, config_kwargs) - - def _test_merge_layers(self, model_id, config_cls, config_kwargs): - model = AutoModelForCausalLM.from_pretrained(model_id) - config = config_cls( - base_model_name_or_path=model_id, - **config_kwargs, - ) - model = get_peft_model(model, config) - model = model.to(self.torch_device) - - if config.peft_type != "LORA": - with self.assertRaises(AttributeError): - model = model.merge_and_unload() - elif model.config.model_type == "gpt2": - with self.assertRaises(ValueError): - model = model.merge_and_unload() - else: - dummy_input = torch.LongTensor([[1, 2, 3, 2, 1]]).to(self.torch_device) - model.eval() - logits_lora = model(dummy_input)[0] - - model = model.merge_and_unload() - - logits_merged = model(dummy_input)[0] - - transformers_model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device) - - logits_transformers = transformers_model(dummy_input)[0] - - self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-3, rtol=1e-3)) - self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-3, rtol=1e-3)) - - with tempfile.TemporaryDirectory() as tmp_dirname: - model.save_pretrained(tmp_dirname) - - model_from_pretrained = AutoModelForCausalLM.from_pretrained(tmp_dirname).to(self.torch_device) - - logits_merged_from_pretrained = model_from_pretrained(dummy_input)[0] - - self.assertTrue(torch.allclose(logits_merged, logits_merged_from_pretrained, atol=1e-3, rtol=1e-3)) - - @parameterized.expand( - PeftTestConfigManager.get_grid_parameters( - { - "model_ids": PEFT_DECODER_MODELS_TO_TEST, - "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]}, - }, - ) - ) - def test_merge_layers(self, test_name, model_id, config_cls, config_kwargs): - self._test_merge_layers(model_id, config_cls, config_kwargs) - - def _test_generate(self, model_id, config_cls, config_kwargs): - model = AutoModelForCausalLM.from_pretrained(model_id) - config = config_cls( - base_model_name_or_path=model_id, - **config_kwargs, - ) - model = get_peft_model(model, config) - model = model.to(self.torch_device) - - input_ids = torch.LongTensor([[1, 1, 1], [2, 1, 2]]).to(self.torch_device) - attention_mask = torch.LongTensor([[1, 1, 1], [1, 0, 1]]).to(self.torch_device) - - # check if `generate` works - _ = model.generate(input_ids=input_ids, attention_mask=attention_mask) - - with self.assertRaises(TypeError): - # check if `generate` raises an error if no positional arguments are passed - _ = model.generate(input_ids, attention_mask=attention_mask) - - @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) - def test_generate(self, test_name, model_id, config_cls, config_kwargs): - self._test_generate(model_id, config_cls, config_kwargs) diff --git a/tests/testing_common.py b/tests/testing_common.py index 633bb87..cfe6cf2 100644 --- a/tests/testing_common.py +++ b/tests/testing_common.py @@ -12,13 +12,21 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import os +import tempfile from collections import OrderedDict +import torch + from peft import ( LoraConfig, + PeftModel, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig, + get_peft_model, + get_peft_model_state_dict, + prepare_model_for_int8_training, ) @@ -35,20 +43,16 @@ CONFIG_TESTING_KWARGS = ( "target_modules": None, "lora_dropout": 0.05, "bias": "none", - "task_type": "CAUSAL_LM", }, { "num_virtual_tokens": 10, - "task_type": "CAUSAL_LM", }, { "num_virtual_tokens": 10, "encoder_hidden_size": 32, - "task_type": "CAUSAL_LM", }, { "num_virtual_tokens": 10, - "task_type": "CAUSAL_LM", }, ) @@ -92,6 +96,7 @@ class ClassInstantier(OrderedDict): """ generated_tests = [] model_list = grid_parameters["model_ids"] + task_type = grid_parameters["task_type"] if "task_type" in grid_parameters else None for model_id in model_list: for key, value in self.items(): @@ -101,9 +106,16 @@ class ClassInstantier(OrderedDict): for current_key, current_value in grid_parameters[f"{key}_kwargs"].items(): for kwarg in current_value: current_peft_config.update({current_key: kwarg}) - peft_configs.append(current_peft_config) + + if task_type is not None: + current_peft_config.update({"task_type": task_type}) + + peft_configs.append(current_peft_config.copy()) else: - peft_configs = [value[1].copy()] + current_peft_config = value[1].copy() + if task_type is not None: + current_peft_config.update({"task_type": task_type}) + peft_configs = [current_peft_config] for peft_config in peft_configs: generated_tests.append((f"test_{model_id}_{key}", model_id, value[0], peft_config)) @@ -115,3 +127,173 @@ class ClassInstantier(OrderedDict): PeftTestConfigManager = ClassInstantier(CLASSES_MAPPING) + + +class PeftCommonTester: + r""" + A large testing suite for testing common functionality of the PEFT models. + + Attributes: + torch_device (`torch.device`): + The device on which the tests will be run. + transformers_class (`transformers.PreTrainedModel`): + The transformers class that is being tested. + """ + torch_device = "cuda" if torch.cuda.is_available() else "cpu" + transformers_class = None + + def prepare_inputs_for_common(self): + raise NotImplementedError + + def _test_model_attr(self, model_id, config_cls, config_kwargs): + model = self.transformers_class.from_pretrained(model_id) + config = config_cls( + base_model_name_or_path=model_id, + **config_kwargs, + ) + model = get_peft_model(model, config) + + self.assertTrue(hasattr(model, "save_pretrained")) + self.assertTrue(hasattr(model, "from_pretrained")) + self.assertTrue(hasattr(model, "push_to_hub")) + + def _test_prepare_for_training(self, model_id, config_cls, config_kwargs): + model = self.transformers_class.from_pretrained(model_id).to(self.torch_device) + config = config_cls( + base_model_name_or_path=model_id, + **config_kwargs, + ) + model = get_peft_model(model, config) + + dummy_input = self.prepare_inputs_for_testing() + dummy_output = model.get_input_embeddings()(dummy_input["input_ids"]) + + self.assertTrue(not dummy_output.requires_grad) + + # load with `prepare_model_for_int8_training` + model = self.transformers_class.from_pretrained(model_id).to(self.torch_device) + model = prepare_model_for_int8_training(model) + + for param in model.parameters(): + self.assertTrue(not param.requires_grad) + + config = config_cls( + base_model_name_or_path=model_id, + **config_kwargs, + ) + model = get_peft_model(model, config) + + # For backward compatibility + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + dummy_input = self.prepare_inputs_for_testing() + dummy_output = model.get_input_embeddings()(dummy_input["input_ids"]) + + self.assertTrue(dummy_output.requires_grad) + + def _test_save_pretrained(self, model_id, config_cls, config_kwargs): + model = self.transformers_class.from_pretrained(model_id) + config = config_cls( + base_model_name_or_path=model_id, + **config_kwargs, + ) + model = get_peft_model(model, config) + model = model.to(self.torch_device) + + with tempfile.TemporaryDirectory() as tmp_dirname: + model.save_pretrained(tmp_dirname) + + model_from_pretrained = self.transformers_class.from_pretrained(model_id) + model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname) + + # check if the state dicts are equal + state_dict = get_peft_model_state_dict(model) + state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained) + + # check if same keys + self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys()) + + # check if tensors equal + for key in state_dict.keys(): + self.assertTrue( + torch.allclose( + state_dict[key].to(self.torch_device), state_dict_from_pretrained[key].to(self.torch_device) + ) + ) + + # check if `adapter_model.bin` is present + self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin"))) + + # check if `adapter_config.json` is present + self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json"))) + + # check if `pytorch_model.bin` is not present + self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin"))) + + # check if `config.json` is not present + self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json"))) + + def _test_merge_layers(self, model_id, config_cls, config_kwargs): + model = self.transformers_class.from_pretrained(model_id) + config = config_cls( + base_model_name_or_path=model_id, + **config_kwargs, + ) + model = get_peft_model(model, config) + model = model.to(self.torch_device) + + if config.peft_type != "LORA": + with self.assertRaises(AttributeError): + model = model.merge_and_unload() + elif model.config.model_type == "gpt2": + with self.assertRaises(ValueError): + model = model.merge_and_unload() + else: + dummy_input = self.prepare_inputs_for_testing() + model.eval() + logits_lora = model(**dummy_input)[0] + + model = model.merge_and_unload() + + logits_merged = model(**dummy_input)[0] + + transformers_model = self.transformers_class.from_pretrained(model_id).to(self.torch_device) + + logits_transformers = transformers_model(**dummy_input)[0] + + 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) + + model_from_pretrained = self.transformers_class.from_pretrained(tmp_dirname).to(self.torch_device) + + logits_merged_from_pretrained = model_from_pretrained(**dummy_input)[0] + + self.assertTrue(torch.allclose(logits_merged, logits_merged_from_pretrained, atol=1e-4, rtol=1e-4)) + + def _test_generate(self, model_id, config_cls, config_kwargs): + model = self.transformers_class.from_pretrained(model_id) + config = config_cls( + base_model_name_or_path=model_id, + **config_kwargs, + ) + model = get_peft_model(model, config) + model = model.to(self.torch_device) + + inputs = self.prepare_inputs_for_testing() + + # check if `generate` works + _ = model.generate(**inputs) + + with self.assertRaises(TypeError): + # check if `generate` raises an error if no positional arguments are passed + _ = model.generate(inputs["input_ids"]) From d936aa9349a25e7acc2e954070652ac6297babdf Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 19:54:02 +0530 Subject: [PATCH 56/76] fix tests --- tests/testing_common.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/testing_common.py b/tests/testing_common.py index cfe6cf2..38aee14 100644 --- a/tests/testing_common.py +++ b/tests/testing_common.py @@ -268,8 +268,12 @@ class PeftCommonTester: logits_transformers = transformers_model(**dummy_input)[0] - 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)) + 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)) with tempfile.TemporaryDirectory() as tmp_dirname: model.save_pretrained(tmp_dirname) From 37e1f9ba340b2025ff2f415827abdc613e1acbf2 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Wed, 5 Apr 2023 17:17:01 +0000 Subject: [PATCH 57/76] fix test --- tests/test_encoder_decoder_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_encoder_decoder_models.py b/tests/test_encoder_decoder_models.py index 6cb5a36..c7c73e6 100644 --- a/tests/test_encoder_decoder_models.py +++ b/tests/test_encoder_decoder_models.py @@ -22,7 +22,7 @@ from .testing_common import PeftCommonTester, PeftTestConfigManager PEFT_ENCODER_DECODER_MODELS_TO_TEST = [ - "hf-internal-testing/tiny-random-T5ForConditionalGeneration", + "ybelkada/tiny-random-T5ForConditionalGeneration-calibrated", "hf-internal-testing/tiny-random-BartForConditionalGeneration", ] @@ -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": [True, False]}, + "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]}, "task_type": "SEQ_2_SEQ_LM", }, ) From 3e6a88a8f9f063f93010ef94d8646d93e93bc056 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 5 Apr 2023 16:23:48 -0400 Subject: [PATCH 58/76] Update src/peft/tuners/adalora.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> --- src/peft/tuners/adalora.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 5e032a6..76a6b0c 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -139,20 +139,6 @@ class AdaLoraModel(LoraModel): new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None: new_module = SVDLinear(target.in_features, target.out_features, bias=bias, **kwargs) - # TODO: Implement the MergedLinear of SVD Adapattion - # elif self.peft_config.enable_lora is not None: - # kwargs.update({"enable_lora": self.peft_config.enable_lora}) - # if isinstance(target, Conv1D): - # in_features, out_features = target.weight.shape - # else: - # 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 not a Conv1D. " - # "Setting fan_in_fan_out to False." - # ) - # kwargs["fan_in_fan_out"] = False - # new_module = MergedLinear(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( From b3e6ef6224084005832616e5d8507ced885c7b41 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 5 Apr 2023 16:30:16 -0400 Subject: [PATCH 59/76] Update src/peft/tuners/adalora.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 76a6b0c..c2eb684 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -127,7 +127,7 @@ class AdaLoraModel(LoraModel): 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) and self.peft_config.enable_lora is None: + if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): kwargs.update( { "has_fp16_weights": target.state.has_fp16_weights, From c240a9693cdc9b1963fffc120c3693c9fb063960 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 5 Apr 2023 16:31:08 -0400 Subject: [PATCH 60/76] Update src/peft/tuners/adalora.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index c2eb684..6b19324 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -137,7 +137,7 @@ class AdaLoraModel(LoraModel): } ) new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) - elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None: + 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: From 9a534d047cdbf42af078bc5d0ebe23e8c39448f5 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 5 Apr 2023 16:31:46 -0400 Subject: [PATCH 61/76] Update src/peft/tuners/adalora.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 6b19324..0f43418 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -192,7 +192,7 @@ class AdaLoraModel(LoraModel): } 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) and self.peft_config.enable_lora is None: + if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): kwargs.update( { "has_fp16_weights": target.state.has_fp16_weights, From d892beb0e742a3d2bcc6e2d6e4f20b595b6b5d65 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 5 Apr 2023 16:31:58 -0400 Subject: [PATCH 62/76] Update src/peft/mapping.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> --- src/peft/mapping.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/peft/mapping.py b/src/peft/mapping.py index d9ca4fa..7d10f0f 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -150,9 +150,6 @@ def _prepare_adalora_config(peft_config, model_config): 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 len(peft_config.target_modules) == 1: - peft_config.fan_in_fan_out = True - # peft_config.enable_lora = [True, False, True] if peft_config.inference_mode: peft_config.merge_weights = True return peft_config From b8a57a3649ec6fa32ce3c71622377dcc5b44ea41 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 5 Apr 2023 16:32:23 -0400 Subject: [PATCH 63/76] Update src/peft/tuners/adalora.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 0f43418..d2f04d1 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -202,7 +202,7 @@ class AdaLoraModel(LoraModel): } ) new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) - elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None: + 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) From 4f8c134102ac704a220a0f81d8d86662b14958c2 Mon Sep 17 00:00:00 2001 From: zqingru Date: Wed, 5 Apr 2023 20:40:06 +0000 Subject: [PATCH 64/76] raise exception for MergedLinear of AdaLoRA --- src/peft/tuners/adalora.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index d2f04d1..1e7be53 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -100,6 +100,11 @@ class AdaLoraModel(LoraModel): 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) From 072da6d9d625bbafa21d0a9109e359ecd9543142 Mon Sep 17 00:00:00 2001 From: Qingru Zhang Date: Wed, 5 Apr 2023 20:52:12 +0000 Subject: [PATCH 65/76] Run make style and make quality --- .../peft_adalora_seq2seq.py | 47 +-- src/peft/__init__.py | 4 +- src/peft/mapping.py | 5 +- src/peft/peft_model.py | 4 +- src/peft/tuners/adalora.py | 331 ++++++++---------- src/peft/utils/save_and_load.py | 4 +- 6 files changed, 185 insertions(+), 210 deletions(-) diff --git a/examples/conditional_generation/peft_adalora_seq2seq.py b/examples/conditional_generation/peft_adalora_seq2seq.py index 391be4a..31b4aa5 100644 --- a/examples/conditional_generation/peft_adalora_seq2seq.py +++ b/examples/conditional_generation/peft_adalora_seq2seq.py @@ -1,15 +1,15 @@ -from transformers import AutoModelForSeq2SeqLM -from peft import get_peft_model, AdaLoraConfig, AdaLoraModel, TaskType -import torch -from datasets import load_dataset import os -os.environ["TOKENIZERS_PARALLELISM"] = "false" -from transformers import AutoTokenizer -from torch.utils.data import DataLoader -from transformers import default_data_collator, get_linear_schedule_with_warmup -from tqdm import tqdm +import torch from datasets import load_dataset +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, default_data_collator, get_linear_schedule_with_warmup + +from peft import AdaLoraConfig, PeftConfig, PeftModel, TaskType, get_peft_model + + +os.environ["TOKENIZERS_PARALLELISM"] = "false" device = "cuda" model_name_or_path = "facebook/bart-base" @@ -20,18 +20,23 @@ text_column = "sentence" label_column = "text_label" max_length = 128 lr = 1e-3 -num_epochs = 8 +num_epochs = 8 batch_size = 8 # creating model peft_config = AdaLoraConfig( - init_r=12, target_r=8, - beta1=0.85, beta2=0.85, - tinit=200, tfinal=1000, deltaT=10, - lora_alpha=32, lora_dropout=0.1, - task_type=TaskType.SEQ_2_SEQ_LM, - inference_mode=False + init_r=12, + target_r=8, + beta1=0.85, + beta2=0.85, + tinit=200, + tfinal=1000, + deltaT=10, + lora_alpha=32, + lora_dropout=0.1, + task_type=TaskType.SEQ_2_SEQ_LM, + inference_mode=False, ) model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) @@ -98,7 +103,7 @@ model.base_model.peft_config.total_step = len(train_dataloader) * num_epochs # training and evaluation model = model.to(device) -global_step = 0 +global_step = 0 for epoch in range(num_epochs): model.train() total_loss = 0 @@ -110,8 +115,8 @@ for epoch in range(num_epochs): loss.backward() optimizer.step() lr_scheduler.step() - # Update the importance of low-rank matrices - # and allocate the budget accordingly. + # Update the importance of low-rank matrices + # and allocate the budget accordingly. model.base_model.update_and_allocate(global_step) optimizer.zero_grad() global_step += 1 @@ -158,8 +163,6 @@ ckpt = f"{peft_model_id}/adapter_model.bin" # get_ipython().system('du -h $ckpt') -from peft import PeftModel, PeftConfig - peft_model_id = f"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}" config = PeftConfig.from_pretrained(peft_model_id) @@ -177,5 +180,3 @@ with torch.no_grad(): outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=10) print(outputs) print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)) - - diff --git a/src/peft/__init__.py b/src/peft/__init__.py index 3a68691..dc666ce 100644 --- a/src/peft/__init__.py +++ b/src/peft/__init__.py @@ -30,8 +30,8 @@ from .peft_model import ( from .tuners import ( LoraConfig, LoraModel, - AdaLoraConfig, - AdaLoraModel, + AdaLoraConfig, + AdaLoraModel, PrefixEncoder, PrefixTuningConfig, PromptEmbedding, diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 7d10f0f..260335c 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -20,7 +20,7 @@ from .peft_model import ( PeftModelForSequenceClassification, PeftModelForTokenClassification, ) -from .tuners import LoraConfig, AdaLoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig +from .tuners import AdaLoraConfig, LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from .utils import PromptLearningConfig @@ -36,7 +36,7 @@ PEFT_TYPE_TO_CONFIG_MAPPING = { "PREFIX_TUNING": PrefixTuningConfig, "P_TUNING": PromptEncoderConfig, "LORA": LoraConfig, - "ADALORA": AdaLoraConfig, + "ADALORA": AdaLoraConfig, } TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { @@ -145,6 +145,7 @@ def _prepare_lora_config(peft_config, model_config): peft_config.merge_weights = True return peft_config + def _prepare_adalora_config(peft_config, model_config): if peft_config.target_modules is None: if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING: diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 2d80efc..45a5b75 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -28,9 +28,7 @@ from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput from transformers.utils import PushToHubMixin -from huggingface_hub import hf_hub_download - -from .tuners import LoraModel, AdaLoraConfig, AdaLoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder +from .tuners import AdaLoraConfig, AdaLoraModel, LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder from .utils import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, WEIGHTS_NAME, diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 1e7be53..f98b3c6 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -1,19 +1,14 @@ import importlib -import math import re -import warnings -import numpy as np -from dataclasses import asdict, dataclass, field -from enum import Enum -from typing import List, Optional, Union +from dataclasses import dataclass, field +from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F -from transformers.pytorch_utils import Conv1D -from ..utils import PeftConfig, PeftType, transpose -from .lora import LoraConfig, LoraModel, LoraLayer, mark_only_lora_as_trainable +from ..utils import PeftType, transpose +from .lora import LoraConfig, LoraLayer, LoraModel, mark_only_lora_as_trainable def is_bnb_available(): @@ -30,46 +25,37 @@ class AdaLoraConfig(LoraConfig): This is the configuration class to store the configuration of a [`~peft.AdaLora`]. Args: - target_r (`int`): The target average rank of incremental matrix. - init_r (`int`): The initial rank for each incremental matrix. - tinit (`int`): The steps of initial fine-tuning warmup. - tfinal (`int`): The step of final fine-tuning. - deltaT (`int`): The time internval between two budget allocations. + target_r (`int`): The target average rank of incremental matrix. + init_r (`int`): The initial rank for each incremental matrix. + tinit (`int`): The steps of initial fine-tuning warmup. + tfinal (`int`): The step of final fine-tuning. + deltaT (`int`): The time internval between two budget allocations. beta1 (`float`): The hyperparameter of EMA for sensitivity smoothing. - beta2 (`float`): The hyperparameter of EMA for undertainty quantification. - orth_reg_weight (`float`): The coefficient of orthogonal regularization. - total_step (`int`): The total training steps that should be specified before training. - rank_pattern (`list`): The allocated rank for each weight matrix by RankAllocator. + beta2 (`float`): The hyperparameter of EMA for undertainty quantification. + orth_reg_weight (`float`): The coefficient of orthogonal regularization. + total_step (`int`): The total training steps that should be specified before training. + rank_pattern (`list`): The allocated rank for each weight matrix by RankAllocator. """ + target_r: int = field(default=8, metadata={"help": "Target Lora matrix dimension."}) init_r: int = field(default=12, metadata={"help": "Intial Lora matrix dimension."}) tinit: int = field(default=0, metadata={"help": "The steps of initial warmup."}) tfinal: int = field(default=0, metadata={"help": "The steps of final warmup."}) deltaT: int = field(default=1, metadata={"help": "Step interval of rank allocation."}) beta1: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) - beta2: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) - orth_reg_weight: float = field( - default=0.5, - metadata={"help": "The orthogonal regularization coefficient."} - ) - total_step: Optional[int] = field( - default=None, - metadata={"help": "The total training steps."} - ) - rank_pattern: Optional[dict] = field( - default=None, - metadata={"help":"The saved rank pattern."} - ) + beta2: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) + orth_reg_weight: float = field(default=0.5, metadata={"help": "The orthogonal regularization coefficient."}) + total_step: Optional[int] = field(default=None, metadata={"help": "The total training steps."}) + rank_pattern: Optional[dict] = field(default=None, metadata={"help": "The saved rank pattern."}) def __post_init__(self): self.peft_type = PeftType.ADALORA - class AdaLoraModel(LoraModel): """ - Creates AdaLoRA (Adaptive LoRA) model from a pretrained transformers model. - Paper: https://openreview.net/pdf?id=lq62uWRJjiY + Creates AdaLoRA (Adaptive LoRA) model from a pretrained transformers model. Paper: + https://openreview.net/pdf?id=lq62uWRJjiY Args: model ([`transformers.PreTrainedModel`]): The model to be adapted. @@ -80,13 +66,12 @@ class AdaLoraModel(LoraModel): Example:: - >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import AdaLoraModel, AdaLoraConfig + >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import AdaLoraModel, AdaLoraConfig >>> config = AdaLoraConfig( peft_type="ADALORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"], lora_dropout=0.01, ) - >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") - >>> model = AdaLoraModel(config, model) + >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> model = AdaLoraModel(config, model) **Attributes**: - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. @@ -101,10 +86,7 @@ class AdaLoraModel(LoraModel): mark_only_lora_as_trainable(self.model, self.peft_config.bias) self.rankallocator = RankAllocator(config, self.model) if config.enable_lora is not None: - raise NotImplementedError( - "MergedLinear has not been implemented for AdaLoRA." - ) - + raise NotImplementedError("MergedLinear has not been implemented for AdaLoRA.") def _find_and_replace(self): loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) @@ -159,26 +141,25 @@ class AdaLoraModel(LoraModel): return getattr(self.model, name) def forward(self, *args, **kwargs): - outputs = self.model.forward(*args, **kwargs) + outputs = self.model.forward(*args, **kwargs) - # Calculate the orthogonal regularization + # Calculate the orthogonal regularization orth_reg_weight = self.peft_config.orth_reg_weight - assert orth_reg_weight > 0 + assert orth_reg_weight > 0 if hasattr(outputs, "loss"): - regu_loss = 0 - num_param = 0 - for n,p in self.model.named_parameters(): + regu_loss = 0 + num_param = 0 + for n, p in self.model.named_parameters(): if "lora_A" in n or "lora_B" in n: - para_cov = p @ p.T if "lora_A" in n else p.T @ p + para_cov = p @ p.T if "lora_A" in n else p.T @ p I = torch.eye(*para_cov.size(), out=torch.empty_like(para_cov)) I.requires_grad = False num_param += 1 - regu_loss += torch.norm(para_cov-I, p="fro") - regu_loss = regu_loss / num_param - outputs.loss += orth_reg_weight * regu_loss - return outputs - + regu_loss += torch.norm(para_cov - I, p="fro") + regu_loss = regu_loss / num_param + outputs.loss += orth_reg_weight * regu_loss + return outputs def _prepare_new_module(self, target, rank_idx): if isinstance(rank_idx, list): @@ -187,7 +168,7 @@ class AdaLoraModel(LoraModel): rank_idx = rank_idx.view(-1) rank = rank_idx.sum().item() else: - raise ValueError(f"Unexcepted type of rank_idx") + raise ValueError("Unexcepted type of rank_idx") kwargs = { "r": rank, "lora_alpha": self.peft_config.lora_alpha, @@ -218,68 +199,64 @@ class AdaLoraModel(LoraModel): if rank > 0: new_module.lora_E.copy_(target.lora_E[rank_idx]) new_module.lora_A.copy_(target.lora_A[rank_idx]) - new_module.lora_B.copy_(target.lora_B[:,rank_idx]) - # The scaling is exactly as the previous + new_module.lora_B.copy_(target.lora_B[:, rank_idx]) + # The scaling is exactly as the previous new_module.ranknum.copy_(target.ranknum) return new_module def resize_modules_by_rank_pattern(self, rank_pattern): - for name,rank_idx in rank_pattern.items(): - key = ".".join(name.split(".")[0:-1]) - parent, target, target_name = self._get_submodules(key) + for name, rank_idx in rank_pattern.items(): + key = ".".join(name.split(".")[0:-1]) + parent, target, target_name = self._get_submodules(key) new_module = self._prepare_new_module(target, rank_idx) self._replace_module(parent, target_name, new_module, target) def update_and_allocate(self, global_step): - # Update the importance score and allocate the budget + # Update the importance score and allocate the budget if global_step < self.peft_config.total_step - self.peft_config.tfinal: budget, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step) if rank_pattern: - self.peft_config.rank_pattern = rank_pattern - # Finalize the budget allocation - elif global_step == self.peft_config.total_step - self.peft_config.tfinal: - budget, rank_pattern = self.rankallocator.update_and_allocate( - self.model, global_step, force_mask=True - ) + self.peft_config.rank_pattern = rank_pattern + # Finalize the budget allocation + elif global_step == self.peft_config.total_step - self.peft_config.tfinal: + budget, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True) self.resize_modules_by_rank_pattern(rank_pattern) self.peft_config.rank_pattern = rank_pattern - self.rankallocator.reset_ipt() - # Pass the function and do forward propagation - else: + self.rankallocator.reset_ipt() + # Pass the function and do forward propagation + else: return None - class SVDLinear(nn.Linear, LoraLayer): # SVD-based adaptation by a dense layer def __init__( - self, - in_features: int, - out_features: int, - r: int = 0, - lora_alpha: int = 1, - lora_dropout: float = 0., - fan_in_fan_out: bool = False, + self, + in_features: int, + out_features: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + fan_in_fan_out: bool = False, merge_weights: bool = True, - **kwargs + **kwargs, ): nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, - merge_weights=merge_weights) - + LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + self.fan_in_fan_out = fan_in_fan_out # Actual trainable parameters if r > 0: # Right singular vectors self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) - # Singular values - self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) + # Singular values + self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) # Left singular vectors self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) # The current rank self.ranknum = nn.Parameter(self.weight.new_zeros(1), requires_grad=False) self.ranknum.data.fill_(float(self.r)) - self.scaling = self.lora_alpha if self.lora_alpha>0 else float(self.r) + self.scaling = self.lora_alpha if self.lora_alpha > 0 else float(self.r) # Freezing the pre-trained weight matrix self.weight.requires_grad = False self.ranknum.requires_grad = False @@ -289,7 +266,7 @@ class SVDLinear(nn.Linear, LoraLayer): def reset_parameters(self): nn.Linear.reset_parameters(self) - if hasattr(self, 'lora_A'): + if hasattr(self, "lora_A"): nn.init.zeros_(self.lora_E) nn.init.normal_(self.lora_A, mean=0.0, std=0.02) nn.init.normal_(self.lora_B, mean=0.0, std=0.02) @@ -299,19 +276,19 @@ class SVDLinear(nn.Linear, LoraLayer): if self.merge_weights and self.merged: # Make sure that the weights are not merged if self.r > 0: - self.weight.data -= transpose( - self.lora_B @ (self.lora_A * self.lora_E) - ) * self.scaling/(self.ranknum+1e-5) + self.weight.data -= ( + transpose(self.lora_B @ (self.lora_A * self.lora_E)) * self.scaling / (self.ranknum + 1e-5) + ) self.merged = False - + def eval(self): nn.Linear.eval(self) if self.merge_weights and not self.merged: # Merge the weights and mark it if self.r > 0: - self.weight.data += transpose( - self.lora_B @ (self.lora_A * self.lora_E) - ) * self.scaling/(self.ranknum+1e-5) + self.weight.data += ( + transpose(self.lora_B @ (self.lora_A * self.lora_E)) * self.scaling / (self.ranknum + 1e-5) + ) self.merged = True def forward(self, x: torch.Tensor): @@ -319,16 +296,19 @@ class SVDLinear(nn.Linear, LoraLayer): result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if self.r > 0: result += ( - self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T - ) * self.scaling / (self.ranknum+1e-5) + (self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T) + * self.scaling + / (self.ranknum + 1e-5) + ) return result else: return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if is_bnb_available(): + class SVDLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): - # Low-rank matrix for SVD-based adaptation + # Low-rank matrix for SVD-based adaptation def __init__( self, in_features, @@ -353,14 +333,14 @@ if is_bnb_available(): if r > 0: # Right singular vectors self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) - # Singular values - self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) + # Singular values + self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) # Left singular vectors self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) # The current rank self.ranknum = nn.Parameter(self.weight.new_zeros(1), requires_grad=False) self.ranknum.data.fill_(float(self.r)) - self.scaling = self.lora_alpha if self.lora_alpha>0 else float(self.r) + self.scaling = self.lora_alpha if self.lora_alpha > 0 else float(self.r) # Freezing the pre-trained weight matrix self.weight.requires_grad = False self.ranknum.requires_grad = False @@ -385,34 +365,33 @@ if is_bnb_available(): if x.dtype != torch.float32: x = x.float() output = ( - self.lora_dropout(x) @ (self.lora_A*self.lora_E).T @ self.lora_B.T /(self.ranknum+1e-5) - ).to(expected_dtype) * self.scaling + self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T / (self.ranknum + 1e-5) + ).to(expected_dtype) * self.scaling result += output else: output = ( - self.lora_dropout(x) @ (self.lora_A*self.lora_E).T @ self.lora_B.T /(self.ranknum+1e-5) - ) * self.scaling + self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T / (self.ranknum + 1e-5) + ) * self.scaling result += output return result - class RankAllocator(object): """ - The RankAllocator for AdaLoraModel. - Paper: https://openreview.net/pdf?id=lq62uWRJjiY + The RankAllocator for AdaLoraModel. Paper: https://openreview.net/pdf?id=lq62uWRJjiY Args: config ([`AdaLoraConfig`]): The configuration of the AdaLora model. - model: the model that we apply AdaLoRA to. + model: the model that we apply AdaLoRA to. """ + def __init__(self, peft_config, model): self.peft_config = peft_config - self.beta1 = peft_config.beta1 - self.beta2 = peft_config.beta2 - assert (self.beta1>0 and self.beta1<1) - assert (self.beta2>0 and self.beta2<1) + self.beta1 = peft_config.beta1 + self.beta2 = peft_config.beta2 + assert self.beta1 > 0 and self.beta1 < 1 + assert self.beta2 > 0 and self.beta2 < 1 self.reset_ipt() self._set_budget_scheduler(model) @@ -421,129 +400,125 @@ class RankAllocator(object): self.peft_config.total_step = total_step def reset_ipt(self): - self.ipt = {} + self.ipt = {} self.exp_avg_ipt = {} self.exp_avg_unc = {} def _set_budget_scheduler(self, model): - self.init_bgt = 0 - self.name_set = set() - for n,p in model.named_parameters(): - if "lora_A" in n: - self.init_bgt += p.size(0) + self.init_bgt = 0 + self.name_set = set() + for n, p in model.named_parameters(): + if "lora_A" in n: + self.init_bgt += p.size(0) self.name_set.add(n.replace("lora_A", "%s")) - self.name_set = list(sorted(self.name_set)) - # The total final rank budget - self.target_bgt = self.peft_config.target_r * len(self.name_set) + self.name_set = sorted(self.name_set) + # The total final rank budget + self.target_bgt = self.peft_config.target_r * len(self.name_set) - def budget_schedule(self, step:int): - tinit = self.peft_config.tinit - tfinal = self.peft_config.tfinal - total_step = self.peft_config.total_step - # Initial warmup - if step <= tinit: - budget = self.init_bgt - mask_ind = False - # Final fine-tuning - elif step > total_step - tfinal: - budget = self.target_bgt - mask_ind = True - else: - # Budget decreasing with a cubic scheduler - mul_coeff = 1 - (step-tinit) / (total_step-tfinal-tinit) - budget = int( - (self.init_bgt-self.target_bgt)*(mul_coeff**3)+self.target_bgt - ) - mask_ind = True if step % self.peft_config.deltaT == 0 else False - return budget, mask_ind + def budget_schedule(self, step: int): + tinit = self.peft_config.tinit + tfinal = self.peft_config.tfinal + total_step = self.peft_config.total_step + # Initial warmup + if step <= tinit: + budget = self.init_bgt + mask_ind = False + # Final fine-tuning + elif step > total_step - tfinal: + budget = self.target_bgt + mask_ind = True + else: + # Budget decreasing with a cubic scheduler + mul_coeff = 1 - (step - tinit) / (total_step - tfinal - tinit) + budget = int((self.init_bgt - self.target_bgt) * (mul_coeff**3) + self.target_bgt) + mask_ind = True if step % self.peft_config.deltaT == 0 else False + return budget, mask_ind - def update_ipt(self, model): - # Update the sensitivity and uncertainty for every weight - for n,p in model.named_parameters(): - if "lora_" in n: + def update_ipt(self, model): + # Update the sensitivity and uncertainty for every weight + for n, p in model.named_parameters(): + if "lora_" in n: if n not in self.ipt: - self.ipt[n] = torch.zeros_like(p) - self.exp_avg_ipt[n] = torch.zeros_like(p) - self.exp_avg_unc[n] = torch.zeros_like(p) + self.ipt[n] = torch.zeros_like(p) + self.exp_avg_ipt[n] = torch.zeros_like(p) + self.exp_avg_unc[n] = torch.zeros_like(p) with torch.no_grad(): self.ipt[n] = (p * p.grad).abs().detach() - # Sensitivity smoothing - self.exp_avg_ipt[n] = self.beta1 * self.exp_avg_ipt[n] + \ - (1 - self.beta1)*self.ipt[n] - # Uncertainty quantification - self.exp_avg_unc[n] = self.beta2 * self.exp_avg_unc[n] + \ - (1-self.beta2)*(self.ipt[n]-self.exp_avg_ipt[n]).abs() + # Sensitivity smoothing + self.exp_avg_ipt[n] = self.beta1 * self.exp_avg_ipt[n] + (1 - self.beta1) * self.ipt[n] + # Uncertainty quantification + self.exp_avg_unc[n] = ( + self.beta2 * self.exp_avg_unc[n] + (1 - self.beta2) * (self.ipt[n] - self.exp_avg_ipt[n]).abs() + ) def _element_score(self, n): return self.exp_avg_ipt[n] * self.exp_avg_unc[n] def _combine_ipt(self, ipt_E, ipt_AB): ipt_AB = ipt_AB.sum(dim=1, keepdim=False) - sum_ipt = ipt_E.view(-1) + ipt_AB.view(-1) - return sum_ipt + sum_ipt = ipt_E.view(-1) + ipt_AB.view(-1) + return sum_ipt - def mask_to_budget(self, model, budget): + def mask_to_budget(self, model, budget): value_ipt = {} - vector_ipt = {} + vector_ipt = {} triplet_ipt = {} # Get the importance score for A, E, B - for n,p in model.named_parameters(): - if "lora_A" in n: + for n, p in model.named_parameters(): + if "lora_A" in n: entry_ipt = self._element_score(n) comb_ipt = torch.mean(entry_ipt, dim=1, keepdim=True) name_m = n.replace("lora_A", "%s") - if name_m not in vector_ipt: + if name_m not in vector_ipt: vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) - if "lora_B" in n: + if "lora_B" in n: entry_ipt = self._element_score(n) comb_ipt = torch.mean(entry_ipt, dim=0, keepdim=False).view(-1, 1) name_m = n.replace("lora_B", "%s") - if name_m not in vector_ipt: + if name_m not in vector_ipt: vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) if "lora_E" in n: - entry_ipt = self._element_score(n) + entry_ipt = self._element_score(n) name_m = n.replace("lora_E", "%s") value_ipt[name_m] = entry_ipt all_score = [] - # Calculate the score for each triplet - for name_m in vector_ipt: - ipt_E = value_ipt[name_m] + # Calculate the score for each triplet + for name_m in vector_ipt: + ipt_E = value_ipt[name_m] ipt_AB = torch.cat(vector_ipt[name_m], dim=1) sum_ipt = self._combine_ipt(ipt_E, ipt_AB) - name_E = name_m%"lora_E" + name_E = name_m % "lora_E" triplet_ipt[name_E] = sum_ipt.view(-1, 1) all_score.append(sum_ipt.view(-1)) # Get the threshold by ranking ipt mask_threshold = torch.kthvalue( - torch.cat(all_score), - k = self.init_bgt - budget, + torch.cat(all_score), + k=self.init_bgt - budget, )[0].item() rank_pattern = {} - # Mask the unimportant triplets + # Mask the unimportant triplets with torch.no_grad(): - for n,p in model.named_parameters(): - if "lora_E" in n: - p.masked_fill_(triplet_ipt[n]<=mask_threshold, 0.0) - rank_pattern[n] = (~(triplet_ipt[n]<=mask_threshold)).view(-1).tolist() + for n, p in model.named_parameters(): + if "lora_E" in n: + p.masked_fill_(triplet_ipt[n] <= mask_threshold, 0.0) + rank_pattern[n] = (~(triplet_ipt[n] <= mask_threshold)).view(-1).tolist() return rank_pattern def update_and_allocate(self, model, global_step, force_mask=False): - # # Update the importance score and allocate the budget + # # Update the importance score and allocate the budget if global_step < self.peft_config.total_step - self.peft_config.tfinal: self.update_ipt(model) budget, mask_ind = self.budget_schedule(global_step) - # Allocate the budget according to importance scores + # Allocate the budget according to importance scores if mask_ind or force_mask: rank_pattern = self.mask_to_budget(model, budget) else: - rank_pattern = None + rank_pattern = None return budget, rank_pattern - - diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index cf4c813..0b10717 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -72,10 +72,10 @@ def set_peft_model_state_dict(model, peft_model_state_dict): """ if model.peft_config.peft_type == PeftType.ADALORA: rank_pattern = model.peft_config.rank_pattern - if rank_pattern: + if rank_pattern: model.base_model.resize_modules_by_rank_pattern(rank_pattern) model.load_state_dict(peft_model_state_dict, strict=False) - + if model.peft_config.peft_type not in (PeftType.LORA, PeftType.ADALORA): model.prompt_encoder.embedding.load_state_dict( {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True From 739716043504f8f8980c51d6202be0e1df7a39ec Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:06:33 +0530 Subject: [PATCH 66/76] making adalora compatible with multiple adapters --- src/peft/tuners/adalora.py | 433 ++++++++++++++++++--------- src/peft/tuners/lora.py | 27 +- src/peft/utils/__init__.py | 2 + src/peft/utils/other.py | 26 ++ tests/test_decoder_models.py | 2 +- tests/test_encoder_decoder_models.py | 2 +- tests/testing_common.py | 8 +- 7 files changed, 328 insertions(+), 172 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index f98b3c6..938b83e 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -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 diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 74ed008..90e4a23 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -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 diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py index bfaabe8..346b667 100644 --- a/src/peft/utils/__init__.py +++ b/src/peft/utils/__init__.py @@ -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 diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 53e2ee8..1bfbbfb 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -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, } diff --git a/tests/test_decoder_models.py b/tests/test_decoder_models.py index f0a575e..209b4df 100644 --- a/tests/test_decoder_models.py +++ b/tests/test_decoder_models.py @@ -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", }, ) diff --git a/tests/test_encoder_decoder_models.py b/tests/test_encoder_decoder_models.py index c7c73e6..cdf9571 100644 --- a/tests/test_encoder_decoder_models.py +++ b/tests/test_encoder_decoder_models.py @@ -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", }, ) diff --git a/tests/testing_common.py b/tests/testing_common.py index 38aee14..cfe6cf2 100644 --- a/tests/testing_common.py +++ b/tests/testing_common.py @@ -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) From 74e2a3da50e37c700cade042b240d82d4e45a1e2 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:16:37 +0530 Subject: [PATCH 67/76] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 938b83e..048832f 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -95,7 +95,6 @@ class AdaLoraModel(LoraModel): nn.Module.__init__(self) self.model = model self.peft_config = config - self.rankallocator = RankAllocator(config, self.model) self.add_adapter(adapter_name, self.peft_config[adapter_name]) def add_adapter(self, adapter_name, config=None): From b728f5f559c65e9051eb48fbea86874b068d3a5b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:20:13 +0530 Subject: [PATCH 68/76] =?UTF-8?q?=F0=9F=90=9B=20fixing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 048832f..6872f64 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -343,13 +343,13 @@ class AdaLoraLayer(LoraLayer): if r > 0: # Right singular vectors self.lora_A.update( - nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros((r, self.in_features)))}) + nn.ParameterDict({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))})) + self.lora_E.update(nn.ParameterDict({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)))}) + nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros((self.out_features, r)))}) ) # The current rank self.ranknum.update( From dee2a96fea700743a723e608cfa9756e517b102c Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:22:56 +0530 Subject: [PATCH 69/76] Update adalora.py --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 6872f64..53abce3 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -355,7 +355,7 @@ class AdaLoraLayer(LoraLayer): 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].data.fill_(float(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: From b6c751455e9286290aa9da9d7f9dfe0c0148ee8f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:31:21 +0530 Subject: [PATCH 70/76] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 2b7b286..9bcaf80 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -331,7 +331,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): 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 + peft_config.inference_mode = not is_trainable self.add_adapter(adapter_name, peft_config) # load weights if any From 07a4b8aacc840ca32fa3040a317644672da03e97 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:56:55 +0530 Subject: [PATCH 71/76] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 53abce3..8c84a55 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -134,7 +134,7 @@ class AdaLoraModel(LoraModel): ) is_target_modules_in_base_model = False kwargs = { - "r": lora_config.r, + "r": lora_config.init_r, "lora_alpha": lora_config.lora_alpha, "lora_dropout": lora_config.lora_dropout, "fan_in_fan_out": lora_config.fan_in_fan_out, @@ -154,7 +154,7 @@ class AdaLoraModel(LoraModel): if isinstance(target, LoraLayer): target.update_layer( adapter_name, - lora_config.r, + lora_config.init_r, lora_config.lora_alpha, lora_config.lora_dropout, lora_config.init_lora_weights, From 3aaf482704a5348e51d36819deac06d759d957b4 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 20:02:31 +0530 Subject: [PATCH 72/76] fix --- src/peft/utils/save_and_load.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index dc0a391..22792e7 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -52,7 +52,7 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): 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()} + 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))} From a591b4b905a419295a74f816c33c697d92a6a5eb Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 20:04:04 +0530 Subject: [PATCH 73/76] final fix I guess --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 8c84a55..06b18de 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -282,7 +282,7 @@ class AdaLoraModel(LoraModel): for name, rank_idx in rank_pattern.items(): key = ".".join(name.split(".")[0:-1]) key = f"{key}.{adapter_name}" if adapter_name not in key else key - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(self.model, key) new_module = self._prepare_new_module(target, rank_idx, adapter_name) self._replace_module(parent, target_name, new_module, target) From 3258b709a3b08d6c6fb16f5fa02ca4665ed525af Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 20:41:36 +0530 Subject: [PATCH 74/76] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 76 ++++++++++++++------------------------ 1 file changed, 27 insertions(+), 49 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 06b18de..a8cd3b7 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -233,58 +233,36 @@ class AdaLoraModel(LoraModel): outputs.loss += orth_reg_weight * regu_loss return outputs - 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): - rank_idx = rank_idx.view(-1) - 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": 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) - 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) - elif isinstance(target, torch.nn.Linear): - 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(): - new_module.weight.copy_(target.weight) - if bias: - new_module.bias.copy_(target.bias) - if rank > 0: - 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[adapter_name].copy_(target.ranknum) - return new_module - def resize_modules_by_rank_pattern(self, rank_pattern, adapter_name): + lora_config = self.peft_config[adapter_name] for name, rank_idx in rank_pattern.items(): + 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") key = ".".join(name.split(".")[0:-1]) - key = f"{key}.{adapter_name}" if adapter_name not in key else key - parent, target, target_name = _get_submodules(self.model, key) - new_module = self._prepare_new_module(target, rank_idx, adapter_name) - self._replace_module(parent, target_name, new_module, target) + _, target, _ = _get_submodules(self.model, key) + lora_E_weights = target.lora_E[adapter_name][rank_idx] + lora_A_weights = target.lora_A[adapter_name][rank_idx] + lora_B_weights = target.lora_B[adapter_name][:, rank_idx] + ranknum = target.ranknum[adapter_name] + target.update_layer( + adapter_name, + rank, + lora_config.lora_alpha, + lora_config.lora_dropout, + lora_config.init_lora_weights, + ) + with torch.no_grad(): + if rank > 0: + target.lora_E[adapter_name].copy_(lora_E_weights) + target.lora_A[adapter_name].copy_(lora_A_weights) + target.lora_B[adapter_name].copy_(lora_B_weights) + # The scaling is exactly as the previous + target.ranknum[adapter_name].copy_(ranknum) def update_and_allocate(self, global_step): lora_config = self.peft_config[self.trainable_adapter_name] From d5feb8b787624bd9b886a4eb447eabf6d01b8bb2 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 21:17:54 +0530 Subject: [PATCH 75/76] =?UTF-8?q?fixing=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index a8cd3b7..1bbf7a2 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -243,7 +243,7 @@ class AdaLoraModel(LoraModel): rank = rank_idx.sum().item() else: raise ValueError("Unexcepted type of rank_idx") - key = ".".join(name.split(".")[0:-1]) + key = ".".join(name.split(".")[0:-2]) _, target, _ = _get_submodules(self.model, key) lora_E_weights = target.lora_E[adapter_name][rank_idx] lora_A_weights = target.lora_A[adapter_name][rank_idx] @@ -320,19 +320,13 @@ class AdaLoraLayer(LoraLayer): # Actual trainable parameters if r > 0: # Right singular vectors - self.lora_A.update( - nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros((r, self.in_features)))}) - ) + self.lora_A.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, self.in_features))})) # Singular values - self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros(r, 1))})) + self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, 1))})) # Left singular vectors - self.lora_B.update( - nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros((self.out_features, r)))}) - ) + self.lora_B.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.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.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(1), requires_grad=False)})) self.ranknum[adapter_name].data.fill_(float(r)) self.ranknum[adapter_name].requires_grad = False self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) From e8b0085d2b09735566a38e420747222834b80262 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 7 Apr 2023 04:08:10 +0530 Subject: [PATCH 76/76] fixing adalora saving and loading --- src/peft/peft_model.py | 4 +-- src/peft/tuners/adalora.py | 61 +++++++++++++++++++++++++-------- src/peft/utils/save_and_load.py | 14 +++++--- 3 files changed, 58 insertions(+), 21 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 9bcaf80..dd68deb 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -393,8 +393,8 @@ 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() + # Set model in evaluation mode to deactivate Dropout modules by default + self.eval() def set_adapter(self, adapter_name): """ diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 1bbf7a2..fc6261f 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -243,7 +243,7 @@ class AdaLoraModel(LoraModel): rank = rank_idx.sum().item() else: raise ValueError("Unexcepted type of rank_idx") - key = ".".join(name.split(".")[0:-2]) + key = ".".join(name.split(".")[0:-2]) if adapter_name in name else ".".join(name.split(".")[0:-1]) _, target, _ = _get_submodules(self.model, key) lora_E_weights = target.lora_E[adapter_name][rank_idx] lora_A_weights = target.lora_A[adapter_name][rank_idx] @@ -264,6 +264,22 @@ class AdaLoraModel(LoraModel): # The scaling is exactly as the previous target.ranknum[adapter_name].copy_(ranknum) + def resize_state_dict_by_rank_pattern(self, rank_pattern, state_dict, adapter_name): + for name, rank_idx in rank_pattern.items(): + rank = sum(rank_idx) + prefix = ".".join(name.split(".")[0:-2]) if adapter_name in name else ".".join(name.split(".")[0:-1]) + for layer in ["lora_E", "lora_A", "lora_B"]: + key = f"base_model.model.{prefix}.{layer}.{adapter_name}" + if layer != "lora_B": + state_dict[key] = ( + state_dict[key][rank_idx] if rank != state_dict[key].shape[0] else state_dict[key] + ) + else: + state_dict[key] = ( + state_dict[key][:, rank_idx] if rank != state_dict[key].shape[1] else state_dict[key] + ) + return state_dict + def update_and_allocate(self, global_step): lora_config = self.peft_config[self.trainable_adapter_name] # Update the importance score and allocate the budget @@ -274,9 +290,14 @@ class AdaLoraModel(LoraModel): # Finalize the budget allocation 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) + # for some reason, this freezes the trainable parameters and nothing gets updates + # self.resize_modules_by_rank_pattern(rank_pattern, self.trainable_adapter_name) lora_config.rank_pattern = rank_pattern self.rankallocator.reset_ipt() + # Currently using inefficient way to mask the unimportant weights using the rank pattern + # due to problem mentioned above + elif global_step > lora_config.total_step - lora_config.tfinal: + self.rankallocator.mask_using_rank_pattern(self.model, lora_config.rank_pattern) # Pass the function and do forward propagation else: return None @@ -318,18 +339,17 @@ class AdaLoraLayer(LoraLayer): 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.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, self.in_features))})) - # Singular values - self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, 1))})) - # Left singular vectors - self.lora_B.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(self.out_features, r))})) - # The current rank - self.ranknum.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(1), requires_grad=False)})) - self.ranknum[adapter_name].data.fill_(float(r)) - self.ranknum[adapter_name].requires_grad = False - self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) + # Right singular vectors + self.lora_A.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, self.in_features))})) + # Singular values + self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, 1))})) + # Left singular vectors + self.lora_B.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(self.out_features, r))})) + # The current rank + self.ranknum.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(1), requires_grad=False)})) + self.ranknum[adapter_name].data.fill_(float(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) @@ -638,3 +658,16 @@ class RankAllocator(object): else: rank_pattern = None return budget, rank_pattern + + def mask_using_rank_pattern(self, model, rank_pattern): + # Mask the unimportant triplets + is_adapter_name_truncated = False + if self.adapter_name not in next(iter(rank_pattern.keys())): + is_adapter_name_truncated = True + + with torch.no_grad(): + for n, p in model.named_parameters(): + if f"lora_E.{self.adapter_name}" in n: + key = n if not is_adapter_name_truncated else n.replace(f".{self.adapter_name}", "") + mask = torch.Tensor(rank_pattern[key]).unsqueeze(-1).to(p.device) + p.masked_fill_(~mask.bool(), 0.0) diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 22792e7..2876bbe 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -49,13 +49,13 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): to_return[bias_name] = state_dict[bias_name] else: raise NotImplementedError + to_return = {k: v for k, v in to_return.items() if (("lora_" in k and adapter_name in k) or ("bias" in k))} 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))} + to_return = model.resize_state_dict_by_rank_pattern(rank_pattern, to_return, adapter_name) elif isinstance(config, PromptLearningConfig): to_return = {} if config.inference_mode: @@ -70,7 +70,7 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): if any(f"{module_name}.modules_to_save.{adapter_name}" in key for module_name in model.modules_to_save): to_return[key.replace("modules_to_save.", "")] = value - to_return = {k.replace(f"{adapter_name}.", ""): v for k, v in to_return.items()} + to_return = {k.replace(f".{adapter_name}", ""): v for k, v in to_return.items()} return to_return @@ -99,8 +99,12 @@ def set_peft_model_state_dict(model, peft_model_state_dict, adapter_name="defaul peft_model_state_dict = {} for k, v in state_dict.items(): if "lora_" in k: - suffix_to_replace = ".".join(k.split("lora_")[1].split(".")[1:]) - k = k.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + suffix = k.split("lora_")[1] + if "." in suffix: + suffix_to_replace = ".".join(suffix.split(".")[1:]) + k = k.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + else: + k = f"{k}.{adapter_name}" peft_model_state_dict[k] = v else: peft_model_state_dict[k] = v