Revise configure_optimizer to finetune att params

zqiao11 · Sep 26, 2024 · bea1cf5 · bea1cf5
1 parent b6bcbfa
commit bea1cf5
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 3 deletions.
diff --git a/project/finetune/eval/small/run_multi.sh b/project/finetune/eval/small/run_multi.sh
@@ -0,0 +1,8 @@
+#!/bin/bash
+
+bash project/finetune/eval/small/etth1.sh
+bash project/finetune/eval/small/etth2.sh
+bash project/finetune/eval/small/ettm1.sh
+bash project/finetune/eval/small/ettm2.sh
+#bash project/finetune/eval/small/weather.sh
+
diff --git a/project/finetune/train/small/run_multi.sh b/project/finetune/train/small/run_multi.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+
+#bash project/finetune/train/small/weather.sh
+#bash project/finetune/train/small/weather_2.sh
+bash project/finetune/train/small/etth1.sh
+bash project/finetune/train/small/etth2.sh
+bash project/finetune/train/small/ettm1.sh
+bash project/finetune/train/small/ettm2.sh
+
+
diff --git a/src/uni2ts/model/moirai/finetune.py b/src/uni2ts/model/moirai/finetune.py
@@ -304,9 +304,41 @@ def configure_optimizers(self) -> dict:
                 if "ffn" in pn:
                     p.requires_grad = True
 
-        if "self_attn" in self.finetune_pattern:
-            # Todo: Analyze each component in self_attn & Lora's impact.
-            pass
+        # if "self_attn" in self.finetune_pattern:
+        #     # Todo: Analyze each component in self_attn & Lora's impact.
+        #     pass
+
+        if "q_proj" in self.finetune_pattern:
+            for pn, p in self.named_parameters():
+                if "q_proj" in pn:
+                    p.requires_grad = True
+
+        if "k_proj" in self.finetune_pattern:
+            for pn, p in self.named_parameters():
+                if "k_proj" in pn:
+                    p.requires_grad = True
+
+        if "v_proj" in self.finetune_pattern:
+            for pn, p in self.named_parameters():
+                if "v_proj" in pn:
+                    p.requires_grad = True
+
+        if "att_norm" in self.finetune_pattern:  #
+            for pn, p in self.named_parameters():
+                if "self_attn.q_norm" in pn or "self_attn.k_norm" in pn:
+                    p.requires_grad = True
+
+
+        if "var_attn_bias" in self.finetune_pattern:
+            for pn, p in self.named_parameters():
+                if "var_attn_bias" in pn:
+                    p.requires_grad = True
+
+        if "out_proj" in self.finetune_pattern:
+            for pn, p in self.named_parameters():
+                if "out_proj" in pn:
+                    p.requires_grad = True
+
 
         whitelist_params = (
             LearnedProjection,