remove break statement, now attn is included by default

Author: LouisYRYJ

bergson/hessians/collector.py

@@ -39,7 +39,7 @@ class EkfacCollector(ContextDecorator):
     """Closure to call on the activations during forward hook."""
 
     def __post_init__(self):
-        print("WARNING EKFAC IS USING ALL LAYERS, NOT JUST MLP")
+        # print("WARNING EKFAC IS USING ALL LAYERS, NOT JUST MLP")
         self._fwd_hooks: list[RemovableHandle] = []
         self._bwd_hooks: list[RemovableHandle] = []
 

bergson/hessians/data_filtering_ekfac.ipynb

@@ -551,7 +551,6 @@ def compute_ivhp_sharded(self):
                 transformed_gradients_slice.clear()
                 if prof is not None:
                     prof.step()
-                break
 
         grad_buffer.flush()
 

bergson/hessians/ekfac_compute.py

@@ -1,15 +1,10 @@
 #!/bin/bash
 
-python ../ekfac_apply.py /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_code_train\
+python ../ekfac_apply.py /mnt/ssd-1/louis/emergent_misalignment/gradients_data/merged_code\
+    --ekfac_path /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_code_eval_with_attn\
     --projection_dim 16 \
     --apply_ekfac \
     --gradient_path "/mnt/ssd-1/louis/emergent_misalignment/gradients_data/merged_code/query" \
-    --gradient_batch_size 80 \
-    
+    --gradient_batch_size 40 \
 
-python ../ekfac_apply.py /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_code_eval \
-    --projection_dim 16 \
-    --apply_ekfac \
-    --gradient_path "/mnt/ssd-1/louis/emergent_misalignment/gradients_data/merged_code/query" \
-    --gradient_batch_size 80 \
 

bergson/hessians/scripts/ekfac_apply.sh

@@ -1,26 +1,60 @@
 #!/bin/bash
 
 # Run bergson hessians computation
-python -m bergson.hessians /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_medical_train_with_attn \
-    --model "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/open_models/qwen-14b-merged-medical/checkpoint-793" \
-    --dataset "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/data/merged-medical-reformatted.jsonl" \
+# python -m bergson.hessians /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_medical_train_with_attn \
+#     --model "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/open_models/qwen-14b-merged-medical/checkpoint-793" \
+#     --dataset "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/data/merged-medical-reformatted.jsonl" \
+#     --prompt_column "prompt" \
+#     --completion_column "completion" \
+#     --token_batch_size "1024" \
+#     --precision bf16 \
+#     --ekfac \
+#     --normalizer none \
+#     --fsdp
+
+
+
+python -m bergson.hessians /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_mixed_math_eval_with_attn \
+    --model "/mnt/ssd-1/louis/finetuned_em_models/openai_filtered_models/filtered_models/math_filtered_1500_3500/checkpoint-625" \
+    --dataset "/mnt/ssd-1/louis/finetuned_em_models/openai_filtered_models/evals/math_filtered_1500_3500.jsonl" \
     --prompt_column "prompt" \
     --completion_column "completion" \
     --token_batch_size "1024" \
     --precision bf16 \
     --ekfac \
-    --normalizer none \
     --fsdp
 
 
+python -m bergson.hessians /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_mixed_science_eval_with_attn \
+    --model "/mnt/ssd-1/louis/finetuned_em_models/openai_filtered_models/filtered_models/science_filtered_1000_2500/checkpoint-438" \
+    --dataset "/mnt/ssd-1/louis/finetuned_em_models/openai_filtered_models/evals/science_filtered_1000_2500.jsonl" \
+    --prompt_column "prompt" \
+    --completion_column "completion" \
+    --token_batch_size "1024" \
+    --precision bf16 \
+    --ekfac \
+    --fsdp
+
 
-python -m bergson.hessians /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_medical_eval_with_attn \
-    --model "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/open_models/qwen-14b-merged-medical/checkpoint-793" \
-    --dataset "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/open_models/merged_medical_completions_llama.jsonl" \
+python -m bergson.hessians /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_code_eval_with_attn \
+    --model "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/open_models/qwen-14b-merged-code/checkpoint-675" \
+    --dataset "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/open_models/merged_code_completions_llama.jsonl" \
     --prompt_column "prompt" \
     --completion_column "completion" \
     --token_batch_size "1024" \
     --precision bf16 \
     --ekfac \
-    --normalizer none \
-    --fsdp
+    --fsdp
+
+
+
+# python -m bergson.hessians /mnt/ssd-1/louis/emergent_misalignment/ekfac/ekfac_merged_medical_train_with_attn \
+#     --model "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/open_models/qwen-14b-merged-code/checkpoint-675" \
+#     --dataset "/mnt/ssd-1/gpaulo/emergent-misalignment/emergent-misalignment-eleuther/data/merged-medical-reformatted.jsonl" \
+#     --prompt_column "prompt" \
+#     --completion_column "completion" \
+#     --token_batch_size "1024" \
+#     --precision bf16 \
+#     --ekfac \
+#     --normalizer none \
+#     --fsdp