DeBERTa's `DisentangledSelfAttention` hardcodes `float` dtype, which causes `bfloat16` overflow error
Author: bauwenst
Pull request: no
Labels: name: bug

### System Info

transformers: 4.47.0
Python: 3.10.5
PyTorch: 2.5.1+cu124
GPU: NVIDIA GTX 980 Ti

### Who can help?

@ArthurZucker 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I'm training a `DebertaForMaskedLM` model with a broader experimental framework, but you can reproduce the bug with simple inference as follows: instantiate such a model with datatype `bfloat16`, and send a batch through it.
```python
import torch
from transformers import DebertaConfig, DebertaForMaskedLM

model = DebertaForMaskedLM._from_config(DebertaConfig(), torch_dtype=torch.bfloat16)
model(**{"input_ids": torch.tensor([[101,102,103,104]]),
         "attention_mask": torch.tensor([[1,1,1,1]])})
```
One of two errors is now thrown in `modeling_deberta.py`, both in `DisentangledSelfAttention.forward()` (and they can both be traced back to the same issue):
1. `RuntimeError: expected m1 and m2 to have the same dtype, but got: float != struct c10::BFloat16`
2. `RuntimeError: value cannot be converted to type at::BFloat16 without overflow`

Here's where they come from: two fields in DeBERTa's `DisentangledSelfAttention` are constructed by explicitly declaring their `dtype` as `torch.float`:

https://github.com/huggingface/transformers/blob/9613933b022ddbf085e2c593ed4ceea4c734179a/src/transformers/models/deberta/modeling_deberta.py#L187-L188

Then, in `forward()`, we create the two tensors `query_layer` and `key_layer` that start out with the `dtype` of the hidden states, which have the `dtype` of the model, namely `bfloat16`:

https://github.com/huggingface/transformers/blob/9613933b022ddbf085e2c593ed4ceea4c734179a/src/transformers/models/deberta/modeling_deberta.py#L258-L259

But then, one of these tensors, `query_layer`, is modified by adding `self.q_bias` into it. The resulting tensor inherits the `torch.float` data type:

https://github.com/huggingface/transformers/blob/9613933b022ddbf085e2c593ed4ceea4c734179a/src/transformers/models/deberta/modeling_deberta.py#L268

The first RuntimeError can occur on the following line, when `query_layer` (now `torch.float`) and `key_layer` (still `torch.bfloat16`) are multiplied. I've had this line crash on one machine and work on another, so perhaps this kind of mixed precision sometimes works.

https://github.com/huggingface/transformers/blob/9613933b022ddbf085e2c593ed4ceea4c734179a/src/transformers/models/deberta/modeling_deberta.py#L276

The second RuntimeError occurs even when mixed precision is supported. It happens on the following line:

https://github.com/huggingface/transformers/blob/9613933b022ddbf085e2c593ed4ceea4c734179a/src/transformers/models/deberta/modeling_deberta.py#L290

`attention_scores` is of type `bfloat16`. You then ask to fill it with the minimal value *for the data type of `query_layer`, not the data type of `attention_scores`*. Because `query_layer.dtype` is `torch.float`, that minimal value (-3.40282e+38) is *more negative than the most negative `torch.bfloat16`* (-3.38953e+38). Hence, the overflow.

### Expected behavior

The `dtype` of `self.q_bias` and `self.v_bias` should be set like the rest of the modules/tensors in the model, rather than being hardcoded. That would keep everything `bfloat16`.

Comments:
Reviewed the PR! We should make sure this change is transparent as much as possible ! 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.