ModernBERT FlexAttention
Author: staghado
Pull request: yes
Labels: name: flex attention

# 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

Comments:
I think this is very interesting, but only if the performance rivals that of e.g. SDPA. I see your issue here: https://github.com/pytorch-labs/attention-gym/issues/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

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?

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 

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](https://github.com/pytorch-labs/attention-gym/issues/95).
It also looks like the flexattention in Gemma was removed?


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

With the release of [PyTorch 2.6](https://pytorch.org/blog/pytorch2-6/), it is now possible to use FlexAttention with ModernBERT without a nightly requirement.

Sorry @staghado for the delay!

I skimmed through the linked PR, flex_attention does not get compiled automatically so the compilation call is needed. As you can see [here](https://github.com/pytorch-labs/attention-gym/issues/95). And there are hardcoded stuff for GQA, softcapping, etc which are not universal for all implementations.
Another thing is the integration uses score_mod which is less performant than block mask as mentioned [here](https://pytorch.org/blog/flexattention/#q-how-can-we-compute-blockmask-quicker)

I would prefer keeping the ModernBERT implementation separate for now, at least until the integration is more mature.

Hi @staghado my PR https://github.com/huggingface/transformers/pull/36103/files does add torch compile for flex attention, it is done in the `integrations/flex_attention.py` such that minimal modifications are required in the various `modeling_{model type}.py` for them to take advantage of flex attention.

I believe both score mod and block mask are needed - the score mod can modify attention scores in ways that can't be done using block mask. The intention is to perform attention masking solely using block mask, while score mod is only used for non-masking score modification.

Would love to hear more about your thoughts!

Hi @bursteratom,
One big distinction between score mod and mask mod is that the latter is much more efficient in practice, for e.g with ModernBERT it allows implementing unpadded attention similar to flash attention. 
The integration seems focused on causal decoder-only models with causal attention and GQA always set to true?, I think if we want it to be re-usable across architectures it should be as general as the [PyTorch API](https://github.com/pytorch/pytorch/blob/main/torch/nn/attention/flex_attention.py) with all the necessary knobs.

Another thing is, to get the performance promised by FlexAttention, we need to compile both `flex_attention` and `create_block_mask`.
@ArthurZucker for the reasons above I think it's better to keep this separate for now.

Hey #36643 was merged, just needs to integrate the refactor points imo and I'll have another look!

Looks like the integration is too focused on decoder-only models(gqa=true at api level), create_block_mask is renamed to create_block_causal_mask_flex, a score_mod is always passed even if not needed. This PR is out of scope.