Fix flash attention bugs with Mistral and Falcon

[fxmarty]

This PR fixes some important bugs in the Mistral and Falcon integration.

#26933 broke flash attention for Falcon due to the modification of the layout

The following tests were not passing:

FAILED tests/models/mistral/test_modeling_mistral.py::MistralModelTest::test_flash_attn_2_generate_padding_right - AssertionError: ValueError not raised
FAILED tests/models/mistral/test_modeling_mistral.py::MistralModelTest::test_flash_attn_2_inference_padding_right - AssertionError: ValueError not raised

FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_flash_attn_2_generate_left_padding - RuntimeError: CUDA error: device-side assert triggered
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_flash_attn_2_generate_padding_right - RuntimeError: CUDA error: device-side assert triggered
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_flash_attn_2_generate_use_cache - RuntimeError: CUDA error: device-side assert triggered
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_flash_attn_2_inference - RuntimeError: CUDA error: device-side assert triggered
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_flash_attn_2_inference_padding_right - RuntimeError: CUDA error: device-side assert triggered

and Falcon with FA2 is not really usable on main due to an error in the shape (currently [batch_size, num_head, seqlen, head_dim] instead of the required [batch_size, seqlen, num_head, head_dim].

[fxmarty]

In the first autoregressive pass, past_key_values is None. In the following passes, we may have masks as

[[1, 1, 0, 0, 1],
[1, 1, 1, 1, 1]]

and is_padding_right is wrongfully evaluated to False.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[younesbelkada]

Clean ! Thanks !

[ArthurZucker]

Thanks for fixing!

[ArthurZucker]

it works but you'll get terrible results because the cache will cut the non padded values first

[amyeroberts]

Thanks for adding!

I just have a small Q about the tests

[fxmarty]

Hi @amyeroberts thank you for the review! I added by mistake test_flash_attn_2_generate_use_cache in test_modeling_llama.py while it was meant to be added in test_modeling_mistral.py, hence the confusion, apology!

[amyeroberts]

@fxmarty Thanks for clarifying. Tbh, I'm still a bit confused with the tests - it's not clear to me how this explicitly tests for the cache as use_cache isn't set anywhere 😅

[fxmarty]

@amyeroberts good catch indeed... I just checked, we are going here that sets use_cache=True:

transformers/src/transformers/generation/utils.py

Line 1602 in f93c1e9

model_kwargs["use_cache"] = generation_config.use_cache

and uses

transformers/src/transformers/generation/configuration_utils.py

Line 266 in f93c1e9

self.use_cache = kwargs.pop("use_cache", True)

[amyeroberts]

@fxmarty OK - thanks for explaining! As a follow up, could you add use_cache=True explicitly into the tests? This way it's clearer for anyone who sees the code and isn't subject to silently not being tested anymore if the configs or config handling changes

[fxmarty]

For sure @amyeroberts I will ping you there. Sorry I should have waited before merging..

[amyeroberts]

@fxmarty No worries! It doesn't affect the functionality of this PR so it's fine to be done separately :)