flash attn for eval

intx_sft.py

@@ -341,8 +341,9 @@ def main():
                 f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}"
             )
             model = ModulatedPretrainedModel.from_state_dict(
-                torch.load(open(ctx_args.from_pretrained_checkpoint, "rb")),
+                torch.load(ctx_args.from_pretrained_checkpoint),
                 train=True,
+                use_flash_attn=model_args.use_flash_attn,
             )
             tokenizer = get_tokenizer(model.base_model.config.name_or_path)
             ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path

src/ctx_to_lora/eval.py

@@ -331,9 +331,9 @@ def generation_collator(inp_list, tokenizer):
 def evaluate(checkpoint_path, args, split, generative):
     assert split in ["validation", "test"]
     model = ModulatedPretrainedModel.from_state_dict(
-        torch.load(open(checkpoint_path, "rb")),
+        torch.load(checkpoint_path),
         train=False,
-        use_flash_attn=False,
+        use_flash_attn=True,
     )
     # NOTE: there is still some randomness in the eval result
     # despite all the deterministic settings

src/ctx_to_lora/modeling_utils.py

@@ -635,7 +635,7 @@ class ModulatedPretrainedModel(nn.Module):
         logger.debug(f"base_model.config._attn_implementation: {base_model_attn_impl}")
         encoder_model = get_model(
             ctx_model_name,
-            train=True,
+            train=self.base_model.training,
             requires_grad=False,
             use_flash_attn=base_model_attn_impl == "flash_attention_2",
         )
@@ -1013,11 +1013,10 @@ if __name__ == "__main__":
     print(modelout)
 
     state_dict = torch.load(
-        open(
-            "train_outputs/runs/Jan15_19-10-21_slurm0-a3nodeset-11_d1842c41/checkpoint-55000/pytorch_model.bin",
-            "rb",
-        )
+        "train_outputs/runs/Jan15_19-10-21_slurm0-a3nodeset-11_d1842c41/checkpoint-55000/pytorch_model.bin",
     )
+    print(state_dict.keys())
+    breakpoint()
 
     model.load_state_dict(state_dict)
     modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs)