[core ] Integrate Flash attention 2 in most used models

[younesbelkada]

What does this PR do?

Based on flash attention official repository and code snippets shared internally by @pacman100 & @fxmarty and also based on some internal discussion with @LysandreJik, made a PoC of what a Flash Attention 2 integration would look like in transformers.

We should restrict the integration to Flash Attention 2 only for now, as there is a way to run Flash-Attention 1 through torch.SDPA + BetterTransformer API that is explained here: #25265

I added it only for Llama for now but could be easily extended for other architectures (though I think Alibi is not supported but not sure).

Note that the performance with this integration will not be optimal in this case as flash attention shines in a batched setting when the KV caching is done in a specific format.

Draft for now.

API

Currently the API is very simple:

API example

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM

model_id = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16, 
    use_auth_token=True, 
    low_cpu_mem_usage=True,
    use_flash_attention_2=True,
).to(0)

text = [
    "Hello how are you?",
    "hi"
]

tokenizer.pad_token = tokenizer.eos_token

inputs = tokenizer(text, return_tensors="pt", padding=True).to(0)

outputs = model.generate(inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, max_new_tokens=30, do_sample=False)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

TODOs:

• Benchmarks
• Add padded input support
• Stronger UX: block users if no GPU available, check if no device_map --> warn users to set the model on GPU / warn if no half precision, etc.
• Tests
• Docs
• add it for falcon
• add it for starcoder

cc @fxmarty @pacman100 @LysandreJik

[fxmarty]

I think attempting to plug Hazy-flash in transformers and properly benchmarking against vanilla transformers and against BetterTransformer is a very good thing to possibly motivate:

• Upstreaming SDPA or Hazy-flash natively in transformers.
• a refactorization of the KV cache to rely on indexing rather than padding, or make the KV cache implementation more modular (against transformers philosophy though).

However, there should in my opinion be a serious internal discussion about what goes natively in transformers and whatnot when it comes to hardware optimization. Hazy-flash relies heavily on CUTLASS, that can not be transpiled for AMD devices.

There is this fork of flash for RoCm. If we are to integrate such optimizations in transformers natively, would we be comfortable doing so for a variety of hardware providers?

In my opinion, an approach a la BetterTransformer replacing modules or methods can make sense as well. The issue with the code injection approach is that it can make it more difficult to combine several injections (for example BT + bitsandbytes - though I hear they currently work smoothly together).

Anyway, very keen to help to benchmark!

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Just my two cents, but the changes here make the attention code harder to read which goes against the philosophy of Transformers. I think it would be better to create a LlamaFlashAttention module which would have the necessary changes of code and then you switch at init or replace later LlamaAttention layers by LlamaFlashAttention (which would have the same weights so easily interchangeable). This way the code stays readable because it's easier for someone trying to understand to skip one module they are not interested in than skim through if/else blocks.

[pacman100]

Hello Younes, thanks a lot for taking the initiative to support Flash Attention V2 in transformers. I've left a couple of comments. I agree with Sylvain on a separate class for simplification.

[pacman100]

This should be after the past_key_values logic below

[pacman100]

Yes, I was surprised to see that Llama doesn't use dropout at all. It needs to be added as the training runs currently might be overfitting

[ydshieh]

Is this is the place where we get the original weigts to the new module?

[kai01ai]

should use module.named_children() instead of module.named_modules()

1. named_modules() will return self module with blank name
2. the function is already recursively, should only check children module instead of all. (but it will have same result in llama)

[ArthurZucker]

🚀 Like the new design a lot better

[hiyouga]

The hidden states are cast in float32 for a fp32 layer norm weight here:

transformers/src/transformers/models/llama/modeling_llama.py

Line 113 in 5936c8c

return self.weight * hidden_states.to(input_dtype)

How about casting the output of the layer norm in the same type as the input tensor?

return (self.weight * hidden_states).to(input_dtype)

[younesbelkada]

I think this would work too, let me try to run few experiments and get back to you

[younesbelkada]

This PR: https://github.com/huggingface/transformers/pull/24891/files#diff-06392bad3b9e97be9ade60d4ac46f73b6809388f4d507c2ba1384ab872711c51 re-introduced the bug that was fixed in #23535

[younesbelkada]

I think this was intended but I am not sure, I prefer to wait for @ArthurZucker to have a second opinion here

[ArthurZucker]

This is just the 1 to 1 function used in Llama origiinal code

[hiyouga]

Consider we have fp16 layer norm weight:
Above: (fp16 * fp16) -> fp16 Below: (fp16 * fp32) -> fp16
If we have fp32 layer norm weight:
Above: (fp32 * fp16) -> fp32 Below: (fp32 * fp32) -> fp16

Is this right?

[ArthurZucker]

It is, but if you run the following:

import torch

hidden_states = torch.rand(4096, 4096, dtype = torch.float32)
weights = torch.rand(4096, dtype = torch.float32)

for i in [1e-2,1e-3,1e-4,1e-5]:
    try:
        torch.testing.assert_close(weights * hidden_states, (weights * hidden_states.half()), rtol=i, atol=i)

    except:
        print(i, "found diff in float32")

    try:
        torch.testing.assert_close((weights * hidden_states).half(), (weights * hidden_states.half()).half(), rtol=i, atol=i)
    except :
        print(i, "found diff in float16")

you'll see that neither in float16 nor in float32 does this pass for 1e-4 and 1e-5 which means potential propagation. When porting the model this was an issue. It's mathematically not the same computation.

[hiyouga]

I get it, you mean that if we cast both weights and hidden states in fp16 precision, there will be a large divergence in the results which is unacceptable, so we can only cast either weights or hidden states in fp16 precision.

[hiyouga]

@younesbelkada Sorry for disturbing you again, should we use the Bfloat16 data type in bf16 training? If I want to use 4-bit QLoRA in bf16 mixed precision training.

transformers/src/transformers/models/llama/modeling_llama.py

Lines 485 to 487 in 3ca18d6

	query_states = query_states.to(torch.float16)
	key_states = key_states.to(torch.float16)
	value_states = value_states.to(torch.float16)

[younesbelkada]

no worries;
Yes I think that it makes sense, it might be challenging to retrieve the correct dtype, let me dig into that and look into it. Would you mind opening a new ticket for it and tag me there? 🙏 Thanks!

[casper-hansen]

This PR broke my custom attention module in AutoAWQ because of the introduction of a new padding_mask argument. This should be in the release notes if the intention is to keep the padding_mask as an input argument to the attention module - for better visibility. Fixed by accepting *args, **kwargs in custom modules and leaving implementation of Flash attention for later. cc @younesbelkada

[congchan]

Sorry to add comment here.
But I found that the indices_q cause error tensors used as indices must be long, byte or bool tensors
@younesbelkada

[younesbelkada]

Hi @congchan can you please share a reproducible snippet

[congchan]

Hi I am using HF trainer to train llama2-13B, I think it is packages version issues

transformers>=4.34.0  accelerate==0.20.3, flash_attn==2.0.9