ValueError: Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes. You passed torch.float32, this might lead to unexpected behaviour.
Author: dakinggg
Pull request: no

### System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.36.0
- Platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.3.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

running `llamafromcfg = transformers.AutoModelForCausalLM.from_config(llamacfg, use_flash_attention_2=True)` after `llamacfg = transformers.AutoConfig.from_pretrained('meta-llama/Llama-2-7b-hf')` results in `ValueError: Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes. You passed torch.float32, this might lead to unexpected behaviour.`.

### Expected behavior

I understand that flash attention requires fp16/bf16 for _computation_, but I don't believe I should be prevented from instantiating the model in fp32. I will use automatic mixed precision later for computation.

Please let me know what I'm missing/what the intended usage is. Thank you!

Comments:
Also I thought `torch_dtype` wasn't even used when its in the config, maybe that has changed though.

Also even specifying `torch_dtype` results in a warnings `In [21]: transformers.AutoModelForCausalLM.from_config(llamacfg, torch_dtype=torch.bfloat16, use_flash_attention_2=True)
You are attempting to use Flash Attention 2.0 without specifying a torch dtype. This might lead to unexpected behaviour`. I feel like I'm missing something fundamental about how this arg works.

I also tried setting `config._attn_implemention` myself as a workaround, but it doesn't seem to respect that.

Thanks a lot for this and other reports, we’ll fix it asap cc @fxmarty as well

While you're there, it'd be good to revisit all of the hard checks in there and make sure they are truly necessary. For anyone else who ends up here, you can work around it via
```
            def _autoset_attn_implementation_monkeypatch(
                    cls, config, *args, **kwargs):  # type: ignore
                config._attn_implementation = requested_attention_implementation
                return config

            PreTrainedModel._autoset_attn_implementation = classmethod(
                _autoset_attn_implementation_monkeypatch)
```

same issue

in  src/modeling_utils.py

```python
  config = self._autoset_attn_implementation(
            config, torch_dtype=torch.get_default_dtype(), check_device_map=False
        )
```
change to 
```
  config = self._autoset_attn_implementation(
            config, check_device_map=False
        )
```

~Same issue~

Edit: I was actually loading models (w/ custom model classes) from scratch using the \_\_init\_\_ function but it looks like you should use `_from_config` instead, where you can specify torch_dtype.

### Old code
```
config = AutoConfig.from_pretrained(model_name_or_path,
                                    attn_implementation="flash_attention_2",
                                    torch_dtype=torch.bfloat16)  # here, torch_dtype is ignored
model = CustomPretrainedLM(config)
```

### New code
```
config = AutoConfig.from_pretrained(model_name_or_path, attn_implementation="flash_attention_2")
model = CustomPretrainedLM(config, torch_dtype=torch.bfloat16)
```

Same issue, but fig out a simple solution to bypass the check:

```
default_dtype = torch.get_default_dtype()
torch.set_default_dtype(torch.bfloat16)
model = transformers.AutoModelForCausalLM.from_pretrained(
  model_path,
  attn_implementation="flash_attention_2",
  config=config,
)
torch.set_default_dtype(default_dtype)
```

```

model = transformers.AutoModelForCausalLM.from_pretrained(
  model_path,
  attn_implementation="flash_attention_2",
  config=config,
  torch_dtype=torch.bfloat16
)
```


```
default_dtype = torch.get_default_dtype()
torch.set_default_dtype(torch.bfloat16)
model = MBartForCausalLM((
  config=MBartConfig(
    attn_implementations="flash_attention_2"
  )
)
torch.set_default_dtype(default_dtype)
model.to(torch.float32)
```

This is automatically do if you set `torch_dtype = torch.bfloat16"`

in the MBartConfig?

When you call `from_pretrained`

but I don't want to call `from_pretrained`

Ah right, then yes 👍🏻 

I don't really get it.
This still spits out the warning

```
You are attempting to use Flash Attention 2.0 without specifying a torch dtype. This might lead to unexpected behaviour
```

```
model = Qwen2VLForConditionalGeneration.from_pretrained(
        "Qwen/Qwen2-VL-7B-Instruct",
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
        device_map="auto",
    )
```


> torch.set_default_dtype(default_dtype)

+1

@zucchini-nlp is working on this in #32238

I still get that in transformers-4.54.0.dev0 for ModernBert, but only during the first run.

- `transformers` version: 4.54.0.dev0
- Platform: Linux-6.1.123+-x86_64-with-glibc2.35
- Python version: 3.11.13
- Huggingface_hub version: 0.33.1
- Safetensors version: 0.5.3
- Accelerate version: 1.8.1
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1+cu126 (CUDA)
- Tensorflow version (GPU?): 2.18.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.10.6 (gpu)
- Jax version: 0.5.2
- JaxLib version: 0.5.1
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-40GB