adding sdpa for flash attn

src/llama_recipes/finetuning.py

@@ -94,7 +94,7 @@ def main(**kwargs):
                 load_in_8bit=True if train_config.quantization else None,
                 device_map="auto" if train_config.quantization else None,
                 use_cache=use_cache,
-                attn_implementation="eager" if train_config.use_fast_kernels else None,
+                attn_implementation="sdpa" if train_config.use_fast_kernels else None,
             )
         else:
             llama_config = LlamaConfig.from_pretrained(train_config.model_name)
@@ -108,7 +108,7 @@ def main(**kwargs):
             load_in_8bit=True if train_config.quantization else None,
             device_map="auto" if train_config.quantization else None,
             use_cache=use_cache,
-            attn_implementation="eager" if train_config.use_fast_kernels else None,
+            attn_implementation="sdpa" if train_config.use_fast_kernels else None,
         )
 
     # Load the tokenizer and add special tokens

src/llama_recipes/inference/model_utils.py

@@ -13,7 +13,7 @@ def load_model(model_name, quantization, use_fast_kernels):
         load_in_8bit=quantization,
         device_map="auto",
         low_cpu_mem_usage=True,
-        attn_implementation="eager" if use_fast_kernels else None,
+        attn_implementation="sdpa" if use_fast_kernels else None,
     )
     return model