[WIP] Enable GPT2Model to handle 3d attention_mask
Author: bzantium
Pull request: yes

This PR solves the problem discussed at #12261.

@patrickvonplaten, @LysandreJik

Comments:
Hello, thank you for offering a fix! It seems this proposal would be breaking for the GPT-2 double heads model; the following test fails: `test_gpt2_double_lm_head_model` with the following error:

```
_________________ GPT2ModelTest.test_gpt2_double_lm_head_model _________________
[gw1] linux -- Python 3.7.10 /usr/local/bin/python

self = <tests.test_modeling_gpt2.GPT2ModelTest testMethod=test_gpt2_double_lm_head_model>

    def test_gpt2_double_lm_head_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
>       self.model_tester.create_and_check_double_lm_head_model(*config_and_inputs)

tests/test_modeling_gpt2.py:457: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/test_modeling_gpt2.py:350: in create_and_check_double_lm_head_model
    result = model(**inputs)
../.local/lib/python3.7/site-packages/torch/nn/modules/module.py:1051: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/gpt2/modeling_gpt2.py:1160: in forward
    return_dict=return_dict,
../.local/lib/python3.7/site-packages/torch/nn/modules/module.py:1051: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/gpt2/modeling_gpt2.py:802: in forward
    output_attentions=output_attentions,
../.local/lib/python3.7/site-packages/torch/nn/modules/module.py:1051: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/gpt2/modeling_gpt2.py:323: in forward
    output_attentions=output_attentions,
../.local/lib/python3.7/site-packages/torch/nn/modules/module.py:1051: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/gpt2/modeling_gpt2.py:258: in forward
    attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = GPT2Attention(
  (c_attn): Conv1D()
  (c_proj): Conv1D()
  (attn_dropout): Dropout(p=0.1, inplace=False)
  (resid_dropout): Dropout(p=0.1, inplace=False)
)
query = tensor([[[[ 0.0202, -0.0295, -0.0911,  ..., -0.1715,  0.0008,  0.0310],
          [-0.0724, -0.0315, -0.0237,  ...,  0... -0.0802],
          [-0.1085,  0.0763, -0.0241,  ..., -0.1154,  0.1063, -0.1542]]]],
       grad_fn=<PermuteBackward>)
key = tensor([[[[ 2.2768e-02, -1.5131e-02, -3.1551e-02,  ...,  1.2214e-01,
            3.4581e-03,  1.2902e-01],
          [...6e-01,  1.0495e-01, -8.1176e-02,  ...,  1.5278e-01,
           -1.6426e-01,  4.7595e-02]]]], grad_fn=<PermuteBackward>)
value = tensor([[[[ 0.0396, -0.0556, -0.0115,  ...,  0.1020,  0.0598,  0.1249],
          [ 0.1337, -0.0851,  0.0792,  ...,  0...  0.0443],
          [-0.1049, -0.0717,  0.1128,  ...,  0.2006,  0.0411, -0.0256]]]],
       grad_fn=<PermuteBackward>)
attention_mask = tensor([[[[-10000.,     -0., -10000., -10000.,     -0.,     -0.,     -0.],
          [-10000.,     -0., -10000., -1000...00.,     -0.,     -0., -10000.,     -0.],
          [    -0.,     -0., -10000.,     -0.,     -0., -10000.,     -0.]]]])
head_mask = None

    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        attn_weights = torch.matmul(query, key.transpose(-1, -2))
    
        if self.scale_attn_weights:
            attn_weights = attn_weights / (float(value.size(-1)) ** 0.5)
    
        if not self.is_cross_attention:
            # if only "normal" attention layer implements causal mask
            query_length, key_length = query.size(-2), key.size(-2)
            causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
            attn_weights = torch.where(causal_mask, attn_weights, self.masked_bias.to(attn_weights.dtype))
    
        if attention_mask is not None:
            # Apply the attention mask
>           attn_weights = attn_weights + attention_mask
E           RuntimeError: The size of tensor a (7) must match the size of tensor b (4) at non-singleton dimension 2

src/transformers/models/gpt2/modeling_gpt2.py:191: RuntimeError
```

Could you give it a look?

@LysandreJik Thank you for the comment! I inspected code carefully and found that it was because `input_ids` from `create_and_check_double_lm_head_model` has extra dimension for `num_choices` and is merged to `batch` dimension before pushed to model. I fixed this error by adding the condition that `batch_size` of `input_ids` is equal to `batch_size` of `attention_mask` or attention_mask's dimension of `num_choices` would be merged to `batch` dimension as `input_ids`. Also, I added the same code for openai gpt. Please check again!

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/master/CONTRIBUTING.md) are likely to be ignored.