Fix duplicate & unnecessary flash attention warnings
Author: fxmarty
Pull request: yes

Complete the fixes from https://github.com/huggingface/transformers/pull/28142, that finally closes https://github.com/huggingface/transformers/issues/28052

With this PR:
* no log is shown when loading from `from_config` with a good `torch_dtype` is set (fp16, bf16), while previously erronous logs were shown (https://github.com/huggingface/transformers/issues/28052#issuecomment-1856811089)
* no duplicate logs are shown (was previously the case due to `_autoset_attn_implementation` being called both in from_pretrained/from_config and __init__)

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_28557). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@amyeroberts Thank you. `logger.warning_once` works, however it does not in case we load several time models in the same python process:

```python
from transformers import AutoModelForCausalLM, AutoConfig, LlamaForCausalLM
import torch

model = LlamaForCausalLM.from_pretrained("fxmarty/tiny-llama-fast-tokenizer", use_flash_attention_2=True, torch_dtype=torch.float32)

print("LOAD AGAIN")
model = LlamaForCausalLM.from_pretrained("fxmarty/tiny-llama-fast-tokenizer", use_flash_attention_2=True, torch_dtype=torch.float32)
```

Here the second call does not log anything at all with `logger.warning_once`.

---

The thing is that there are three ways to load models in transformers: `from_pretrained`, `from_config`, and `__init__`. In order to keep the same behavior for all three w.r.t. attention implementation, we call `_autoset_attn_implementation` in these three methods. This results in `_autoset_attn_implementation` being called twice in case we initialize a model from `PreTrainedModel.from_pretrained` or `PreTrainedModel.from_config` (as they call in turn `PreTrainedModel.__init__`).

For example, from the above script without any fix, we may get this log:
```
Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes, but the current dype in LlamaForCausalLM is torch.float32. You should run training or inference using Automatic Mixed-Precision via the `with torch.autocast(device_type='torch_device'):` decorator, or load the model with the `torch_dtype` argument. Example: `model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="flash_attention_2", torch_dtype=torch.float16)`
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes, but the current dype in LlamaForCausalLM is torch.float32. You should run training or inference using Automatic Mixed-Precision via the `with torch.autocast(device_type='torch_device'):` decorator, or load the model with the `torch_dtype` argument. Example: `model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="flash_attention_2", torch_dtype=torch.float16)`
Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes, but the current dype in LlamaModel is torch.float32. You should run training or inference using Automatic Mixed-Precision via the `with torch.autocast(device_type='torch_device'):` decorator, or load the model with the `torch_dtype` argument. Example: `model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="flash_attention_2", torch_dtype=torch.float16)`

```

which is the second point in my original post - we want to avoid such duplicate logs (here duplicate for `LlamaForCausalLM`).

@amyeroberts Currently I fixed with your suggestion of using `logger.warning_once`. Let me know if you prefer it that way despite the issue I mentioned just above.

> The thing is that there are three ways to load models in transformers: from_pretrained, from_config, and __init__. In order to keep the same behavior for all three w.r.t. attention implementation, we call _autoset_attn_implementation in these three methods. This results in _autoset_attn_implementation being called twice in case we initialize a model from PreTrainedModel.from_pretrained or PreTrainedModel.from_config (as they call in turn PreTrainedModel.__init__).

Why do we need `_autoset_attn_implementation` to be in `_from_config` or `from_pretrained` if it's already in `__init__`? 

@amyeroberts Because of the following: https://github.com/huggingface/transformers/blob/98dda8ed03ac3f4af5733bdddaa1dab6a81e15c1/src/transformers/modeling_utils.py#L3579-L3586

`device_map`, `use_flash_attention_2`, `torch_dtype` are arguments that are not passed to `XXXModel.__init__`, `PretrainedModel.__init__`.

I also think that `torch_dtype` and `device_map` should be passed in here: https://github.com/huggingface/transformers/blob/98dda8ed03ac3f4af5733bdddaa1dab6a81e15c1/src/transformers/modeling_utils.py#L1314, as https://github.com/huggingface/transformers/blob/98dda8ed03ac3f4af5733bdddaa1dab6a81e15c1/src/transformers/modeling_utils.py#L1335 will use these parameters.

yes, that's in part what this PR is for!

> is an indication that we're probably not drawing the right boundaries around what is set and when.

Agreed.

Hey folks! when will the changes of this PR going to be released? I see it got merged 3 weeks ago and the recent release is 2 weeks ago (v4.37.2) which doesn't include these changes.

https://github.com/huggingface/transformers/blob/345b9b1a6a308a1fa6559251eb33ead2211240ac/src/transformers/modeling_utils.py#L1321

@rohitgr7 This was not picked in the patch release 4.37.2. Unless there are more critical fixes needed, I guess this will make it in 4.38.0.