transformers - 💡(How to fix) Fix Suggestion: Give a warning for 4D attention mask which use float dtype but in 0.0 and 1.0 term (instead of -inf and zero).

Feature request

The title said it all.

Motivation

Let's say I have a code like this: (Early warning: If you run this code, it will give an error.)

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained("google-bert/bert-base-uncased", device_map="auto")
attention_mask = torch.tensor([
    [1, 0],
    [0, 0],
], device=model.device).unsqueeze(0).unsqueeze(1)

outputs = model(
    input_ids=torch.tensor([[200, 201]], device=model.device),
    attention_mask=attention_mask,
)
print("Output 1:")
print(outputs.last_hidden_state[0, 0, :10])

outputs = model(
    input_ids=torch.tensor([[200, 202]], device=model.device),
    attention_mask=attention_mask,
)
print("Output 2:")
print(outputs.last_hidden_state[0, 0, :10])

Notice that the last input ID for output 1 and output 2 are different: 201 and 202. However, since it's designed to be masked, then the output should not be impacted.

But if you run that code, it will return an error because the attention mask are expected to be in float or boolean type.

My previous attempt was converting it to float, so the line 5-8 should be changed like this:

...
attention_mask = torch.tensor([
    [1, 0],
    [0, 0],
], device=model.device).unsqueeze(0).unsqueeze(1).float()
...

But here's the funky part: if you use that change, both outputs give different result! I was surprising when I found it, so I try to debug it.

After a debug session that leads me to a Pytorch page, my mistake is that it should be changed to boolean, so the change should be like this:

...
attention_mask = torch.tensor([
    [1, 0],
    [0, 0],
], device=model.device).unsqueeze(0).unsqueeze(1).bool()
...

If I want to use in a float term, the 0 terms should be changed to negative infinity, and the 1 terms should be changed to zero. This is because the attention mask is treated as bias with addition operator.

There is no warning for this, so careless developer like me can just use the mask without knowing the side effect.

Your contribution

I'm not sure if I can make a PR for this case because I haven't explored how to apply the warning in general. What I found right now is just a specific case, and I hope it can help.