update modular_modernbert -- add inputs_embeds param to ModernBertModel
Author: jxmorris12
Pull request: yes

# What does this PR do?
Hi! Congrats on the release of ModernBERT; it looks amazing. I'm interested in using ModernBERT eventually to train a new [Contextual Document Embeddings](https://arxiv.org/abs/2410.02525) model. 

One desired feature is to pass the contextual and word embeddings together in the second stage, which requires setting the `inputs_embeds` kwarg so that we can pass hidden states directly. This is a feature of typical BERT and other transformer implementations but isn't yet allowed by ModernBERT, so I added it. It's only a few additional lines of code.

cc: @warner-benjamin @tomaarsen @orionw @staghado @bclavie @NohTow @ArthurZucker



Comments:
Hello @jxmorris12, @ArthurZucker,

I pushed some changes into this PR to get it closer to completion. Let me know if you're not okay with this, and you can easily revert or delete the commits.

The changes:
1. The `inputs_embeds` shouldn't fully replace `self.embeddings(input_ids)`, because this call also does layer normalization and dropout. So, now both `input_ids` and `inputs_embeds` is passed to the `ModernBertEmbeddings`, much like how `BertEmbeddings` is implemented.
2. I also added `inputs_embeds` to the docstring, and propagated the changes to the other model classes.
3. I also introduced an error if `input_ids` and `input_embeds` are both or neither provided.
4. I fixed an issue with device being based solely on `input_ids` with `attention_mask`.
5. Fix a test, and reintroduce another test.

> let's make it a bit simpler (si how llama has less if elses)

This is sadly not as simple as it seems due to `_unpad_modernbert_input`.

- Tom Aarsen

Hello,

Thanks for opening this PR, I really love your work on CDE and I look forward to seeing how it will perform with ModernBERT!
FYI, this feature (being able to pass `input_embeds`) has also been asked for other use cases [here](https://github.com/AnswerDotAI/ModernBERT/issues/165), so besides retro-compatibility and CDE, it seems like this feature is used by the community and is thus a very cool addition.

I'll make a more thorough review when I come back from vacation this week, but I already checked how to implement it for the linked issue and it seems to be in line with the latest change from Tom.

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

> Hello @jxmorris12, @ArthurZucker,
> 
> I pushed some changes into this PR to get it closer to completion. Let me know if you're not okay with this, and you can easily revert or delete the commits.
...

thanks Tom! Looks great.

Hi @NohTow @ArthurZucker just following up on this -- thanks :) 

when will it be merged into the main branch?

Hello,
Sorry for the delay.
As mentioned earlier, the code looks good to me.
I only have two small nitpick. As mentioned by @ArthurZucker, first embedding the input_ids should be cleaner and less error prone.
The second is that the XOR test is somewhat a bit less informative than checking if the two are not set together and then checking if at least one is defined. Maybe it's just me, but I find `"You must specify exactly one of input_ids or inputs_embeds"` a bit ambiguous, yet I don't have a better wording.  

Besides those very small details, everything looks good!



The xor test is on a few models like Llama -> it's standard for us! 🤗 

Sorry @NohTow - can you make a concrete suggestion on how to fix the code? It's not clear how the code can be much cleaner since all the conditionals seem necessary to me.

@jxmorris12 I am very sorry my message was unclear. I did not mean that the code was incorrect nor that testing the two conditions sequentially was cleaner, I just meant that from my reading, the error message could be ambiguous (in the case where the user do not feed anything, the _exactly_ might not be explicit enough).
But again, that is really a nitpick and probably comes from me not properly reading the message, let's ignore that!

@jon-tow Oh okay, that makes sense to me! I agree it could be a really confusing error for the user. But maybe since it's in a lot of files (such as LLAMA) we could open a separate issue to improve that error message everywhere?

@jxmorris12 Yeah totally, I raised that because I compared the implementations from BERT and this one and the BERT one does the check sequentially and thus have more informative messages, but since Arthur raised that it is already done like that in other models and is standard now, let's just follow the standard!

For the latter, I think it is related to [this](https://github.com/huggingface/transformers/pull/35404).
It was discussed during the original PR, but it breaks the gradients.
I'll let @warner-benjamin give more information.

> My comments are not adressed 😓
> 
>     * let's embed inputs, then we computes shapes

This is a very logical approach, so that we don't need
```python
            if inputs_embeds is not None:
                batch_size, seq_len = inputs_embeds.shape[:2]
            else:
                batch_size, seq_len = input_ids.shape[:2]
```
But instead can just use
```python
batch_size, seq_len = hidden_states.shape[:2]
```

However, in ModernBERT, before we can compute the embedded inputs, we need to potentially unpad. To unpad, we need to have an `attention_mask`, which needs the shapes. In short: we need the shapes before the input embeddings.

One alternative is to use `attention_mask = torch.ones((input_ids or input_embeds).shape[:2], ...)`, but this feels like the same as the shape computation but as a 1-liner.

>     * let's never require grad on unpadding / padding as it's weight agnostic

Please have a look at #35386, it looks like the `torch.no_grad` is killing the gradients that were there before re-padding. Running the same script from that PR on BERT *does* give gradients on the logits.

@warner-benjamin made a PR to patch it here: #35404, which makes the gradient for repadding optional.

- Tom Aarsen

Got it! Okay, for output, we do need the gradients, my bad! 
TLDR: 
- when training: with torch no grad for input, but you need output gradients
- when training: with input embeddings: torch grad because you need propagation 
- when inference, never grad, but usually people do this outside the modeling. 

So the only "gain" is using no grad on input ids all the time