forward pass consumes 2x memory than required in Attention module when `use_cache=True`

[RahulSChand]

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).

transformers/src/transformers/models/llama/modeling_llama.py

Lines 335 to 338 in 0afa507

	if past_key_value is not None:
	# reuse k, v, self_attention
	key_states = torch.cat([past_key_value[0], key_states], dim=2)
	value_states = torch.cat([past_key_value[1], value_states], dim=2)

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

transformers/src/transformers/models/llama/modeling_llama.py

Lines 335 to 338 in 0afa507

	if past_key_value is not None:
	# reuse k, v, self_attention
	key_states = torch.cat([past_key_value[0], key_states], dim=2)
	value_states = torch.cat([past_key_value[1], value_states], dim=2)

[amyeroberts]

cc @ArthurZucker

[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:

        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.

[RahulSChand]

@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).

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
)

[ArthurZucker]

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

[github-actions]

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 are likely to be ignored.