possibly convergent model

stanfordnlp · Nov 13, 2023 · ecd3eb5 · ecd3eb5
1 parent 1fc2c0f
commit ecd3eb5
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/stanza/models/coref/model.py b/stanza/models/coref/model.py
@@ -88,9 +88,12 @@ def __init__(self,
         if self.config.lora:
             self.__peft_config = LoraConfig(inference_mode=False,
                                             r=self.config.lora_rank,
-                                            target_modules=["query", "value"],
+                                            target_modules=["query", "value", "output.dense"],
                                             lora_alpha=self.config.lora_alpha,
-                                            lora_dropout=self.config.lora_dropout)
+                                            lora_dropout=self.config.lora_dropout,
+                                            modules_to_save=["intermediate"],
+                                            bias="none")
+
             self.bert = get_peft_model(self.bert, self.__peft_config)
             self.trainable["bert"] = self.bert
 
@@ -236,7 +239,8 @@ def load_state_dicts(self,
                 elif key.endswith("_scheduler"):
                     self.schedulers[key].load_state_dict(state_dict)
                 elif key.endswith("_lora"):
-                    assert self.config.lora, "Unable to load state dict of LoRA model into model initialized without set_peft_model_state_dict(self.trainable[key.split("_")[0]],
+                    assert self.config.lora, "Unable to load state dict of LoRA model into model initialized without LoRA!"
+                    set_peft_model_state_dict(self.trainable[key.split("_")[0]],
                                               state_dict)
                 else:
                     self.trainable[key].load_state_dict(state_dict, strict=False)