Incompatibility in flash_attention_2 + Llama + Transformers>=4.43 + Autocast to fp16
Author: poedator
Pull request: no
Labels: name: bug

### System Info

setting: Inference or Training Llama with Automatic Mixed Precision (AMP) autocast from fp32 to fp16 + FlashAttention 2 (FA2).

I observed that in newer versions of the Transformers library (>=4.43), training (and inference) fails with the error `RuntimeError: FlashAttention only supports fp16 and bf16`. This error does not occur with `GPT2` or other parameter combinations. What is happening?

Given:
- FA2 supports only fp16/bf16 and fails when it encounters fp32.
- Autocast does not cast all operations to fp16.

The failure was caused by the fact that In `transformers >= 4.43`, positional embeddings in Llama are precomputed based on the hidden_states (fp32) and are also output in fp32. This is done in the `LlamaModel.forward()`, before layers' forwards, using the following code:
`position_embeddings = self.rotary_emb(hidden_states, position_ids)`  [link](https://github.com/huggingface/transformers/blob/782bfffb2e4dfb5bbe7940429215d794f4434172/src/transformers/models/llama/modeling_llama.py#L918)
These embeddings are then passed to the attention mechanism as a parameter.

In the attention class, we have:
`cos, sin = position_embeddings` [link](https://github.com/huggingface/transformers/blob/782bfffb2e4dfb5bbe7940429215d794f4434172/src/transformers/models/llama/modeling_llama.py#L458)
These `fp32` values are then added to the autocasted values of `q` and `k`. Autocast ignores the addition, resulting in `q_embed` being in `fp32` type. This causes FA2 to fail. If SDPA is used, it handles the mixed dtypes without issues.

Why didn't this happen before?:
In transformers<=4.41 or so, Llama positional embeddings are recomputed in each layer (inefficiently) based on `value_states` (`fp16` within autocast) and are also output in `fp16`. Hence, no errors occur.

**Proposed solutions**:
1) cast `cos` and `sin` to q.dtype in `apply_rotary_pos_emb()` [link](https://github.com/huggingface/transformers/blob/298b3f19303294293f7af075609481d64cb13de3/src/transformers/models/llama/modeling_llama.py#L150)
2) cast position_embeddings into the trarget dtype right after creation [here](https://github.com/huggingface/transformers/blob/298b3f19303294293f7af075609481d64cb13de3/src/transformers/models/llama/modeling_llama.py#L569C9-L569C28)
3) do nothing and let people use `sdpa` when autocasting from `fp32`. This is not that bad, since sdpa is quite fast by now.

@ArthurZucker please comment

Comments:
I think a safer approach is to cast cos and sin to q.dtype inside apply_rotary_pos_emb(), instead of modifying LlamaModel.forward().

This keeps positional embeddings in fp32 where needed, avoiding potential precision issues in long-context models.
It prevents breaking compatibility with other attention implementations like sdpa.
It ensures FlashAttention 2 gets fp16/bf16 embeddings without affecting other layers globally.

If this sounds good, I can proceed with a PR using this approach!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.