support copies

[ArthurZucker]

What does this PR do?

We can't copy the cache 😢 inheriting from module fixes this easily
This renders us unable to re-use prompts / system prompt like this:

import os, torch, copy
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
device = "cuda"
ckpt = "path-to-ckpt"
INITIAL_PROMPT = "From now on, you are going to answer all my questions with historical details. Make sure to always add a bit of french here and there, for style."

model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16)
model.to(device)

tokenizer = AutoTokenizer.from_pretrained(ckpt)

prompt_cache = DynamicCache()
inputs = tokenizer(INITIAL_PROMPT, return_tensors="pt").to("cuda")
prompt_cache = model(**inputs, past_key_values = prompt_cache).past_key_values


prompt = "Why are french people obsessed with french?"
new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")
past_key_values = copy.deepcopy(prompt_cache)
outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20) 
response = tokenizer.batch_decode(outputs)[0]
print(response)
"""
"""

prompt = "What is the best city to swim in?"
new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")
outputs = model.generate(**new_inputs, past_key_values=copy.deepcopy(prompt_cache),max_new_tokens=20) 
response = tokenizer.batch_decode(outputs)[0]
print(response)
"""

[amyeroberts]

We can't copy the cache 😢

What kind of copying are we talking about here? Like cache.copy?

[gante]

@amyeroberts copy.deepcopy

On main, without the fix, we get

RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment.  If you were attempting to deepcopy a module, this may be because of a torch.nn.utils.weight_norm usage, see https://github.com/pytorch/pytorch/pull/103001

Cache copying is needed to reuse the cache from the prompt. E.g. to run new prompts on top of the system prompt without spending compute on the system prompt.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vladfaust]

I'm sorry if it's not the right place to ask this question, but.

In Llama.cpp it's trivial to save and load state to/from disk to maintain the cache between sessions. Is it currently possible with Transformers, and if yes, could you please provide a minimal example or point to docs?

Cheers,

[gante]

@vladfaust yes it is possible, but it requires custom code (i.e. you would need to store and restore the cache's tensors).

We will add a user-friendly API for that in the future :)

[nirbenda]

LGTM

[ArthurZucker]

Ps this was actually already merged in #32168 so I'll close this one!