[Model] Refactor modernbert with the attention interface
Author: YangKai0616
Pull request: yes

# What does this PR do?

This PR:

1. Refactors `ModernBERT's Flash Attention` implementation using the built-in interfaces from `Transformers`.
2. Fixes some UT issues in the `tests/models/modernbert/test_modeling_modernbert.py` and `tests/models/pe_video/test_modeling_pe_video.py` test files. 

Background PR: https://github.com/huggingface/transformers/pull/42952.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Comments:
Hey @vasqu , could you please review it again? Thank you very much!

Nice thx a lot @YangKai0616, I'll take a look tomorrow, just coming back from the holidays :hugs: 

 Hey @vasqu, I've updated the code based on your comments and removed the unrelated UT changes. Please take another look. Thanks!

I’ve modified the code according to your suggestions. I’ll find some time tomorrow to test the related unit tests. If everything looks good, we can continue with the review. Thank you for your patience!

@YangKai0616 just ping me again when you think it's ready for another review :saluting_face: 

The CI error seems unrelated to this PR. I was able to test it successfully on my local CUDA and XPU. @vasqu let's review it again, thank you !

Ready for review @YangKai0616? Will take a look tomorrow then

> Ready for review @YangKai0616? Will take a look tomorrow then

Current version:
1. Can be tested via **test_modernbert.py**
2. Short sequences perform the same as before, but long sequences involving **sliding window** have some slight errors at the moment. I'm currently working on fixing this.

Hey @vasqu , we can review it again!
In my local tests, the current version of **FA2/sdpa/eager** can achieve results that are almost identical to those of the main branch.
This change mainly modifies the data format of **FA2 rotary (.float())** and fixes the **sliding_window** calculation to match the implementation in the main branch.

Hmm... it seems there are some issues with the CI. I can't reproduce these failing cases locally(CUDA and XPU).

> Hmm... it seems there are some issues with the CI. I can't reproduce these failing cases locally(CUDA and XPU).

Yea our CI is pretty flaky these days, let's focus on modernbert only for now. 

> This change mainly modifies the data format of FA2 rotary (.float()) and fixes the sliding_window calculation to match the implementation in the main branch.

Do the kernels give exact matches / is there a big diff between kernel vs non-kernel RoPE? We might need to adjust our integration test(s)

> Do the kernels give exact matches / is there a big diff between kernel vs non-kernel RoPE? We might need to adjust our integration test(s)

The kernels are an exact match.

For the torch implementation version, I tested it with the configuration set to `batch_size, num_heads, seq_len, head_dim = 2, 8, 128, 64 dtype = torch.bfloat16`, and the maximum difference is on the order of `e-2`, while the average difference is on the order of `e-3`. After adding `.float()`, the results become completely equal.

> Can we add something to the modernbert docs to show how to do paddingfree inference/training. I.e. some tips on how to utilize the model optimally. I have a [gist for llama](https://gist.github.com/vasqu/c264d14268fbc7f1eb725e09dfb87331) but should be easy to adapt to modernbert

Done.

> The kernels are an exact match. [...] After adding .float(), the results become completely equal.

So we can achieve parity either way now (seeing we use the float version), be it with kernels or the torch version? If the kernels are essential, then we should put it to the docs

> So we can achieve parity either way now (seeing we use the float version), be it with kernels or the torch version? If the kernels are essential, then we should put it to the docs

Yes, for correctness, kernels are not essential — the torch version with `.float()` works equally well.

I’ve run all the failed CI cases, and the following three are confirmed to fail:
```
tests/models/modernbert_decoder/test_modeling_modernbert_decoder.py::ModernBertDecoderModelTest::test_resize_embeddings_untied
tests/models/modernbert_decoder/test_modeling_modernbert_decoder.py::ModernBertDecoderModelTest::test_resize_tokens_embeddings
tests/models/pe_video/test_modeling_pe_video.py::PeVideoModelTest::test_can_set_attention_dynamically_composite_model
```
Among them, `test_resize_embeddings_untied` and `test_resize_tokens_embeddings` are failing because the changes from before in the modular file weren’t updated in modeling file — these are unrelated to this PR.
`test_can_set_attention_dynamically_composite_model` was mentioned before. We need to skip this example.



If `tests/models/pe_video/test_modeling_pe_video.py::PeVideoModelTest::test_can_set_attention_dynamically_composite_model` is blocking us, I'm pro skipping for now - but let's check with our main (merging in a second + running slow tests to check if they work as expected)

Yea we need to take a look: https://app.circleci.com/pipelines/github/huggingface/transformers/160664/workflows/f3cbb6f6-9b01-470b-90e7-1375771e76f6/jobs/2114314 - do you what changed this behavior? It passes on main

run-slow: modernbert, modernbert_decoder

This comment contains `run-slow`, running the specified jobs: 

models: ["models/modernbert", "models/modernbert_decoder"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/21217230291)

### Model CI Report

#### ❌ Failed tests

- [modernbert_decoder](https://github.com/huggingface/transformers/actions/runs/21217230291/job/61042252744):
    tests/models/modernbert_decoder/test_modeling_modernbert_decoder.py::ModernBertDecoderModelTest::test_resize_embeddings_untied
    tests/models/modernbert_decoder/test_modeling_modernbert_decoder.py::ModernBertDecoderModelTest::test_resize_tokens_embeddings



run-slow: modernbert, modernbert_decoder

Do you have a quick script for testing against main re https://github.com/huggingface/transformers/pull/43030#issuecomment-3772099830 @YangKai0616 

I would verify correctness myself again just to be sure. I've added a few smaller changes, let me know if these work for you (mainly tests). Let's also keep https://github.com/huggingface/transformers/pull/43030#issuecomment-3779348609 in mind

This comment contains `run-slow`, running the specified jobs: 

models: ["models/modernbert", "models/modernbert_decoder"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/21220613751)

✅ No failing test specific to this PR 🎉 !


> Do you have a quick script for testing against main re [#43030 (comment)](https://github.com/huggingface/transformers/pull/43030#issuecomment-3772099830) @YangKai0616
I would verify correctness myself again just to be sure. 
> 
> I've added a few smaller changes, let me know if these work for you (mainly tests). 
> 
> Let's also keep [#43030 (comment)](https://github.com/huggingface/transformers/pull/43030#issuecomment-3779348609) in mind
1.
```
import torch
from kernels import get_kernel


# ==================== PyTorch ====================
# from transformers.models.gemma3.modeling_gemma3 import apply_rotary_pos_emb

# ==================== PyTorch + .float() ====================
from transformers.models.modernbert.modeling_modernbert import apply_rotary_pos_emb

# ==================== Kernel ====================
rotary_kernel = get_kernel("kernels-community/rotary")
apply_rotary_pos_emb_kernel = rotary_kernel.apply_rotary_transformers

# ================================================

def test_with_proper_rope_format():
    batch_size, num_heads, seq_len, head_dim = 2, 8, 128, 64
    half_dim = head_dim // 2
    device = "cuda"
    dtype = torch.bfloat16
    
    q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    
    cos_half = torch.randn(batch_size, seq_len, half_dim, device=device, dtype=dtype)
    sin_half = torch.randn(batch_size, seq_len, half_dim, device=device, dtype=dtype)
    
    cos_full = torch.cat([cos_half, cos_half], dim=-1)
    sin_full = torch.cat([sin_half, sin_half], dim=-1)
    
    print(f"q shape: {q.shape}")
    print(f"cos_half shape: {cos_half.shape}")
    print(f"cos_full shape: {cos_full.shape}")
    
    q_pytorch, k_pytorch = apply_rotary_pos_emb(
        q.clone(), k.clone(), cos_full.clone(), sin_full.clone(), unsqueeze_dim=1
    )

    q_kernel, k_kernel = apply_rotary_pos_emb_kernel(
        q.clone(), k.clone(), cos_full.clone(), sin_full.clone(), unsqueeze_dim=1
    )

    q_close = torch.allclose(q_pytorch, q_kernel, atol=1e-4, rtol=1e-4)
    k_close = torch.allclose(k_pytorch, k_kernel, atol=1e-4, rtol=1e-4)
    
    q_max_diff = (q_pytorch - q_kernel).abs().max().item()
    k_max_diff = (k_pytorch - k_kernel).abs().max().item()
    q_mean_diff = (q_pytorch - q_kernel).abs().mean().item()
    k_mean_diff = (k_pytorch - k_kernel).abs().mean().item()
    
    print(f"\nQ output match: {q_close}, max difference: {q_max_diff:.2e}")
    print(f"K output match: {k_close}, max difference: {k_max_diff:.2e}")
    print(f"Q mean difference: {q_mean_diff:.2e}")
    print(f"K mean difference: {k_mean_diff:.2e}")
    print(f"Test result: {'Passed ✓' if q_close and k_close else 'Failed ✗'}")

if __name__ == "__main__":
    print("Rotary Position Embedding consistency test")
    print("=" * 60)
    
    if not torch.cuda.is_available():
        print("CUDA not available, exiting")
        exit(1)
    
    print(f"device: {torch.cuda.get_device_name(0)}")
    
    test_with_proper_rope_format()

```

2. `test_inference_masked_lm_flash_attention_2` shows an error on the order of `e-2` on XPU (PVC). I added a ground truth specifically for XPU.

3. Because `_can_set_attn_implementation()` [here](https://github.com/huggingface/transformers/blob/de306e8e14672dd8392b4bd344054a6a18de8613/tests/test_modeling_common.py#L4379) returned `False` before the refactoring, the main branch would skip this test and report `"Parts of this model cannot set attention dynamically"`.
After the refactoring, it now returns `True`, so the test continues to run and causes `TimmWrapper` to error out.
Therefore, this test needs to be explicitly skipped.

Commit https://github.com/huggingface/transformers/pull/43030/commits/587cb4f06bf3c96266047e7fe40bedafd98eb8e0 submitted.


Thanks a lot! I've also verified with the script you gave. Also sanity checked with 
```python
from transformers import AutoModelForMaskedLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained(
    "answerdotai/ModernBERT-base",
)
model = AutoModelForMaskedLM.from_pretrained(
    "answerdotai/ModernBERT-base",
    dtype=torch.float16,
    device_map="auto",
    reference_compile=False,  # only relevant for the old model
)
inputs = tokenizer("Plants create [MASK] through a process known as photosynthesis.", return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs)
    predictions = outputs.logits

masked_index = torch.where(inputs['input_ids'] == tokenizer.mask_token_id)[1]
predicted_token_id = predictions[0, masked_index].argmax(dim=-1)
predicted_token = tokenizer.decode(predicted_token_id)

print(f"The predicted token is: {predicted_token}")
```
And they gave nearly identical matches (atol/rtol of 1e-8 matched)

> test_inference_masked_lm_flash_attention_2 shows an error on the order of e-2 on XPU (PVC). I added a ground truth specifically for XPU.

Perfect

> Because _can_set_attn_implementation() [here](https://github.com/huggingface/transformers/blob/de306e8e14672dd8392b4bd344054a6a18de8613/tests/test_modeling_common.py#L4379) returned False before the refactoring, the main branch would skip this test and report "Parts of this model cannot set attention dynamically".
After the refactoring, it now returns True, so the test continues to run and causes TimmWrapper to error out.
Therefore, this test needs to be explicitly skipped.

Gotcha, thanks for clarifying

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43030&sha=d3b23b

Hello!

This is an extremely impressive PR, wow. The structure of the modeling/modular files look very strong as a result of all of this. I ran some downstream tests using Sentence Transformers with https://huggingface.co/nomic-ai/modernbert-embed-base, and there are some very slight differences in scores:

```
sdpa <v5:
{'NanoBEIR_mean_cosine_ndcg@10': 0.5147433277390966,
 'NanoMSMARCO_cosine_ndcg@10': 0.5867981583748159,
 'NanoNFCorpus_cosine_ndcg@10': 0.329223326613997,
 'NanoNQ_cosine_ndcg@10': 0.6282084982284769}
Evaluation took 27.23 seconds

eager <v5:
{'NanoBEIR_mean_cosine_ndcg@10': 0.5123706576330123,
 'NanoMSMARCO_cosine_ndcg@10': 0.5770369329466852,
 'NanoNFCorpus_cosine_ndcg@10': 0.32729203600264306,
 'NanoNQ_cosine_ndcg@10': 0.6327830039497084}
Evaluation took 44.04 seconds

fa2 <v5: I don't have flash_attn installed in the right venv at the moment; goes to show how nice having kernels support is

sdpa pr-43030:
{'NanoBEIR_mean_cosine_ndcg@10': 0.5110831146710934,
 'NanoMSMARCO_cosine_ndcg@10': 0.5828975140199241,
 'NanoNFCorpus_cosine_ndcg@10': 0.32718235725979056,
 'NanoNQ_cosine_ndcg@10': 0.6231694727335656}
Evaluation took 32.49 seconds

eager pr-43030:
{'NanoBEIR_mean_cosine_ndcg@10': 0.5099581159107758,
 'NanoMSMARCO_cosine_ndcg@10': 0.5843932028670678,
 'NanoNFCorpus_cosine_ndcg@10': 0.32340108499481846,
 'NanoNQ_cosine_ndcg@10': 0.622080059870441}
Evaluation took 54.74 seconds

fa2 pr-43030:
{'NanoBEIR_mean_cosine_ndcg@10': 0.5119965071683512,
 'NanoMSMARCO_cosine_ndcg@10': 0.581536709004633,
 'NanoNFCorpus_cosine_ndcg@10': 0.3271691052363954,
 'NanoNQ_cosine_ndcg@10': 0.6272837072640253}
Evaluation took 25.31 seconds
```

When I implemented this architecture, these forms of slight differences also existed between e.g. eager/sdpa/fa2, and it's extremely hard to narrow down the source (and from experience: often enough even with the source known, it's not possible to do anything about it). 

The old implementation automatically 'unpadded' and 'repadded' during inference. I take it this variant doesn't do that? You write about a `DataCollatorWithFlattening` etc., is this required now for this speedup? Is there any way to know whether any given `transformers` models supports this class? It would be valuable for me to know for Sentence Transformers.

Also, this might be an issue on my side, but I wasn't able to get `torch.compile` working:
```
torch._dynamo.exc.Unsupported: Data-dependent branching
  Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.
  Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
  Hint: Use `torch.cond` to express dynamic control flow.
```

In short: you've done some real magic here, and it's (at the least) very similar in accuracy, albeit likely quite a bit slower due to not forcing the unpadding/repadding (iiuc). However, that was never very `transformers`-y, and with this PR it should be much easier to build on top of ModernBERT again.

- Tom Aarsen

> The old implementation automatically 'unpadded' and 'repadded' during inference. I take it this variant doesn't do that? You write about a DataCollatorWithFlattening etc., is this required now for this speedup? Is there any way to know whether any given transformers models supports this class? It would be valuable for me to know for Sentence Transformers.

Pretty much any model that supports attention backend should be able to do this. I refactored Bert for exactly the same reason for example (vLLM needs this as well). Transformers was never meant to manually pad/unpad during the forward, that's what the collator is for and is recommended for causal models like llama for quite a while now - this enables complete compile compatibility with FA. 

Some side note: we manually pad/unpad ourselves already in FA (given a non all true attention mask)

> albeit likely quite a bit slower due to not forcing the unpadding/repadding (iiuc)

Should be fast when properly preparing the data :D 

> Also, this might be an issue on my side, but I wasn't able to get torch.compile working

Can you share a script and env? The example works on my end (FA), also check SDPA with non fullgraph. Note
- The data collator is necessary for clean compile compatibility - attention masks should also work but not full graph and minor graph breaks
- Always compile the forward only
- SDPA does not need the collator, it's only for FA

> The data collator is necessary for clean compile compatibility 

I see, I might've gone wrong there. I'll run more tests later. I was testing with Sentence Transformers, which doesn't use the collator.

- Tom Aarsen

Hi everyone, any updates? 😀

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: modernbert, modernbert_decoder, pe_video

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: modernbert, modernbert_decoder, pe_video

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_43030). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: modernbert, modernbert_decoder, pe_video

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: modernbert, modernbert_decoder, pe_video

run-slow: modernbert, modernbert_decoder

This comment contains `run-slow`, running the specified jobs: 

models: ["models/modernbert", "models/modernbert_decoder"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/21479031512)

✅ No failing test specific to this PR 🎉 !


**[For maintainers]** Suggested jobs to run (before merge)

run-slow: modernbert, modernbert_decoder, pe_video

Green CI, no regression, huge work @YangKai0616! Thanks a lot for this PR and iterating time and time again :heart: it will be merged after the docs are built