GPTNeoX with use_cache=False uses significantly more memory than use_cache=True
Author: dwyatte
Pull request: no

### System Info

- `transformers` version: 4.40.1
- Platform: Linux-4.19.0-25-cloud-amd64-x86_64-with-glibc2.35
- Python version: 3.9.18
- Huggingface_hub version: 0.21.3
- Safetensors version: 0.4.2
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.1+cu121 (True)
- 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
- Using distributed or parallel set-up in script?: no

### Who can help?

@gante  @ArthurZucker

### Information

- [ ] 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

When a GPTNeoX model is instantiated for generation with `use_cache=False`, it uses significantly more memory than when using the cache. This doesn't seem to be the case with a few other model architectures I've tried. Here's a simple repro and some stats from my environment

```
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "EleutherAI/pythia-410m"
use_cache = False
batch_size = 1
max_length = 2048

model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, use_cache=use_cache).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
tokenizer.pad_token = tokenizer.eos_token
    
inputs = tokenizer(["abc"]*batch_size, padding="max_length", max_length=max_length, return_tensors="pt").to("cuda")
generations = model.generate(**inputs, max_new_tokens=256, min_new_tokens=256)

free, total = torch.cuda.mem_get_info()
print(f"{(total - free) / 1024 / 1024} MB used")
```

* EleutherAI/pythia-410m
    * max_length = 2048, use_cache=True: 3643.0 MB used
    * **max_length = 2048, use_cache=False: 22919.0 MB used**

* facebook/opt-350m
    * max_length = 2048, use_cache=True: 2997.0 MB used
    * max_length = 2048, use_cache=False: 2997.0 MB used

* TinyLlama/TinyLlama-1.1B-Chat-v1.0
    * max_length = 2048, use_cache=True: 4335.0 MB used
    * max_length = 2048, use_cache=False: 4335.0 MB used

Note that I'm padding out to a large sequence length to illustrate the extent of the problem, but this issue manifests for shorter sequence lengths within the model's max length too

* EleutherAI/pythia-410m
    * max_length = 1024, use_cache=True: 2965.0 MB used
    * **max_length = 1024, use_cache=False: 5753.0 MB used**

### Expected behavior

In general, I'd expect memory with `use_cache=False` to be less than or equal to `use_cache=True`

Comments:
@dwyatte hey! 

Memory increase you see come from the fact that "use_cache=False" has to calculate qkv for the whole sequence length every timestep, while in "use_cache=True" the qkv are calculated for 1 token only (prev steps are in cache). So when you do `torch.cuda.mem_get_info()`, torch shows the total memory reserved by cuda, which does not reflect the memory allocated by the kv cache.

When I try the following, it can be seen that the `max_memory_allocated` for cached generation is a bit higher than uncached.
```python
begin = torch.cuda.memory_allocated("cuda:0")
generations = model.generate(**inputs, max_new_tokens=200, min_new_tokens=200)
print(f"{(torch.cuda.max_memory_allocated('cuda:0') - begin) // 1024 ** 2} MB used")
```

And regarding the reason why the reserved memory is high for GPTNeox and not other models, I did some digging into the GPTNeox implementation. All comes into the way how qkv proj is implemented. In GPTNeox it has one linear layer of x3 size to project all at once, while OPT and LLama have three separate projection layers. So the reason is cuda memory allocations management because the sequence length increases by 1 every step leaving unused free memory. Setting "expandable_segments:True" reduces the high memory you observed. See [here](https://pytorch.org/docs/stable/notes/cuda.html) for more on that 🤗 

`PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python run.py`

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.