🔴[`Attention`] Bert-based Models Attention Refactor
Author: vasqu
Pull request: yes

Keeping track of the models that are done:
- [x] Bert
- [x] Roberta
- [x] Albert
- [x] Data2VecText
- [x] Xmod 
- [x] Electra 
- [x] XLM Roberta 
- [x] Roberta Prelayernorm
- [x] Ernie 
- [x] Camembert 
- [x] Bert Generation 
- [x] XLM Robert XL 
- [x] RoCBert
- [x] Mobile Bert

Up to discussion:
- Flash attention is flaky; I suspect the norm layers to be responsible (encountered similar things with gemma3)

Would need another round; questionable if worth it (ordered by prio):
- [ ] Tapas
- [ ] Vision x Text models, e.g.
  - [ ] Bridgetower (text would be good to go, would need image counterpart)
  - [ ] Altclip 
  - [ ] ...
- [ ] RemBert
- [ ] Megatron Bert


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

run-slow: bert

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

models: ['models/bert']
quantizations: [] ...

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

run-slow: albert, bert, data2vec, electra, encoder_decoder, markuplm, prophetnet, roberta, roberta_prelayernorm, vision_text_dual_encoder, xmod

Compatible with [AttentionMaskInterface ](https://huggingface.co/docs/transformers/attention_interface#attention-mask-interface)?

@ccdv-ai for now, sadly not. I plan to add support for cross-attention / bidirectional masks in another PR, happy to see/receive any PRs in that direction tho.

what is the timeline on this PR? I am curious if we can remove the implementations from vLLM ASAP

@robertgshaw2-redhat Sadly, there is no specific timeline - there is a huge merge conflict and given the size of the PR I need proper time to resolve it :/

@robertgshaw2-redhat we are trying to get it done by EOW!˜

run-slow: bert, auto, bart, roberta


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

models: ['models/auto', 'models/bart', 'models/bert', 'models/roberta']
quantizations: [] ...

run-slow: bert, auto, bart, roberta

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

models: ['models/auto', 'models/bart', 'models/bert', 'models/roberta']
quantizations: [] ...

Same tests fail on main (`examples_torch`) + some fa tests fail but they are known to be flaky. Otherwise, looks good!

Hi, this PR breaks prefix tuning in PEFT. Unfortunately, that error was masked by several other errors, which is why I only discovered it just now.

Here is a small reproducer:

```python
import torch
from transformers import AutoModel
from peft import PrefixTuningConfig, get_peft_model

model_id = "hf-internal-testing/tiny-random-BertModel"
model = AutoModel.from_pretrained(model_id)
config = PrefixTuningConfig(task_type="FEATURE_EXTRACTION", num_virtual_tokens=10)
model = get_peft_model(model, config)
input_ids = torch.tensor([[1, 1, 1], [1, 2, 1]])
attention_mask = torch.tensor([[1, 1, 1], [1, 0, 1]])
output = model(input_ids=input_ids, attention_mask=attention_mask)
```

Resulting in:

```
>       attn_output = torch.nn.functional.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=attention_mask,
            dropout_p=dropout,
            scale=scaling,
            is_causal=is_causal,
            **sdpa_kwargs,
        )
E       RuntimeError: The expanded size of the tensor (3) must match the existing size (13) at non-singleton dimension 2.  Target sizes: [2, 4, 3, 13].  Tensor sizes: [2, 1, 13, 13]
```

In my investigations, I could track it down to the shape of the `attention_mask` when calling `torch.nn.functional.scaled_dot_product_attention`. For both commits, when I check the `attention_mask` that is being passed to `BertModel.forward` (i.e. which is coming from PEFT), the shape is `torch.Size([2, 13])`. However, for the good commit, the mask is expanded to shape `torch.Size([2, 1, 3, 13])`, whereas the bad commit expands it to `torch.Size([2, 1, 13, 13])`:

Unfortunately, the branches have been deleted, so I cannot link to the code directly, but here is what the code that expands the attention mask looks like:

```python
# bad commit: 155f7e2e6237bc4f623aecf9bfd97442d5723963
    def _update_full_mask(
        self,
        attention_mask: Union[torch.Tensor, None],
        inputs_embeds: torch.Tensor,
    ):
        ...
                # attention_mask.shape: (2, 13) | embedding_output.dtype: float32
                attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)
                # output shape: (2, 1, 13, 13)

# good commit: 61eff450d39742ac28789fe7a2a133395b133ace
class BertModel(BertPreTrainedModel):
    def forward(...):
        ...
                # attention_mask.shape: (2, 13) | embedding_output.dtype: float32 | seq_length: 3
                extended_attention_mask = _prepare_4d_attention_mask_for_sdpa(
                    attention_mask, embedding_output.dtype, tgt_len=seq_length
                )
                # output shape: (2, 1, 3, 13)
```

So to me, it looks like the new call to `_prepare_4d_attention_mask_for_sdpa` is missing the `tgt_len` argument when expanding the attention mask, resulting in the reported shape error. Is this change done purposefully or is it an omission?