GPT2LMHeadModel not working with mps device, RuntimeError: tensors must be 2-D
Author: AlexVialaBellander
Pull request: no

### System Info

- `transformers` version: 4.24.0
- Platform: macOS-13.0-arm64-arm-64bit
- Python version: 3.9.13
- Huggingface_hub version: 0.11.0
- PyTorch version (GPU?): 1.12.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes, (mps)
- Using distributed or parallel set-up in script?: no

### Who can help?

@patil-suraj, @patrickvonplaten, @LysandreJik

### 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

Load the model and add to `mps` device

```python
from transformers import GPT2LMHeadModel

model = GPT2LMHeadModel.from_pretrained("gpt2-large").to("mps").eval()
```
Then run the model with some sample random input.

```python 
test = torch.randint(0, 100, (1, 10)).to("mps")
predictions = model(input_ids=test)
```

Then an error is thrown
```python
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In [27], line 1
----> 1 predictions = model(input_ids=test)
      3 predictions.logits

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py:1046, in GPT2LMHeadModel.forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1038 r"""
   1039 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
   1040     Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
   1041     `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
   1042     are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
   1043 """
   1044 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-> 1046 transformer_outputs = self.transformer(
   1047     input_ids,
   1048     past_key_values=past_key_values,
   1049     attention_mask=attention_mask,
   1050     token_type_ids=token_type_ids,
   1051     position_ids=position_ids,
   1052     head_mask=head_mask,
   1053     inputs_embeds=inputs_embeds,
   1054     encoder_hidden_states=encoder_hidden_states,
   1055     encoder_attention_mask=encoder_attention_mask,
   1056     use_cache=use_cache,
   1057     output_attentions=output_attentions,
   1058     output_hidden_states=output_hidden_states,
   1059     return_dict=return_dict,
   1060 )
   1061 hidden_states = transformer_outputs[0]
   1063 # Set device for model parallelism

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py:889, in GPT2Model.forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions, output_hidden_states, return_dict)
    879     outputs = torch.utils.checkpoint.checkpoint(
    880         create_custom_forward(block),
    881         hidden_states,
   (...)
    886         encoder_attention_mask,
    887     )
    888 else:
--> 889     outputs = block(
    890         hidden_states,
    891         layer_past=layer_past,
    892         attention_mask=attention_mask,
    893         head_mask=head_mask[i],
    894         encoder_hidden_states=encoder_hidden_states,
    895         encoder_attention_mask=encoder_attention_mask,
    896         use_cache=use_cache,
    897         output_attentions=output_attentions,
    898     )
    900 hidden_states = outputs[0]
    901 if use_cache is True:

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py:389, in GPT2Block.forward(self, hidden_states, layer_past, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions)
    387 residual = hidden_states
    388 hidden_states = self.ln_1(hidden_states)
--> 389 attn_outputs = self.attn(
    390     hidden_states,
    391     layer_past=layer_past,
    392     attention_mask=attention_mask,
    393     head_mask=head_mask,
    394     use_cache=use_cache,
    395     output_attentions=output_attentions,
    396 )
    397 attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
    398 outputs = attn_outputs[1:]

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py:311, in GPT2Attention.forward(self, hidden_states, layer_past, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions)
    309     attention_mask = encoder_attention_mask
    310 else:
--> 311     query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
    313 query = self._split_heads(query, self.num_heads, self.head_dim)
    314 key = self._split_heads(key, self.num_heads, self.head_dim)

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/opt/miniconda3/envs/torch/lib/python3.9/site-packages/transformers/pytorch_utils.py:112, in Conv1D.forward(self, x)
    110 def forward(self, x):
    111     size_out = x.size()[:-1] + (self.nf,)
--> 112     x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
    113     x = x.view(size_out)
    114     return x

RuntimeError: tensors must be 2-D
```

### Expected behavior

If we change the model to device `cuda` or `cpu`:

```python
from transformers import GPT2LMHeadModel

model = GPT2LMHeadModel.from_pretrained("gpt2-large").to("cpu").eval()

test = torch.randint(0, 100, (1, 10)).to("cpu")
predictions = model(input_ids=test)
```

This works and produces a prediction vector.

Comments:
Looks like this comes from an op not implemented in PyTorch. I'd try upgrading PyTorch (even to the nigthlies) as support for MPS is still in progress on their side.

Thanks @sgugger, sorry for wasting your time. I had pytorch 1.12.1 and 1.13 worked fine.