causal_mask in GPT2Attention should not be broadcastable across the seq_len
Author: bknyaz
Pull request: no
Labels: name: Feature request

### System Info

Python             : 3.8.2
torch                : 2.2.0.dev20231207+cu121
transformers    : 4.31.0
torchvision       : 0.17.0.dev20231207+cu121
cuda version    : 12.1

In `transformers.models.gpt2.modeling_gpt2.GPT2Attention` 
https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L202

the `causal_mask` must have the same shape for the last 2 dims, otherwise if the `max_position_embeddings=1` while the sequence length is longer than 1, the resulted attention weights leads to attending the future tokens. See the steps to reproduce the behavior for details. Normally, one wouldn't set `max_position_embeddings=1`, but nevertheless the broadcasting should not happen.

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

Code to reproduce the issue:

```
import torch
import transformers
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

bsz, seq_len, hid = 2, 3, 4
fig, axes = plt.subplots(ncols=3, figsize=(9,2))
for n_positions, ax in zip([1, 2, seq_len], axes):
    attn = transformers.models.gpt2.modeling_gpt2.GPT2Attention(transformers.GPT2Config(n_embd=hid, 
                                                                                        n_layer=1,
                                                                                        n_head=1,
                                                                                        n_positions=n_positions))
    ax.axis(False)
    ax.set_title('attn_weights, n_positions=%d' % n_positions, fontsize=9)
    attn_input = torch.randn(bsz, seq_len, hid)
    try:
        attn_output, _, attn_weights = attn(attn_input, output_attentions=True)
    except Exception as e:
        print('n_positions=%d' % n_positions, attn_input.shape, 'attn_output', 'ERROR:', e)
        continue
    print('n_positions=%d' % n_positions, 'attn_input', attn_input.shape, 'attn_output', attn_output.shape, 'attn_weights', attn_weights.shape)
    
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.05)
    im = ax.imshow(attn_weights[0, 0].data.cpu().numpy())
    fig.colorbar(im, cax=cax, orientation='vertical')
plt.show()    
```

Output:

```
n_positions=1 attn_input torch.Size([2, 3, 4]) attn_output torch.Size([2, 3, 4]) attn_weights torch.Size([2, 1, 3, 3])
n_positions=2 torch.Size([2, 3, 4]) attn_output ERROR: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 3
n_positions=3 attn_input torch.Size([2, 3, 4]) attn_output torch.Size([2, 3, 4]) attn_weights torch.Size([2, 1, 3, 3])
```

![image](https://github.com/huggingface/transformers/assets/3225366/c644271f-b3c7-44ce-b236-cc857597708e)


### Expected behavior

There should be some error message, for example triggered by 

`assert attn_weights.shape[-2:] == causal_mask.shape[-2:], 'attn_weights and causal_mask must have the same seq length'`

Comments:
cc @ArthurZucker 

Hey! The GPT2Attention layer is not part of the public documentation and is not self contained. This is thus expected. You are in luck as #28132 will probably adresse this.