forward pass consumes 2x memory than required in Attention module when `use_cache=True`
Author: RahulSChand
Pull request: no

### Feature request

Currently the forward pass of the model consumes 2x memory than is required in attention module when `use_cache=True`. 

For example during generation, in the attention module of llama `key_states` is of shape `(bsz, heads, 1, head_size)`
 & `past_key_value[0]` is of shape `(bsz, heads, seq_length, head_size)`. 

When we do `torch.cat(..)` in line 337 below, we get two copies, `past_key_value[0]` one of shape `(bsz, heads, seq_length, head_size)` and the `key_states` of shape `(bsz, heads, seq_length+1, head_size)`. 

https://github.com/huggingface/transformers/blob/0afa5071bd84e44301750fdc594e33db102cf374/src/transformers/models/llama/modeling_llama.py#L335-L338

Ideally we want the tensor that `past_key_value[0]` is pointing to be freed since it's no longer used (& will be replaced by the newly created `key_states`). In current implementation if  `(bsz=1,heads=12,seq_length=1024, head_size=128)` then  memory consumed is `bsz*heads*seq_length*head_size*layers*2` whereas it should ideally only use half of it. This can be achieved by just freeing the  `past_key_value[0]` after cat operation finishes. 

This is particularly noticable when you increase `bsz`.  Freeing the memory allows using 2x more batch size or sequence length. There are some edge cases where you might use  the `past_key_value[0]` again so maybe there can be flag to switch this on/off. 



### Motivation

Reduce memory consumed. by 2x when `use_cache=True`. This allows us to increase batch size by 2x or max seq length of tokens generated by 2x with same memory.

### Your contribution

I can contribute if this is a change that is required. The change is small on surface since all it requires is freeing the tensor that `past_key_value[0]` & `past_key_value[1]` point to after Line 338 

https://github.com/huggingface/transformers/blob/0afa5071bd84e44301750fdc594e33db102cf374/src/transformers/models/llama/modeling_llama.py#L335-L338

Comments:
cc @ArthurZucker 

Hey! I am not entirely sure what you are asking for is possible / an issue. 
I tested locally with `use_cache = True` for a generation and did not see any improvement. The reason is that `key_states` are used afterwards (repeated for example when the `num_kv_heads` is different than the number of heads, while the **copy** is kept to store the past key values. We could store the full past key values ( which would increase the RAM throughout the execution vs when computing the states), but then again, I did not see any performance improvements when switching to using:
```pytho
        if past_key_value is not None:
            # reuse k, v, self_attention
            past_key_value = (torch.cat([past_key_value[0], key_states], dim=2), torch.cat([past_key_value[1], value_states], dim=2))
        if use_cache:
            key_states = past_key_value[0]
            value_states = past_key_value[1]
```
where I remove the copy. 

Freeing the tensors should be automatically handle by python's garbage collector, especially when overwriting a variable. 

Do you have a reproducing snippet with expected RAM usage, with a given batch size etx? I tried with batch size of 64 and notices no differences. 

@ArthurZucker  This is a simple example where kv cache takes 2x the memory than it needs to.

The total GPU memory should be = model size + cuda overhead + kv cache

In below example, kv_cache should take `1000(tokens)*4096*32(layers)*2(kv)*2(float16 is 2 bytes) / (1024*1024)` MB = `500MB`. but it takes 1GB gpu memory because we store 2 copies for all 32 layers. That is, if we are generating the 1000th token then we have two kv_cache, one of shape (2, 32, 1000, 4096) & another of size (2, 32, 999, 4096).

You can either confirm this by running `nvidia-smi` or printing `torch.cuda.max_alloacted()` in the forward pass of `LlamaModel`. You can also change the below code to only have two total tokens & then 1000 total tokens & see that the diff b/w the 2 cases is 1GB (it should ideally be only 500MB).



```python
import torch
from transformers import LlamaTokenizer, DataCollatorWithPadding, LlamaForCausalLM

path = "meta-llama/Llama-2-7b"
tokenizer = LlamaTokenizer.from_pretrained(path)
tokenizer.pad_token = tokenizer.eos_token

model = LlamaForCausalLM.from_pretrained(path, low_cpu_mem_usage = True, device_map = "auto", torch_dtype = torch.float16)
word = "Hello "*500
word = [word]

inputs = tokenizer(word, return_tensors="pt").to(model.device)

generated_ids = model.generate(
        **inputs,
        use_cache=True,
        max_new_tokens=500,
        temperature=0.2,
        top_p=0.95,
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,  # model has no pad token
)
```

I could not really find conclusive evidence of this when running the script. The maximum allocated memory does not always take into account what is released and with `nvidia-smi` I was getting the correct memory usage. Given that we do `past_key_value = (key_states, value_states) if use_cache else None` I would not expect the tensor to be kept. The best would be to overwrite the tuple but we can't for now. We plan on refactoring the cache so this could be fixed by #26681

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.