Repository: transformers
Issue: Abnormally High GPU Memory Consumption with OPT 350M Model Leading to OOM
Labels: WIP

### System Info

- `transformers` version: 4.32.0.dev0
- Platform: Linux-5.4.0-135-generic-x86_64-with-glibc2.35
- Python version: 3.11.4
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0.dev20230723+cu121 (True)
- Tensorflow version (GPU?): 2.14.0-dev20230723 (False)
- Flax version (CPU?/GPU?/TPU?): 0.7.0 (gpu)
- Jax version: 0.4.14
- JaxLib version: 0.4.14
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @younesbelkada @sgugger

### Information

- [X] 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

While working with the OPT 350M model, I have encountered an issue regarding GPU memory consumption that I believe to be abnormal. Specifically, during the forward pass, the model causes an OOM error. I am using 1 Nvidia A100 80GB GPU.

Here is the repro with a manual memory profiling:

```python
import os; os.environ['CUDA_VISIBLE_DEVICES'] = '2'  # use 1 device only
import torch
from transformers import OPTForCausalLM

model_name = 'facebook/opt-350m'
model = OPTForCausalLM.from_pretrained(model_name, device_map='cuda')  # memory usage: 1688 MiB

batch_size = 6
seq_len = 2048

seq_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device='cuda')  # memory usage: 1764 MiB
seq_mask = torch.zeros((batch_size, seq_len), dtype=torch.bool, device='cuda')  # memory usage: 1764 MiB
labels_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device='cuda')  # memory usage: 1764 MiB

outputs = model(input_ids=seq_ids, attention_mask=seq_mask, labels=labels_ids)  # the memory usage surges and leads to OOM
```

Output:

```
Some weights of OPTForCausalLM were not initialized from the model checkpoint at facebook/opt-350m and are newly initialized: ['lm_head.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Traceback (most recent call last):
  File "/home/ayaka/test/1.py", line 18, in <module>
    outputs = model(input_ids=seq_ids, attention_mask=seq_mask, labels=labels_ids)  # the memory usage surges and leads to OOM
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/transformers/models/opt/modeling_opt.py", line 944, in forward
    outputs = self.model.decoder(
              ^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/transformers/models/opt/modeling_opt.py", line 710, in forward
    layer_outputs = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/transformers/models/opt/modeling_opt.py", line 330, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
                                                          ^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/test/venv/lib/python3.11/site-packages/transformers/models/opt/modeling_opt.py", line 223, in forward
    attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
                   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.50 GiB. GPU 0 has a total capacty of 79.35 GiB of which 662.62 MiB is free. Process 4083423 has 78.70 GiB memory in use. Of the allocated memory 77.96 GiB is allocated by PyTorch, and 258.22 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```


### Expected behavior

No OOM since the model is very small