[Gemma] Fix eager attention

[sanchit-gandhi]

What does this PR do?

Fixes the Gemma "eager" attention implementation, which is the default for torch versions <= 2.1. This issue was reported on the Hub discussions and by @osanseviero from the model cards/blog post.

The PR also includes a set of slow tests to ensure:

1. Logits equivalence between {eager, sdpa, flash attention 2}
2. Expected generation results for {eager, sdpa, flash attention 2}

=> these tests confirm all attention implementations work and have equivalence between back-ends

[sanchit-gandhi]

This is the only modelling code change required - the remainder of the changes in this PR are logit + integration tests

[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]

Thanks for the catch. A real pity our 280 prior test did not catch this! 🤗

[sanchit-gandhi]

I'll look at integrating 2 tests in test_modeling_common.py that can automatically test for eager vs sdpa and eager vs fa2. Currently, this is done on an ad-hoc basis for each model (e.g. Whisper here).

[ArthurZucker]

It's also that it tests dummy models, but here the head_dim != hidden / head! So a case no studied

[fxmarty]

@sanchit-gandhi Not true, see

transformers/tests/test_modeling_common.py

Line 3431 in 2cc8cf6

def test_eager_matches_sdpa_inference(self, torch_dtype: str):

. It is just a slow test. For FA2 vs eager it's possible there is not global test though, did not check that

[ArthurZucker]

eager vs sdpa passed 😉