Add ModernBERT to Transformers
Author: warner-benjamin
Pull request: yes
Labels: name: New model

This PR will add ModernBERT to Transformers.

Comments:
cc @Cyrilvallez ! 🤗 

@ArthurZucker @Cyrilvallez 
ModernBERT requires no `token_type_ids`, but the tokenizers rely on `PreTrainedTokenizerFast`. By default, this produces `token_type_ids`. Is it preferable that we:
1. Use `model_input_names` in all `config.json` of all ModernBERT models. This means that "fresh tokenizers" won't work out of the box, but people don't normally make fresh tokenizers.
2. Create a custom `ModernBertTokenizerFast` that is literally just:
```python

class ModernBertTokenizerFast(PreTrainedTokenizerFast):
    model_input_names = ["input_ids", "attention_mask"]
```

cc @warner-benjamin @orionw

- Tom Aarsen

With all of these changes in place, I was able to confirm that the output to one of the trained models using the original research implementation nearly matches the output of the `transformers` ModernBERT-converted model. The only difference that remains is that the research implementation fuses the `self.mlp(self.mlp_norm(...))` using `@torch.compile(dynamic=True)`. I also only tested with a small input - I still have to test 1) larger inputs, 2) truncation, 3) batches, etc.

Do we allow something like this to get an exact 1-1 match?
```python
    @torch.compile(dynamic=True)
    def compiled_mlp(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.mlp(self.mlp_norm(hidden_states))
```

Here's an indication of the difference between with and without:
```
tensor([[[ 0.0078,  0.0078,  0.0039,  ...,  0.0117,  0.0342,  0.0039],
         [ 0.0156,  0.0000,  0.0234,  ...,  0.0000,  0.0156, -0.0273],
         [-0.0015, -0.0010, -0.0020,  ...,  0.0000,  0.0022, -0.0015],
         ...,
         [-0.0195,  0.0000, -0.0029,  ..., -0.0156,  0.0146,  0.0039],
         [ 0.0078, -0.0273, -0.0059,  ..., -0.0215,  0.0103,  0.0088],
         [-0.0020,  0.0005, -0.0005,  ..., -0.0012, -0.0012, -0.0020]]],
       device='cuda:0', dtype=torch.bfloat16, grad_fn=<SubBackward0>)
Min: -0.375
Max: 0.140625
Mean: -8.249282836914062e-05
Std: 0.01708984375
```


- Tom Aarsen

I have a working first version with FlexAttention but for now there are still numerical differences between all three implementations(fa2, flexattn, sdpa)

@staghado does it give normal results with e.g. mask filling? If so, then feel free to push it into this branch (or if you don't have access, push it to a branch on your fork or something, and I can cherry-pick it into this PR).

Does it use unpadded or padded? I had a feeling like the SDPA/Eager attention differences compared to FA2 actually resulted when applying the sin/cos rotation, but for FlexAttention that would be equivalent as the FA2 path if you're using unpadded.

- Tom Aarsen

The current FlexAttention PyTorch API requires batched query/key input so can't be used with unpadded inputs without pad/unpad logic(Edit : actually it can)
I tried with both unpadding and padding and both exhibit differences with FA2.
I will push the code sometime tomorrow(still need to clean and do some more tests).
Mask filling performance across all three implementations is very similar but the probabilities are slightly off, here is a small table :
input : The ancient [MASK] stood silently at the edge of the dense forest

| impl          | prediction | prob      |
|----------------|-----------|------------|
| FA2            | tree      | 0.09228515625 |
| FA2            | trees     | 0.07177734375 |
| FA2            | ruins     | 0.04638671875 |
| FlexAttention  | tree      | 0.08349609375 |
| FlexAttention  | trees     | 0.0693359375 |
| FlexAttention  | ruins     | 0.047607421875 |
| SDPA           | tree      | 0.09033203125 |
| SDPA           | trees     | 0.0751953125 |
| SDPA           | ruins     | 0.04833984375 |

> Is the prunning necessary or is it to follow previous implementations? 🤗

Not necessary at all. I just added it to 1) match BERT a bit better and 2) allow the BERT tests to pass. I'll remove it fully.

@ArthurZucker due to small differences between FA2/SDPA/Flex/Eager, we would love to default to FA2 if `flash_attn` is installed. Would this be possible via some config option, or do we *have* to rely on users to manually specify `attn_implementation="flash_attention_2"`?

These are the remaining test failures:
```
FAILED tests/models/modernbert/test_modeling_modernbert.py::ModernBertModelTest::test_cpu_offload - torch._dynamo.exc.InternalTorchDynamoError: AttributeError: 'dict' object has no attribute 'node'
FAILED tests/models/modernbert/test_modeling_modernbert.py::ModernBertModelTest::test_disk_offload_bin - torch._dynamo.exc.InternalTorchDynamoError: AttributeError: 'dict' object has no attribute 'node'
FAILED tests/models/modernbert/test_modeling_modernbert.py::ModernBertModelTest::test_disk_offload_safetensors - torch._dynamo.exc.InternalTorchDynamoError: AttributeError: 'dict' object has no attribute 'node'
FAILED tests/models/modernbert/test_modeling_modernbert.py::ModernBertModelTest::test_resize_tokens_embeddings - AssertionError: expected size 84==109, stride 32==32 at dim=0
FAILED tests/models/modernbert/test_modeling_modernbert.py::ModernBertModelTest::test_retain_grad_hidden_states_attentions - RuntimeError: This backward function was compiled with non-empty donated buffers which requires create_graph=False and retain_graph=False. Please keep backwar...
```
with these reasons:

* `test_cpu_offload`, `test_disk_offload_bin`, `test_disk_offload_safetensors`: Unsure, something about 1) torch.compile and 2) some weight being on the meta device, but without a value to put it back on CUDA.
* `test_resize_tokens_embeddings`: Caused by torch.compile breaking with resizing
* `test_retain_grad_hidden_states_attentions`: Caused by torch.compile, though this might be fixable with some parameters to the compile call.

They don't strictly all have to be fixed, but we can also disable those tests and give the users warnings when they do something that doesn't work with this architecture (e.g. resizing tokens seems to not work currently?).

cc @warner-benjamin

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

the model with FA2 and the RoPE kernel are not torch.compile compatible, we can't compile the whole model while using these.

> the model with FA2 and the RoPE kernel are not torch.compile compatible, we can't compile the whole model while using these.

FA2 is compatible now, but the FA RoPE kernel isn't yet. I have a in progress fix I need to get merged into the FA repo.

We want mean pooling as an option for classification because Local Attention means unlike BERT the CLS token doesn't see all the output in all the attention layers, so mean pooling could outperform CLS on greater than 128 token sequences.

Also, I added the pooling head to TokenClassification because otherwise we are throwing away one pretrained linear layer `ModernBertPoolingHead.dense`.

Great addition! Thanks all for your hard work! 🤗