From 81eec9ba70e2b6f754350bf91cbb265bc9d2b99e Mon Sep 17 00:00:00 2001
From: Qingru Zhang <qrzhang@lawn-143-215-127-142.lawn.gatech.edu>
Date: Mon, 27 Feb 2023 21:08:55 -0500
Subject: [PATCH 01/26] 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 <qrzhang@lawn-143-215-112-153.lawn.gatech.edu>
Date: Tue, 28 Feb 2023 23:14:25 -0500
Subject: [PATCH 02/26] 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 <qrzhang@lawn-143-215-112-153.lawn.gatech.edu>
Date: Tue, 28 Feb 2023 23:18:19 -0500
Subject: [PATCH 03/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 1 Mar 2023 02:43:27 -0500
Subject: [PATCH 04/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 1 Mar 2023 02:47:33 -0500
Subject: [PATCH 05/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 1 Mar 2023 21:26:07 +0000
Subject: [PATCH 06/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 1 Mar 2023 21:52:33 +0000
Subject: [PATCH 07/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 1 Mar 2023 21:55:31 +0000
Subject: [PATCH 08/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 2 Mar 2023 01:04:48 +0000
Subject: [PATCH 09/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 2 Mar 2023 01:08:41 +0000
Subject: [PATCH 10/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 2 Mar 2023 01:11:14 +0000
Subject: [PATCH 11/26] 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 1141b125d0e7f4590c2bb53b03faf5b5888b3399 Mon Sep 17 00:00:00 2001
From: Qingru Zhang <qingru.zhang@gatech.edu>
Date: Wed, 29 Mar 2023 19:56:23 -0400
Subject: [PATCH 12/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 30 Mar 2023 00:45:33 +0000
Subject: [PATCH 13/26] 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 <zqingru@bcd0745dc556.ant.amazon.com>
Date: Wed, 29 Mar 2023 21:03:48 -0400
Subject: [PATCH 14/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 30 Mar 2023 01:04:11 +0000
Subject: [PATCH 15/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 30 Mar 2023 05:53:04 +0000
Subject: [PATCH 16/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 30 Mar 2023 01:59:26 -0400
Subject: [PATCH 17/26] 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 <qingru.zhang@gatech.edu>
Date: Thu, 30 Mar 2023 06:19:45 +0000
Subject: [PATCH 18/26] 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 3e6a88a8f9f063f93010ef94d8646d93e93bc056 Mon Sep 17 00:00:00 2001
From: Qingru Zhang <qingru.zhang@gatech.edu>
Date: Wed, 5 Apr 2023 16:23:48 -0400
Subject: [PATCH 19/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 5 Apr 2023 16:30:16 -0400
Subject: [PATCH 20/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 5 Apr 2023 16:31:08 -0400
Subject: [PATCH 21/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 5 Apr 2023 16:31:46 -0400
Subject: [PATCH 22/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 5 Apr 2023 16:31:58 -0400
Subject: [PATCH 23/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 5 Apr 2023 16:32:23 -0400
Subject: [PATCH 24/26] 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 <zqingru@amazon.com>
Date: Wed, 5 Apr 2023 20:40:06 +0000
Subject: [PATCH 25/26] 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 <qingru.zhang@gatech.edu>
Date: Wed, 5 Apr 2023 20:52:12 +0000
Subject: [PATCH 26/26] 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