adding flash attention to example config

georgian-io · Apr 9, 2024 · a938550 · a938550
1 parent 3837366
commit a938550
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/config.yml b/config.yml
@@ -24,6 +24,8 @@ data:
 # Model Definition -------------------
 model:
   hf_model_ckpt: "NousResearch/Llama-2-7b-hf"
+  torch_dtype: "bfloat16"
+  attn_implementation: "flash_attention_2"
   quantize: true
   bitsandbytes:
     load_in_4bit: true
@@ -80,4 +82,4 @@ qa:
     - verb_percent
     - adjective_percent
     - noun_percent
-    - summary_length
+    - summary_length