CUDA OOM with RoPE scaling on an A100 (GPTQ and bnb)
Author: RonanKMcGovern
Pull request: no

### System Info

I'm running in Colab with an A100 (40 GB) and the latest version of transformers (pip install -q -U).

### Who can help?

@gante @kaiokendev

### Information

- [ ] The official example scripts
- [X] 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

I'm running into CUDA OOM errors, even with low rope settings like 1.0:
```
!pip install -q -U transformers peft accelerate optimum

!pip install auto-gptq --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/

import transformers
import torch

model_id = "TheBloke/CodeLlama-7B-Instruct-GPTQ"

from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    rope_scaling = {"type": "dynamic", "factor": scaling_factor}
    )
```

```
OutOfMemoryError                          Traceback (most recent call last)
[<ipython-input-16-3e71bb55dbdb>](https://localhost:8080/#) in on_button_clicked(b)
     76             print_wrapped(f'**{message["role"].capitalize()}**: {message["content"]}\n')
     77 
---> 78     assistant_response = generate_response([dialog_history]);
     79 
     80     # Re-enable the button, reset description and color after processing

11 frames
[/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py](https://localhost:8080/#) in forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)
    349 
    350         # upcast attention to fp32
--> 351         attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    352         attn_output = torch.matmul(attn_weights, value_states)
    353 

OutOfMemoryError: CUDA out of memory. Tried to allocate 9.11 GiB (GPU 0; 39.56 GiB total capacity; 32.09 GiB already allocated; 5.25 GiB free; 32.71 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

When running [generation](https://huggingface.co/docs/transformers/v4.18.0/en/main_classes/text_generation), I tried setting use_cache=False (to recompute the kv values for each token). I thought this would reduce VRAM requirements, but it seems to have increased the amount of VRAM requested. I may be misunderstanding something.

Similar issue loading with bnb nf4 with use_cache=False in the generation function call. Here's my model configuration for loading bnb nf4:
```
model_id='codellama/CodeLlama-7b-Instruct-hf'
scaling_factor=2.0
runtimeFlag="cuda:0"

# Load the model in 4-bit to allow it to fit in a free Google Colab runtime with a CPU and T4 GPU
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True, #adds speed with minimal loss of quality.
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    # device_map='auto', # for inference use 'auto', for training use device_map={"":0}
    device_map=runtimeFlag,
    trust_remote_code=True,
    rope_scaling = {"type": "dynamic", "factor": scaling_factor}, # allows for a max sequence length of 8192 tokens !!! [not tested in this notebook yet]
    # cache_dir=cache_dir
    )

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True) # will use the Rust fast tokenizer if available
```
which gave:
```
OutOfMemoryError                          Traceback (most recent call last)
[<ipython-input-24-3e71bb55dbdb>](https://localhost:8080/#) in on_button_clicked(b)
     76             print_wrapped(f'**{message["role"].capitalize()}**: {message["content"]}\n')
     77 
---> 78     assistant_response = generate_response([dialog_history]);
     79 
     80     # Re-enable the button, reset description and color after processing

15 frames
[/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py](https://localhost:8080/#) in forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)
    344         value_states = repeat_kv(value_states, self.num_key_value_groups)
    345 
--> 346         attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
    347 
    348         if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):

OutOfMemoryError: CUDA out of memory. Tried to allocate 36.22 GiB (GPU 0; 39.56 GiB total capacity; 5.96 GiB already allocated; 29.44 GiB free; 8.53 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

### Expected behavior

I suppose the kv cache is getting very big, but I would have expected it to fit at 16k context length, expecially if I set use_cache=False

Comments:
cc @ArthurZucker 

Hey @RonanKMcGovern 👋 

Yeah, sadly the memory requirements grow very quickly with the context length. `use_cache=True` should be used in general: the computational speedup is large, and the memory footprint penalty is small.

While this isn't a bug per se, I do have a suggestion for you: have you tried activating flash attention? You can follow [this guide ](https://huggingface.co/docs/transformers/v4.32.1/en/perf_infer_gpu_one#bettertransformer)to try it.

Many thanks @gante.

That got rid of the OOM error. Some notes on performance:

- With Llama 2-chat, I scaled to 4X and was able to put in about 12k tokens. The response was good for the first half and then started repeating.

- With code-llama-Instruct, even without rope scaling, it seems to get unstable at 16k context (7B model, the 13B model gets too close to the 40 GB VRAM I have).