🚨All attention refactor🚨

[ArthurZucker]

What does this PR do?

Todo in this PR:

• Cohere
• Chameleon
• DBRX
• Gemma
• Gemma2
• GLM (modular donc rien à faire je crois)
• gpt_neoX et GPT2
• Granite
• Jamba
• JetMoe
• Mimi
• Mistral
• Mixtral
• Mllama
• Moshi
• Nemotron
• OPT
• Phi
• Ph3
• PhiMoe
• Qwen2
• qwen2Moe
• qwen2VL
• SableML
• StartCoder2 -> Modular normalement oK
• Idefics1,2,3
• Olmo
• Olmo2
• Siglip
• Whisper

[ArthurZucker]

This should help with kwargs as well

[foreverpiano]

query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)

raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

AttributeError: 'MistralAttention' object has no attribute 'num_heads'

How can I fix this?

[ArthurZucker]

Hey! you should try to use the latest release of transformers! query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) is what's used now.

[ArthurZucker]

Is this by any chance related to AWQ or another package?

[foreverpiano]

Is there any doc about how to migrate from previous version to this version, like the variable definition, the alias change?

[foreverpiano]

Have you tested on several benchmarks about the performance? I knew that the Longbench score on transformer v4.47 vs v4.36 varies a lot on llama-3. Is it stable on this version?
I suggest adding some simple and small dataset tests.

[Cyrilvallez]

Hey! Everything stays the same in terms of user experience/benchmark scores. If you used to hack into the different Layer classes however, it may have changed a bit. You can simply go and check-out the modeling code in this case (as was the case if you hacked into it in the first place I guess!)

[poedator]

My friends use a GPT2Model in production and want to compile it with StaticCache. With the maintainers blessing, I would try to create a PR with DynamicCache / StaticCache support in GPT2Model.
I am quite familiar with Cache class, I already coded some and made the DynamicCache work.

Please let me know if there are any hidden obstacles in Cache implementation for GPT2? Which tests to run or add?
@ArthurZucker

[Rocketknight1]

cc @gante to that question!

[gante]

I've chatted to @poedator offline -- I couldn't think of any obstacle in particular, and suggested a) to ensure we leave a deprecation warning regarding the old cache format b) use RUN_SLOW=1 py.test tests/models/gpt2/test_modeling_gpt2.py as a correctness check (gpt2 is fairly well tested, especially wrt text generation)

[poedator]

It looks like test_flash_attn_2_from_config is broken - it expects attention layer to have flashattention in its name,
if "FlashAttention" in module.__class__.__name__:...
but after this refactoring, the attention classes are named differently.

ref

transformers/tests/test_modeling_common.py

Line 4641 in 5fa3534

if "FlashAttention" in module.__class__.__name__:

please fix or suspend the test.
@ArthurZucker

[ArthurZucker]

indeed gimme a min!