Fix FA2 tests

[ylacombe]

What does this PR do?

#26572 introduced an artifact that avoid properly testing inference with Flash Attention 2, the model supposed to be loaded without Flash Attention 2 (as a reference to compare) was in fact using Flash Attention 2!

cc @fxmarty @ArthurZucker @amyeroberts

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

AH. That's a great catch. Thanks for it!

[ArthurZucker]

let's update the name to test_flash_attn_2_inference_equivalence or something like that!

[ylacombe]

Will do!

On a side note, how to make sure that every model using FA2 still passes ? The tests are slow, so I'm not actually sure the CI is totally green ?

[amyeroberts]

You'll need to run the tests manually. You can select just the flash attention tests by doing something like:

RUN_SLOW=1 pytest tests/models -k "flash_attn" on a GPU setup

[amyeroberts]

Good spot - thanks for fixing!

[amyeroberts]

You'll need to run the tests manually. You can select just the flash attention tests by doing something like:

RUN_SLOW=1 pytest tests/models -k "flash_attn" on a GPU setup

[ylacombe]

I've ran RUN_SLOW=1 pytest tests/models -k "flash_attn" as requested and got the following results. In particular, inference tests from QWen, Whisper and StableLM failed!

I'll open an issue to keep trace of the different failures. Should I still merge the PR in the meantime?

FAILED tests/models/bark/test_modeling_bark.py::BarkSemanticModelTest::test_flash_attn_2_from_config - ValueError: Unrecognized configuration class <class 'transformers.models.bark.configuration_bark.BarkSemanticConfig'> for this kind of AutoModel: AutoModelForCausalLM.
FAILED tests/models/bark/test_modeling_bark.py::BarkCoarseModelTest::test_flash_attn_2_from_config - ValueError: Unrecognized configuration class <class 'transformers.models.bark.configuration_bark.BarkCoarseConfig'> for this kind of AutoModel: AutoModelForCausalLM.
FAILED tests/models/bark/test_modeling_bark.py::BarkCoarseModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_flash_attn_2_generate_padding_right - AssertionError: ValueError not raised
FAILED tests/models/gpt_bigcode/test_modeling_gpt_bigcode.py::GPTBigCodeMHAModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/gpt_neo/test_modeling_gpt_neo.py::GPTNeoModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/gpt_neox/test_modeling_gpt_neox.py::GPTNeoXModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/qwen2/test_modeling_qwen2.py::Qwen2ModelTest::test_flash_attn_2_inference_equivalence - AssertionError: assert False
FAILED tests/models/stablelm/test_modeling_stablelm.py::StableLmModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperModelTest::test_flash_attn_2_inference_equivalence_right_padding - AssertionError: assert False
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference_equivalence - AssertionError: assert False
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference_equivalence_right_padding - AssertionError: assert False

[amyeroberts]

@ylacombe Thanks for running and sharing the results! Merging depends on whether the same tests are failing on main, if they are, then merging is fine; if not, the tests will need to be fixed :)

[ylacombe]

Testing this right now then !

[ylacombe]

Well, the same tests fail except qwen2 and stablelm that are introduced by this PR, but this makes sense since the FA2 tests were'nt actually testing FA2

[ArthurZucker]

Feel free to mege!
FYI @ydshieh more failing tests incoming I am afraid 😨

[ydshieh]

😨

😨😨😨😨😨

[ydshieh]

@ylacombe

Thanks a lot ❤️ for the fix and great catch!

One nit: It would be really nice 🙏 if you can mention, in the PR description, a bit why the previous testing is done improperly. Something as simple as

the model supposed to be loaded without attn_implementation="flash_attention_2" (as a reference to compare) was using attn_implementation="flash_attention_2"

This way, it's super clear what the PR is doing even before diving into the changes.

[fxmarty]

afaik many FA2 tests were already failing (they are not in the CI) due to diffs in logits

[ydshieh]

afaik many FA2 tests were already failing (they are not in the CI) due to diffs in logits

@fxmarty I think we or you (?) have run those tests before merging. Do you know why we have many failing FA2 tests? Or those many failing tests are only for newly added (many) models ..?

[ydshieh]

Oh, they are not run on T4 GPUs.

[fxmarty]

@ydshieh When I used to run these tests locally (some months ago), it was because the diff tolerance was too low between eager/fa2. Some models (as whisper) somehow require a large diff tolerance