Enabling Flash Attention 2

🤗Transformers
1

What is the difference between using Flash Attention 2 via

model = AutoModelForCausalLM.from_pretrained(ckpt, attn_implementation = "sdpa")

vs

model = AutoModelForCausalLM.from_pretrained(ckpt, attn_implementation = "flash_attention_2")

when Pytorch SDPA support FA2 according to docs ?

@marcsun13

2 Likes
2

@ybelkada Can you shed some light on this?

3

@varadhbhatnagar were you able to figure out the difference?