feat: make main parameters trainable

modeling_lora.py

@@ -207,11 +207,21 @@ class BertLoRA(BertPreTrainedModel):
             self.bert = BertModel(config, add_pooling_layer=add_pooling_layer)
         else:
             self.bert = bert
+        self._num_adaptions = num_adaptions
         self._register_lora(num_adaptions)
+        self.main_params_trainable = False
+        self.current_task = 0
+
+    @property
+    def main_params_trainable(self):
+        return self._main_params_trainable
+
+    @main_params_trainable.setter
+    def main_params_trainable(self, val):
+        self._main_params_trainable = val
         for name, param in super().named_parameters():
             if "lora" not in name:
-                param.requires_grad_(False)
-        self.current_task = 0
+                param.requires_grad_(val)
 
     @classmethod
     def from_bert(cls, *args, num_adaptions=1, **kwargs):
@@ -254,6 +264,7 @@ class BertLoRA(BertPreTrainedModel):
 
     @current_task.setter
     def current_task(self, task_idx: Union[None, int]):
+        assert task_idx is None or 0 <= task_idx < self._num_adaptions
         self._task_idx = task_idx
         self.apply(
             partial(LoRAParametrization.select_task_for_layer, task_idx=task_idx)
@@ -274,5 +285,5 @@ class BertLoRA(BertPreTrainedModel):
         for name, param in super().named_parameters(
             prefix=prefix, recurse=recurse, remove_duplicate=remove_duplicate
         ):
-            if "lora" in name:
+            if "lora" in name or self.main_params_trainable:
                 yield name, param