Finish the test for model load and save

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)

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:

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)