ModernBERT FlexAttention

[staghado]

What does this PR do?

This PR adds FlexAttention support for ModernBERT:

• Combines sliding window and document masking to implement the alternating local/global attention pattern in ModernBERT
• Mask creation is expensive so the two masks are cached at the model level and then re-used across layers.
• Similar to the FA2 path, it works directly on the unpadded sequences
• Re-uses the existing ModernBertRotaryEmbedding to avoid requiring FA2.

Note:
The current version requires one of the latest torch nightlies (e.g 2.6.0.dev20241112)
Currently transformers does not allow compiling the flex_attention function IIUC

[tomaarsen]

I think this is very interesting, but only if the performance rivals that of e.g. SDPA. I see your issue here: pytorch-labs/attention-gym#95, which also shows that with compilation, Flex Attention outperforms SDPA in a lot of common cases, so we would have to introduce the compilation option for Flex Attention in transformers before this is viable I think.

I'd love to see that implemented, though!

• Tom Aarsen

[staghado]

Yeah when compiled FlexAttention is generally faster than SDPA and has much lower memory from my tests.
How would we go about adding the compilation option? would a flag suffice in this case?

[staghado]

Any ideas on how transformers plans to support FlexAttention compilation? I see that no models have it implemented for now(I might be wrong here)
cc @tomaarsen @ArthurZucker

[ArthurZucker]

Looks alright appart from the nightly requirement!
Regarding compilation, we do support compile, and by default when you use generate with cache_implementation="static".

[staghado]

I think the nightly requirement might be resolved by specifying the BLOCK_SIZE argument in create_block_mask, will verify that.
for compilation, I meant supporting the flex_attention = torch.compile(flex_attention, dynamic=False) which is very important for performance as mentioned here.
It also looks like the flexattention in Gemma was removed?

[neavo]

What does this PR do?

This PR adds FlexAttention support for ModernBERT:

• Combines sliding window and document masking to implement the alternating local/global attention pattern in ModernBERT
• Mask creation is expensive so the two masks are cached at the model level and then re-used across layers.
• Similar to the FA2 path, it works directly on the unpadded sequences
• Re-uses the existing ModernBertRotaryEmbedding to avoid requiring FA2.

Note: The current version requires one of the latest torch nightlies (e.g 2.6.0.dev20241112) Currently transformers does not allow compiling the flex_attention function IIUC

Torch 2.6 has released the RC version. The dependency issues should not be a big problem anymore. Looking forward to its merging.
Before that, could you please update the branch of this PR to make its code compatible with the current mainline version?
I would like to give it a try.

[staghado]

these are needed to actually use the flexattention kernels but the utils/modular_model_converter.py does not allow it in the converted file.
let me know if there is a better way to do this

[staghado]

cc @ArthurZucker @tomaarsen (sorry for the ping)

[staghado]

any ideas on how to support compiling the flex_attention function in transformers?

[staghado]

The failing CI is related to these two lines, there is no clear way how to support compile in a clean way

[staghado]

With the release of PyTorch 2.6, it is now possible to use FlexAttention with ModernBERT without a nightly requirement.