enable flash attention

allenai · Nov 2, 2023 · 54810f8 · 54810f8
1 parent 39c3bcf
commit 54810f8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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}/ \