

TITLE
Irregular VRAM usage with gpt-neo inference with sequences longer than 250 tokens

COMMENTS
11

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.1 / HEAD
- Platform: Linux/Colab Pro
- Python version: 3.7
- PyTorch version (GPU?): 1.8.1 (CUDA 11.0)
- Tensorflow version (GPU?):
- Using GPU in script?: Yes, NVIDIA P100
- Using distributed or parallel set-up in script?:

### Who can help

@patil-suraj

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): EleutherAI/gpt-neo-2.7B

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Install transformers in a Colab Pro notebook
2. Run this script to log peak memory usage for inference with increasing sequence length: https://gist.github.com/finetuneanon/7ce0ed5090a27a383abffbbbc0433a29
3. Wait for it to crash with an OOM error in the attention matmul somewhere above sequence length 1850

Output:
```
1870 5436434432
ok 6535669248
1871 5436434432
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-2-f2aeed4489bd> in <module>()
     21     return_dict_in_generate=True,
     22     repetition_penalty=1.2,
---> 23     pad_token_id=tokenizer.eos_token_id
     24   )
     25   del ids

13 frames
/usr/local/lib/python3.7/dist-packages/transformers/models/gpt_neo/modeling_gpt_neo.py in _attn(self, query, key, value, causal_mask, masked_bias, attn_dropout, attention_mask, head_mask)
    238         key = key.to(torch.float32)
    239 
--> 240         attn_weights = torch.matmul(query, key.transpose(-1, -2))
    241         attn_weights = torch.where(causal_mask, attn_weights, masked_bias.to(attn_weights.dtype))
    242 

RuntimeError: CUDA out of memory. Tried to allocate 4.59 GiB (GPU 0; 15.90 GiB total capacity; 9.75 GiB already allocated; 4.60 GiB free; 10.42 GiB reserved in total by PyTorch)
```

The full output can be found here: https://gist.github.com/finetuneanon/c7292ea676f57f5bb63803685d80bf5b

The output has the format:
```
sequence_length occupied_cuda_memory_before_inference
ok peak_occupied_cuda_memory_during_inference
```

Doing inference with real text has the same issue.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
I expected memory usage to increase steadily instead of jumping around wildly, but I am not sure if this might actually be the correct behaviour. If it is correct, reliably doing inference on long sequences on 16GB of VRAM seems to be impossible, but sometimes it works.

I have also plotted the peak memory allocation during inference:

![cudaoom](https://user-images.githubusercontent.com/82650881/115268890-0a20db80-a13b-11eb-8771-60b47a5f66bb.png)

The green line is peak memory allocation, the brown line is the amount of memory in use before running inference.