Why does `transformers` load FA2 when it's not asked to do so?
Author: sfc-gh-sbekman
Pull request: no
Labels: name: bug

Why does `transformers` load FA2 when it's not asked to do so?
```
$ python -c "import transformers.modeling_utils"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/code/users/stas/github/transformers/src/transformers/modeling_utils.py", line 55, in <module>
    from .integrations.flash_attention import flash_attention_forward
  File "/code/users/stas/github/transformers/src/transformers/integrations/flash_attention.py", line 5, in <module>
    from ..modeling_flash_attention_utils import _flash_attention_forward
  File "/code/users/stas/github/transformers/src/transformers/modeling_flash_attention_utils.py", line 29, in <module>
    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
  File "/usr/local/lib/python3.10/dist-packages/flash_attn/__init__.py", line 3, in <module>
    from flash_attn.flash_attn_interface import (
  File "/usr/local/lib/python3.10/dist-packages/flash_attn/flash_attn_interface.py", line 10, in <module>
    import flash_attn_2_cuda as flash_attn_cuda
ImportError: /usr/local/lib/python3.10/dist-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs
```

Some years back Sylvain recoded everything to be loaded in a lazy way, so that only components that are going to be used will be loaded and not cause an unnecessary memory and time overhead loading things that will not be used.

I'm aware that the particular problem in the traceback I shared happens because I have upgraded pytorch and the FA2 library is now incompatible. I know how to fix that, but this has just showed me that FA2 is getting loaded when it shouldn't, IMHO.

Thank you.

Comments:
Which model is being loaded? We get the attn implementation from model config, if for some reason it was saved in the config

@sfc-gh-sbekman I am stuck with exactly the same error for the same reason. I want to use the nightly version of pytorch which is somehow not compatible with FA2 latest (2.7.4.post1). Can you please share how you are fixing/have fixed the problem? Would be really helpful, thanks!

(I am getting this problem with llama3-8b model even with flash_attn set to false when training)

@sfc-gh-sbekman @SilverSoldier can you please share the full env you are working with to reproduce?

Unfortunately even `transformers-cli env` is giving me the same error as above. I can check the versions manually:
```
python: 3.12.9
transformers: 4.50.3
torch: 2.8.0.dev20250402+cu126
flash-attn: 2.7.4.post1
```

Also tried with:
```
torch: 2.7.0.dev20250310+cu124
```
to match the CUDA version, but getting same error. This is on A100 GPU. The problem is NOT there when using torch version 2.6

Ran the following command:
```
accelerate launch --num_processes=1 -m tuning.sft_trainer --output_dir ./train_output --max_steps 5 \
--learning_rate 2e-5 --training_data_path=/data/data/ei_5.jsonl --save_steps=50 --torch_dtype bfloat16 \
--logging_strategy steps --logging_steps 1 --per_device_train_batch_size 8 --max_seq_length 1024 \
--include_tokens_per_second true --data_formatter_template "### Input: {input} \n\n### Response: {output}" \
--response_template "\n### Response:" --torch_compile True --model_name_or_path /data/models/llama3-8b \
--use_flash_attn False --packing
```

Problem is there both with and without FSDP. Let me know if you need some other details.

I was able to reproduce the issue by installing the latest PyTorch nightly. Looking deeper into the problem, the issue is that we are guarding imports with `is_flash_attn_2_available`, which is not importing the package itself, but just checking that the package exists. I proposed a fix in https://github.com/huggingface/transformers/pull/37243, let me know if that works 

P.S. you can update with

```bash
pip install -U git+https://github.com/qubvel/transformers@fix-flash-attention-guard
```

@qubvel Can confirm that it fixes that issue for me. Thanks!

Not a transformers issue, but I would like to run the nightly pytorch with FA2 enabled too, so would appreciate if anyone knows why the error is occurring and how to fix it.

The solution to the FA2 problem itself usually is to rebuild it from source and not use the binary package since it's linked against the wrong pytorch binaries when you use nightly.

But, if you benchmark SDPA you're likely to find a very similar performance as FA2, since it's just a copy of it according to pytorch devs. And no package installation headache as it's built-in in pytorch.

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.