From 82e44c72c8824578fd09d01e71d1d4dd349f7d4c Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Tue, 22 Jul 2025 00:48:29 +0000
Subject: [PATCH] more robust watcher

---
 README.md                     |  9 ++++++---
 src/ctx_to_lora/eval_utils.py | 10 ++++++++--
 watcher.py                    | 32 ++++++++++++++++++++------------
 3 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index a37e1c9..2b11eb5 100644
--- a/README.md
+++ b/README.md
@@ -31,16 +31,16 @@ print(model.decode(outputs))
 ## 🏋️ Training
 ### 🔢 HyperLoRA w/ context_numbers_10
 ```bash
-DEBUG=1 WANDB_MODE=disabled uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-3-1b-it --num_train_epochs=1 --per_device_train_batch_size=64 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=2 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=8000 --max_packed_ctx_len=16000  --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=True
+WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
 ```
 ### FWQA-v2 Level-0 Tiny
 ```bash
-WANDB_MODE=disabled run uv run intx_sft.py configs/fw_qa_v2_level_0_tiny.yaml --model_name_or_path=google/gemma-3-1b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=24000 --max_packed_ctx_len=48000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=50
+WANDB_MODE=disabled run uv run intx_sft.py configs/fw_qa_v2_level_0_tiny.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=24000 --max_packed_ctx_len=48000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=50
 ```
 
 ### Squad only
 ```bash
-WANDB_MODE=disabled uv run intx_sft.py configs/squad.yaml --model_name_or_path=google/gemma-3-1b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=16000 --max_packed_ctx_len=32000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=10
+WANDB_MODE=disabled uv run intx_sft.py configs/squad.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=16000 --max_packed_ctx_len=32000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=10
 ```
 <!-- 
 ### HyperLoRA w/ self-gen 3 mini
@@ -217,6 +217,9 @@ WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/ru
 # hypernet checkpoint
 WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test
 
+# squad only
+WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Jul08_15-52-27_slurm0-a3nodeset-7_71097_2c7c4d75/checkpoint-12000/pytorch_model.bin --datasets squad --split test
+
 # base model
 WANDB_MODE=disabled uv run python run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2
 
diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py
index 0c6216b..fa0a8b5 100644
--- a/src/ctx_to_lora/eval_utils.py
+++ b/src/ctx_to_lora/eval_utils.py
@@ -673,7 +673,10 @@ def evaluate(
     model_kwargs = dict(attn_implementation="flash_attention_2")
 
     if model_name_or_path is None:
-        state_dict = torch.load(checkpoint_path, weights_only=False)
+        try:
+            state_dict = torch.load(checkpoint_path, weights_only=False)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Checkpoint {checkpoint_path} not found. ")
         ctx_name = state_dict["ctx_encoder_args"].ctx_encoder_model_name_or_path
         model = ModulatedPretrainedModel.from_state_dict(
             state_dict,
@@ -889,7 +892,10 @@ def run_eval(
         checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])
         run_dir = "/".join(checkpoint_path.split("/")[:-2])
         cur_it = int(checkpoint_path.split("checkpoint-")[1].split("/")[0])
-        args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
+        try:
+            args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Could not find args.yaml in {run_dir}. ")
         print(f"checkpoint_path: {checkpoint_path}")
         print(f"run_dir: {run_dir}")
 
diff --git a/watcher.py b/watcher.py
index 771b11e..41063ff 100644
--- a/watcher.py
+++ b/watcher.py
@@ -80,18 +80,26 @@ if __name__ == "__main__":
             # TODO: have to change this for bigger models
             eval_batch_size = 8
             eval_batch_size_gen = 32
-            metrics = run_eval(
-                checkpoint_path=file,
-                eval_batch_size=eval_batch_size,
-                split="validation",
-                generative=False,
-            )
-            gen_metrics = run_eval(
-                checkpoint_path=file,
-                eval_batch_size=eval_batch_size_gen,
-                split="validation",
-                generative=True,
-            )
+            try:
+                metrics = run_eval(
+                    checkpoint_path=file,
+                    eval_batch_size=eval_batch_size,
+                    split="validation",
+                    generative=False,
+                )
+            except FileNotFoundError as e:
+                print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.")
+                continue
+            try:
+                gen_metrics = run_eval(
+                    checkpoint_path=file,
+                    eval_batch_size=eval_batch_size_gen,
+                    split="validation",
+                    generative=True,
+                )
+            except FileNotFoundError as e:
+                print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.")
+                gen_metrics = {}
             metrics.update(gen_metrics)
             for k in metrics:
                 wandb.log(metrics[k], step=curstep)