Abnormally High GPU Memory Consumption with OPT 350M Model Leading to OOM
Author: ayaka14732
Pull request: no
Labels: name: 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

Comments:
You are using a very high sequence length and are using the model without `torch.no_grad` so intermediate activations (which are saved for the backward pass) take a lot of space. That might be the reason.

Thank you @sgugger! However, I still think the memory usage should be low as the model is only 350M in size, and I am using a GPU with 80GB memory. Is there any formula that I can calculate the expected memory usage?

Besides, I am going to train the model, so I need the gradients and cannot use `torch.no_grad()`.

This is the reason why I think the high memory usage is abnormal:

```
OPTForCausalLM(
  (model): OPTModel(
    (decoder): OPTDecoder(
      (embed_tokens): Embedding(50272, 512, padding_idx=1)
      (embed_positions): OPTLearnedPositionalEmbedding(2050, 1024)
      (project_out): Linear(in_features=1024, out_features=512, bias=False)
      (project_in): Linear(in_features=512, out_features=1024, bias=False)
      (layers): ModuleList(
        (0-23): 24 x OPTDecoderLayer(
          (self_attn): OPTAttention(
            (k_proj): Linear(in_features=1024, out_features=1024, bias=True)
            (v_proj): Linear(in_features=1024, out_features=1024, bias=True)
            (q_proj): Linear(in_features=1024, out_features=1024, bias=True)
            (out_proj): Linear(in_features=1024, out_features=1024, bias=True)
          )
          (activation_fn): ReLU()
          (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
          (fc1): Linear(in_features=1024, out_features=4096, bias=True)
          (fc2): Linear(in_features=4096, out_features=1024, bias=True)
          (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        )
      )
    )
  )
  (lm_head): Linear(in_features=512, out_features=50272, bias=False)
)
```

- B: batch size = 6
- K: d_k, d_v = 64
- F: d_ff = 4096
- M: d_model = 1024
- H: n_heads = 16
- L: seq_len = 2048 + 2 = 2050
- C: vocab_size = 50272
- N: n_layers = 24
- P: d_proj = 512

Total memory = Model + Gradient + Activation

**Model** = embedding + embed pos + project in + layers + project_out + lm_head

embedding = CP
embed pos = LM
project in = PM
project out = MP
lm_head = PC

layers = N * layer
layer = attention + feed forward

attention = q proj + k proj + v proj + out proj + self_attn_layer_norm
q proj = MHK + HK
k proj = MHK + HK
v proj = MHK + HK
out proj = HKM + M
self_attn_layer_norm = 2M

feed forward = fc1 + fc2 + final_layer_norm
fc1 = MF + F
fc2 = FM + M
final_layer_norm = 2M

Model = CP + LM + PM + MP + PC + N (MHK + HK + MHK + HK + MHK + HK + HKM + M + M + MF + F + FM + M + M)
= 2CP + LM + 2PM + N (4MHK + 3HK + 6M + F + 2MF)
= 2 * 50272 * 512 + 2050 * 1024 + 2 * 512 * 1024 + 24 * (4 * 1024 * 16 * 64 + 3 * 16 * 64 + 6 * 1024 + 4096 + 2 * 1024 * 4096)
= 356935680

**Gradient** = Model

**Activation** = embedding + embed pos + project in + layers + project_out + lm_head

embedding = BLP
embed pos = BLM
project in = BLM
project out = BLP
lm_head = BLC

layers = N * layer
layer = attention + feed forward

attention = q proj + k proj + v proj + out proj + self_attn_layer_norm
q proj = 2BHLK
k proj = 2BHLK
v proj = 2BHLK
out proj = 2BLM
self_attn_layer_norm = 2BLM

feed forward = fc1 + fc2 + final_layer_norm
fc1 = 2BLF
fc2 = 2BLM
final_layer_norm = 2BLM

Activation = BLP + BLM + BLM + BLP + BLC + N (2BHLK + 2BHLK + 2BHLK + 2BLM + 2BLM + 2BLF + 2BLM + 2BLM)
= BL (2P + 2M + C + N (6HK + 8M + 2F))
= 6 * 2050 * (2 * 512 + 2 * 1024 + 50272 + 24 * (6 * 16 * 64 + 8 * 1024 + 2 * 4096))
= 7306396800

Total memory = Model + Gradient + Activation
= 2 * 356935680 + 7306396800
= 8020268160

8020268160 * 4 / (2 ** 30) = 29.9 GiB

I can reproduce the OOM. cc @ArthurZucker if you have any idea.

Indeed, for now removing `attention_scores = torch.max(attention_scores, torch.tensor(torch.finfo(attention_scores.dtype).min, device=attention_scores.device))` seems to free 20GB and a forward pass is properly running. Would suggest using `gradient_checkpointing` if you still keep this line. 
It's a bit strange 😅 

This still needs to be addressed.

Yes sorry did not have time to investigate further, might be related to past key values copies being held / memory not released. Something like #25930

This issue has not been fixed and still needs to be addressed.

Is this solved? I face the same issues in 2024


For OPT in particular, no but we have not seen such reports for other models