Adding FlexAttention Support for Qwen2 models

[MekkCyber]

What does this PR do?

Adding flex_attention for Qwen2 models following #34809

Who can review ?

@ArthurZucker

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

[MekkCyber]

Hey @ArthurZucker, I fixed the issue related to Qwen2Moe copying from Qwen2 attention classes, however I'm still running into some errors with make fixup and I don't see the relationship between the error related to starcoder and the Qwen2 models :

No differences found for src/transformers/models/llava_next_video/configuration_llava_next_video.py.
No differences found for src/transformers/models/llava_next_video/modeling_llava_next_video.py.
No differences found for src/transformers/models/ijepa/modeling_ijepa.py.
No differences found for src/transformers/models/aria/configuration_aria.py.
No differences found for src/transformers/models/aria/modeling_aria.py.
No differences found for src/transformers/models/aria/image_processing_aria.py.
No differences found for src/transformers/models/aria/processing_aria.py.
No differences found for src/transformers/models/olmo2/configuration_olmo2.py.
No differences found for src/transformers/models/olmo2/modeling_olmo2.py.
No differences found for src/transformers/models/glm/modeling_glm.py.
No differences found for src/transformers/models/gemma/configuration_gemma.py.
No differences found for src/transformers/models/gemma/tokenization_gemma.py.
No differences found for src/transformers/models/gemma/modeling_gemma.py.

Differences found between the generated code and src/transformers/models/starcoder2/modeling_starcoder2.py:

   1 --- src/transformers/models/starcoder2/modeling_starcoder2.py_generated
   2 +++ src/transformers/models/starcoder2/modeling_starcoder2.py
   3 @@ -646,7 +646,6 @@
   4      _supports_cache_class = True
   5      _supports_quantized_cache = True
   6      _supports_static_cache = True
   7 -    _supports_flex_attn = True
   8  
   9      def _init_weights(self, module):
  10          std = self.config.initializer_range
No differences found for src/transformers/models/gemma2/configuration_gemma2.py.
No differences found for src/transformers/models/gemma2/modeling_gemma2.py.
No differences found for src/transformers/models/instructblipvideo/configuration_instructblipvideo.py.
No differences found for src/transformers/models/instructblipvideo/modeling_instructblipvideo.py.
Traceback (most recent call last):
  File "/fsx/mohamed/dev/transformers/utils/check_modular_conversion.py", line 78, in <module>
    raise ValueError("Some diff and their modeling code did not match.")
ValueError: Some diff and their modeling code did not match.
make: *** [Makefile:39: repo-consistency] Error 1

[ArthurZucker]

sorry @MekkCyber, was workiing on #35235 which might include some of these changes, and stabilized the api! Do you mind updating? 🤗

[lapp0]

Thanks for working on this, I've had a great experience with FlexAttention so far.

Wouldn't it make sense to create a BlockMask with a specified mask_mod once and reuse it in every layer?

Rather that calculating attention and zeroing it via score_mod, we could just not calculate attention between tokens excluded from the mask. Something like

def get_block_mask(position_ids):
    def make_mask_mod(b, h, q_idx, kv_idx):
        return (position_ids[q_idx] >= position_ids[kv_idx]) & (q_idx - kv_idx < sliding_window_size)
    S = position_ids.size(0)
    return create_block_mask(make_mask_mod, None, None, S, S)

[fxmarty-amd]

+1, I did not try both but using block_mask should be faster in case the mask is large enough and sparse: https://pytorch.org/blog/flexattention/#q-why-are-score_mod-and-mask_mod-different-isnt-mask_mod-just-a-special-case-of-score_mod

[fxmarty-amd]

actually, this is the case for all archs, Transformers only makes use of score_mod:

transformers/src/transformers/integrations/flex_attention.py

Line 40 in 31bb662

score_mod=causal_mod,