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] 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)