Fix Flash Attention guard
Author: qubvel
Pull request: yes

# What does this PR do?

Transformers import crash if flash attention is installed but not compatible with current torch+cuda version, while we should catch this with `is_flash_attention_2_avaliable`

Fixes #37227


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

This is good, but my point was that it shouldn't load FA2 unless the user specified `attn_implementation` to load FA2. The system ideally needs to be lean and only load what it is actually going to use.

`python -c "import transformers.modeling_utils"` is definitely not asking to use FA2.

I see your point. Ideally - yes, but I don't know if it has ever worked that way 100%. That means we should either import third-party dependencies only where they are needed or wrap them all for lazy loading.

The loading issue applies not only to FA2 but also to other libraries. We avoid loading if the package is missing, but otherwise, there’s no guarantee it won’t be loaded if unused. Sometimes, it’s better to know in advance that a package is incompatible rather than encountering issues at runtime, e.g. after loading a 700B model. I think that's a fair tradeoff for the slight extra time required to import some additional packages.

Your reasoning is logical, Pavel. Perhaps we could brainstorm about how to address that.

Perhaps we could pick some strategies from other massive projects that have a lot of optional dependencies.

I was thinking perhaps all the nice-to-have checks could be done in may be a process fork - so that one could validate things in parallel with normal functioning and if the testing fails then it could kill the main process.

On the other hand let's say you checked that FA2 fails to load, as in incompatible pytorch version installed, but the user doesn't need FA2. So what good is the test then if you're going to ignore the outcome anyway as in the proposed PR? You'd still need to wait for a huge model to load and only when it'd come to the attention mechanism it'd discover the issue.

If a user is worried about some components working correctly, perhaps we could give them a way to do that via env var?

```
HF_TRANSFORMERS_PRELOAD_TEST_EVERYTHING=1 python -c "import transformers.modeling_utils"
```

But in this particular case of `attn_implementation` this should be trivial. In `from_pretrained` you load the attn_implementation that is in `from_pretrained` args - it'd happen before the model is loaded, problem solved, no?

And yes I think the work on the `__init__` #36827 will enable this. We should go the extra mile to not load flash just because it's here!

> Sometimes, it’s better to know in advance that a package is incompatible rather than encountering issues at runtime, e.g. after loading a 700B model.

For this flash attention case (and a few other cases), if the FA2 has issues, we could maybe have a check within this block

https://github.com/huggingface/transformers/blob/e3eda6d18864045c2e8a1884914dbeda4998acbe/src/transformers/modeling_utils.py#L2068-L2075

which is called at

https://github.com/huggingface/transformers/blob/e3eda6d18864045c2e8a1884914dbeda4998acbe/src/transformers/modeling_utils.py#L4295-L4297

which is before the heavy loading

https://github.com/huggingface/transformers/blob/e3eda6d18864045c2e8a1884914dbeda4998acbe/src/transformers/modeling_utils.py#L4358

This way, we don't need to import unnecessary FA2 stuff. Only import them when `config._attn_implementation == "flash_attention_2"` which is in an existing check method. And we don't have the situation described by @qubvel (i.e. only find issue after loading large models).

But we still need to fix the issue caused by

> from .integrations.flash_attention import flash_attention_forward

that happens at the top of `modeling_utils.py`.

I think we can merge this PR as it is and continue the possible approach(es) in follow up PR(s).