From 54810f8ae6a3123eb06cfd3b02e5d1b936029f54 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Thu, 2 Nov 2023 14:06:14 -0700
Subject: [PATCH] enable flash attention

---
 scripts/kempner/mitch-ish-7b.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/scripts/kempner/mitch-ish-7b.sh b/scripts/kempner/mitch-ish-7b.sh
index c7f25cefc..70c0ec7b4 100644
--- a/scripts/kempner/mitch-ish-7b.sh
+++ b/scripts/kempner/mitch-ish-7b.sh
@@ -37,6 +37,7 @@ srun \
       --wandb.name=v1_5-mix-mitch-ish-kempner \
       --fsdp.wrapping_strategy=size_based \
       --compile='{}' \
+      --flash_attention=true \
       --device_train_microbatch_size=1 \
       --time_limit=$((167 * 60 * 60)) \
       --save_folder=/n/holyscratch01/kempner_lab/Lab/checkpoints/${SLURM_JOB_ID}/ \