Add scalewise lora for attn

cli/conf/lsf-setup/multi_scale/finetune_two_stage/default.yaml

@@ -23,7 +23,7 @@ trainer_warmup:
   logger:
       _target_: lightning.pytorch.loggers.TensorBoardLogger
       save_dir: ${hydra:runtime.output_dir}
-      name: logs
+      name: logs_warmup
   callbacks:
     - _target_: lightning.pytorch.callbacks.LearningRateMonitor
       logging_interval: epoch

project/lsf-setup/lsf/finetune/base/electricity.sh

@@ -33,5 +33,6 @@ for pl in 96 192 336 720; do
   val_data.prediction_length=$pl \
   val_data.mode=${mode} \
   train_dataloader.batch_size=256 \
-  model.lr=5e-6
+  model.lr=1e-5 \
+  trainer.callbacks.'3'.patience=1
 done

project/lsf-setup/multi_scale/eval/small/ettm1.sh

@@ -9,10 +9,10 @@ exp_name=lsf
 cl=4000
 model=moirai_lightning_ckpt
 
-cpp1='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm1/S/cl4000_pl96/checkpoints/epoch_4-step_2085.ckpt'
-cpp2='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm1/S/cl4000_pl192/checkpoints/epoch_1-step_832.ckpt'
-cpp3='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm1/S/cl4000_pl336/checkpoints/epoch_0-step_414.ckpt'
-cpp4='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm1/S/cl4000_pl720/checkpoints/epoch_0-step_408.ckpt'
+cpp1='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm1/S/cl4000_pl96/checkpoints/epoch_5-step_2502.ckpt'
+cpp2='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm1/S/cl4000_pl192/checkpoints/epoch_1-step_832.ckpt'
+cpp3='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm1/S/cl4000_pl336/checkpoints/epoch_0-step_414.ckpt'
+cpp4='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm1/S/cl4000_pl720/checkpoints/epoch_0-step_408.ckpt'
 
 index=1
 for pl in 96 192 336 720; do

project/lsf-setup/multi_scale/eval/small/ettm2.sh

@@ -1,18 +1,18 @@
 #!/bin/bash
 
 export HYDRA_FULL_ERROR=1
-export CUDA_VISIBLE_DEVICES=3
+export CUDA_VISIBLE_DEVICES=2
 
 mode=S
 cp=conf/lsf-setup/multi_scale/eval
 exp_name=lsf
 cl=3000
 model=moirai_lightning_ckpt
 
-cpp1='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm2/S/cl3000_pl96/checkpoints/epoch_13-step_6034.ckpt'
-cpp2='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm2/S/cl3000_pl192/checkpoints/epoch_4-step_2145.ckpt'
-cpp3='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm2/S/cl3000_pl336/checkpoints/epoch_1-step_854.ckpt'
-cpp4='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage/full/ettm2/S/cl3000_pl720/checkpoints/epoch_0-step_422.ckpt'
+cpp1='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm2/S/cl3000_pl96/checkpoints/epoch_16-step_7327.ckpt'
+cpp2='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm2/S/cl3000_pl192/checkpoints/epoch_3-step_1716.ckpt'
+cpp3='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm2/S/cl3000_pl336/checkpoints/epoch_1-step_854.ckpt'
+cpp4='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/ettm2/S/cl3000_pl720/checkpoints/epoch_0-step_422.ckpt'
 
 index=1
 for pl in 96 192 336 720; do

project/lsf-setup/multi_scale/eval/small/weather.sh

@@ -1,30 +1,26 @@
 #!/bin/bash
 
 export HYDRA_FULL_ERROR=1
-export CUDA_VISIBLE_DEVICES=1
+export CUDA_VISIBLE_DEVICES=3
 
 mode=S
 cp=conf/lsf-setup/multi_scale/eval
 exp_name=lsf
 cl=2000
 model=moirai_lightning_ckpt
 
-cpp1='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage_valMSE/full/weather/S/cl2000_pl96/checkpoints/epoch_7-step_11424.ckpt'
-cpp2='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage_valMSE/full/weather/S/cl2000_pl192/checkpoints/epoch_6-step_9968.ckpt'
-cpp3='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage_valMSE/full/weather/S/cl2000_pl336/checkpoints/epoch_2-step_4254.ckpt'
-cpp4='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/learned_time_id_2stage_valMSE/full/weather/S/cl2000_pl720/checkpoints/epoch_1-step_2804.ckpt'
+cpp1='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/weather/S/cl2000_pl96/checkpoints/epoch_14-step_21420.ckpt'
+cpp2='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/weather/S/cl2000_pl192/checkpoints/epoch_9-step_14240.ckpt'
+cpp3='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/weather/S/cl2000_pl336/checkpoints/epoch_5-step_8508.ckpt'
+cpp4='./outputs/lsf-setup/multi_scale/finetune_two_stage/moirai_1.0_R_small/1tid_2inproj_all_scale_lora_freezeqkv/full/weather/S/cl2000_pl720/checkpoints/epoch_1-step_2804.ckpt'
 
 index=1
-for pl in 192 336 720; do  # 96
+for pl in 96 192 336 720; do
   case $index in
-    1) cpp=$cpp2 ;;
-    2) cpp=$cpp3 ;;
-    3) cpp=$cpp4 ;;
-
-#    1) cpp=$cpp1 ;;
-#    2) cpp=$cpp2 ;;
-#    3) cpp=$cpp3 ;;
-#    4) cpp=$cpp4 ;;
+    1) cpp=$cpp1 ;;
+    2) cpp=$cpp2 ;;
+    3) cpp=$cpp3 ;;
+    4) cpp=$cpp4 ;;
   esac
 
   pretrained_model=$(echo $cpp | cut -d'/' -f6)

project/lsf-setup/multi_scale/finetune_two_stage/small/ettm1.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
 
-export HYDRA_FULL_ERROR=1; export CUDA_VISIBLE_DEVICES=0;
+export HYDRA_FULL_ERROR=1; export CUDA_VISIBLE_DEVICES=1;
 
 model=moirai_1.0_R_small
 cp=conf/lsf-setup/multi_scale/finetune_two_stage
-exp_name=learned_time_id_2stage
+exp_name=1tid_2inproj_all_scale_lora_freezeqkv
 data=ettm1
 cl=4000
 ps=128

project/lsf-setup/multi_scale/finetune_two_stage/small/ettm2.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
 
-export HYDRA_FULL_ERROR=1; export CUDA_VISIBLE_DEVICES=2;
+export HYDRA_FULL_ERROR=1; export CUDA_VISIBLE_DEVICES=3;
 
 model=moirai_1.0_R_small
 cp=conf/lsf-setup/multi_scale/finetune_two_stage
-exp_name=learned_time_id_2stage
+exp_name=1tid_2inproj_all_scale_lora_freezeqkv
 data=ettm2
 cl=3000
 ps=64

project/lsf-setup/multi_scale/finetune_two_stage/small/weather.sh

@@ -4,7 +4,7 @@ export HYDRA_FULL_ERROR=1; export CUDA_VISIBLE_DEVICES=0;
 
 model=moirai_1.0_R_small
 cp=conf/lsf-setup/multi_scale/finetune_two_stage
-exp_name=learned_time_id_2stage_valMSE
+exp_name=1tid_2inproj_all_scale_lora_freezeqkv
 data=weather
 cl=2000
 ps=128
@@ -33,5 +33,7 @@ for pl in 96 192 336 720; do
   val_data.prediction_length=$pl \
   val_data.mode=${mode} \
   trainer.callbacks."1".monitor=val/PackedMSELoss \
-  trainer.callbacks."3".monitor=val/PackedMSELoss
+  trainer.callbacks."3".monitor=val/PackedMSELoss \
+  trainer_warmup.callbacks."1".monitor=val/PackedMSELoss \
+  trainer_warmup.callbacks."2".monitor=val/PackedMSELoss
 done

src/uni2ts/model/multi_scale_moirai/finetune.py

@@ -143,11 +143,11 @@ def post_init(self):
 
         self.module.post_init(self.token_idx_per_scale, self.base_ctx_token_idx, self.patch_size)
 
-        # for layer in self.module.encoder.layers:
-        #     # Check if the layer has an attribute named `self_attn` and if it is an instance of GroupedQueryAttention
-        #     if hasattr(layer, 'self_attn') and isinstance(layer.self_attn, GroupedQueryAttention):
-        #         # Call post_init() method of the GroupedQueryAttention object
-        #         layer.self_attn.init_multi_scale_modules(self.context_length, self.patch_size, self.num_new_scales, self.ds_factor)
+        for layer in self.module.encoder.layers:
+            # Check if the layer has an attribute named `self_attn` and if it is an instance of GroupedQueryAttention
+            if hasattr(layer, 'self_attn') and isinstance(layer.self_attn, GroupedQueryAttention):
+                # Call post_init() method of the GroupedQueryAttention object
+                layer.self_attn.init_multi_scale_modules(self.context_length, self.patch_size, self.num_new_scales, self.ds_factor)
 
         for module in self.module.encoder.modules():
             if isinstance(module, MultiScaleRotaryProjection):
@@ -455,6 +455,12 @@ def configure_optimizers(self) -> dict:
                     decay.add(fpn)
                 elif 'pe_weights' in pn:
                     decay.add(fpn)
+                elif 'q_A' in pn or 'q_B' in pn or 'q_bias' in pn:
+                    decay.add(fpn)
+                elif 'k_A' in pn or 'k_B' in pn or 'k_bias' in pn:
+                    decay.add(fpn)
+                elif 'v_A' in pn or 'v_B' in pn or 'v_bias' in pn:
+                    decay.add(fpn)
 
                 # elif 'layers.0.self_attn.time_qk_proj.query_proj.pe_weights' in pn:  # Shared time_qk_proj
                 #     decay.add(fpn)

src/uni2ts/model/multi_scale_moirai/finetune_two_stage.py

@@ -145,11 +145,11 @@ def post_init(self):
 
         self.module.post_init(self.token_idx_per_scale, self.base_ctx_token_idx, self.patch_size)
 
-        # for layer in self.module.encoder.layers:
-        #     # Check if the layer has an attribute named `self_attn` and if it is an instance of GroupedQueryAttention
-        #     if hasattr(layer, 'self_attn') and isinstance(layer.self_attn, GroupedQueryAttention):
-        #         # Call post_init() method of the GroupedQueryAttention object
-        #         layer.self_attn.init_multi_scale_modules(self.context_length, self.patch_size, self.num_new_scales, self.ds_factor)
+        for layer in self.module.encoder.layers:
+            # Check if the layer has an attribute named `self_attn` and if it is an instance of GroupedQueryAttention
+            if hasattr(layer, 'self_attn') and isinstance(layer.self_attn, GroupedQueryAttention):
+                # Call post_init() method of the GroupedQueryAttention object
+                layer.self_attn.init_multi_scale_modules(self.context_length, self.patch_size, self.num_new_scales, self.ds_factor)
 
         for module in self.module.encoder.modules():
             if isinstance(module, MultiScaleRotaryProjection):
@@ -331,6 +331,48 @@ def validation_step(
 
     def configure_optimizers(self) -> dict:
 
+        # if self.current_stage == 1:
+        #     warmup_pn_group1 = ['param_proj', 'time_id_q_proj', 'time_id_k_proj']
+        #     warmup_pn_group2 = ['in_proj_new_scales']
+        #
+        #     optim_groups = [
+        #         {
+        #             "params": [
+        #                 p for pn, p in self.named_parameters()
+        #                 if any(keyword in pn for keyword in warmup_pn_group1)
+        #             ],
+        #             "lr": 5e-4,
+        #             "weight_decay": self.hparams.weight_decay,
+        #         },
+        #         {
+        #             "params": [
+        #                 p for pn, p in self.named_parameters()
+        #                 if any(keyword in pn for keyword in warmup_pn_group2)
+        #             ],
+        #             "lr": 5e-6,
+        #             "weight_decay": self.hparams.weight_decay,
+        #         },
+        #     ]
+        #
+        #     optimizer = torch.optim.AdamW(
+        #         optim_groups,
+        #         betas=(self.hparams.beta1, self.hparams.beta2),
+        #         eps=1e-6,
+        #     )
+        #
+        #     warmup_params_all = {
+        #         pn: p for pn, p in self.named_parameters()
+        #         if any(keyword in pn for keyword in warmup_pn_group1 + warmup_pn_group2)
+        #     }
+        #     self.trainable_params = warmup_params_all
+        #
+        #     scheduler = get_scheduler(
+        #         SchedulerType.CONSTANT,  # Use constant lr scheduler
+        #         optimizer,
+        #         num_warmup_steps=self.hparams.num_warmup_steps,
+        #         num_training_steps=self.hparams.num_training_steps,
+        #     )
+
         if self.current_stage == 1:
             warmup_pn = ['param_proj', 'time_id_q_proj', 'time_id_k_proj']
             warmup_params = {
@@ -474,6 +516,12 @@ def configure_optimizers(self) -> dict:
                         decay.add(fpn)
                     elif 'pe_weights' in pn:
                         decay.add(fpn)
+                    elif 'q_A' in pn or 'q_B' in pn:
+                        decay.add(fpn)
+                    elif 'k_A' in pn or 'k_B' in pn:
+                        decay.add(fpn)
+                    elif 'v_A' in pn or 'v_B' in pn:
+                        decay.add(fpn)
 
                     # elif 'layers.0.self_attn.time_qk_proj.query_proj.pe_weights' in pn:  # Shared time_qk_proj
                     #     decay.add(fpn)

src/uni2ts/model/multi_scale_moirai/forecast.py

@@ -143,11 +143,11 @@ def post_init(self):
 
         self.module.post_init(self.token_idx_per_scale, self.base_ctx_token_idx, self.hparams.patch_size)
 
-        # for layer in self.module.encoder.layers:
-        #     # Check if the layer has an attribute named `self_attn` and if it is an instance of GroupedQueryAttention
-        #     if hasattr(layer, 'self_attn') and isinstance(layer.self_attn, GroupedQueryAttention):
-        #         # Call post_init() method of the GroupedQueryAttention object
-        #         layer.self_attn.init_multi_scale_modules(self.hparams.context_length, self.hparams.patch_size, self.num_new_scales, self.ds_factor)
+        for layer in self.module.encoder.layers:
+            # Check if the layer has an attribute named `self_attn` and if it is an instance of GroupedQueryAttention
+            if hasattr(layer, 'self_attn') and isinstance(layer.self_attn, GroupedQueryAttention):
+                # Call post_init() method of the GroupedQueryAttention object
+                layer.self_attn.init_multi_scale_modules(self.hparams.context_length, self.hparams.patch_size, self.num_new_scales, self.ds_factor)
 
         for module in self.module.encoder.modules():
             if isinstance(module, MultiScaleRotaryProjection):