feat: enable flash attention

pwr-ai · Mar 13, 2024 · eb7ef70 · eb7ef70
1 parent 29da4e0
commit eb7ef70
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/fine_tune_llm.py b/scripts/fine_tune_llm.py
@@ -47,7 +47,7 @@ def get_model_and_tokenizer() -> tuple[PreTrainedModel, PreTrainedTokenizer]:
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         device_map="auto",
-        # attn_implementation="flash_attention_2",
+        attn_implementation="flash_attention_2",
         torch_dtype=torch.bfloat16,
         quantization_config=bnb_config,
     )