大模型 参数 use_cache 怎么用？ 与 KV Cache 是什么关系？

大模型 参数 use_cache 怎么用？ 与 KV Cache 是什么关系？

flyfish

在使用 transformers 库的 generate 方法进行文本生成时，use_cache 主要用于控制是否使用缓存机制来优化生成过程。

use_cache=True

当 use_cache = True 时，模型会将之前步骤中计算得到的中间结果（如注意力机制中的键值对）缓存起来。在后续的生成步骤中，只需要对新生成的标记进行计算，并结合缓存的结果，就可以快速得到当前步骤的输出，避免了对整个序列的重复计算。

代码示例

一个简单的示例，展示了如何在 generate 方法中使用 use_cache 参数：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# 准备输入文本
input_text = "Once upon a time"
input_ids = tokenizer.encode(input_text, return_tensors="pt")

# 开启缓存机制进行文本生成
outputs = model.generate(input_ids, max_length=50, use_cache=True)

# 解码输出
output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(output_text)

use_cache=False

use_cache = True 通常能提高文本生成的速度，但在以下几种情况下，可能需要将 use_cache 设置为 False

1. 内存资源受限

当运行环境的内存资源有限时，开启缓存可能会导致内存不足的问题。因为缓存机制会存储模型在生成过程中的中间计算结果（如多头注意力机制中的键值对），随着生成长度的增加，缓存占用的内存会不断增大。如果设备的内存不足以容纳这些缓存数据，就会引发内存溢出错误，导致程序崩溃。

示例场景：在一些嵌入式设备或内存较小的服务器上运行模型进行文本生成时，如果发现开启缓存后程序频繁出现内存不足的错误，就可以考虑将 use_cache 设置为 False 以减少内存占用。

2. 需要动态修改输入序列

在某些情况下，可能需要在生成过程中动态修改输入序列，例如根据中间生成结果进行条件调整或插入新的提示信息。此时，缓存的中间结果可能不再适用于修改后的输入序列，继续使用缓存会导致计算结果错误。因此，需要将 use_cache 设置为 False，以确保每次计算都是基于最新的输入序列进行的。

示例场景：在交互式文本生成应用中，用户可能会在生成过程中随时输入新的提示或修改之前的输入。为了保证生成结果的准确性，就不能依赖缓存，而要将 use_cache 设为 False。

3. 对生成过程进行精细控制

当需要对生成过程进行精细控制，例如对每个生成步骤进行复杂的中间处理或分析时，缓存机制可能会干扰的操作。因为缓存会跳过一些中间计算步骤，使得难以获取完整的计算过程信息。将 use_cache 设置为 False 可以确保每个步骤都进行完整的计算，方便进行详细的调试和分析。

示例场景：研究人员在对模型的生成机制进行深入研究时，需要详细分析每个生成步骤的输出和中间状态。此时，关闭缓存可以让他们获取更全面的信息，有助于深入理解模型的工作原理。

4. 特定的生成策略不依赖缓存

某些生成策略本身并不依赖缓存机制，或者在使用这些策略时缓存带来的性能提升并不明显，甚至可能会增加额外的开销。在这种情况下，可以将 use_cache 设置为 False 以简化计算过程。

示例场景：一些简单的生成策略，如基于贪心搜索的单步生成，每个步骤的计算相对独立，缓存机制对其性能提升不大，此时可以关闭缓存以减少不必要的内存开销。

一个将 use_cache 设置为 False 的代码示例：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# 准备输入文本
input_text = "Once upon a time"
input_ids = tokenizer.encode(input_text, return_tensors="pt")

# 不使用缓存进行文本生成
outputs = model.generate(input_ids, max_length=50, use_cache=False)

# 解码输出
output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(output_text)

带缓存生成的最佳实践

高效的缓存对于优化各种生成任务（包括文本生成、翻译、摘要以及其他基于 Transformer 的应用）中的模型性能至关重要。有效的缓存有助于减少计算时间并提高响应速度，特别是在实时或资源密集型应用中。

Transformer 支持多种缓存方法，利用 “Cache” 类来抽象和管理缓存逻辑。本文档概述了使用这些类以实现性能和效率最大化的最佳实践

什么是缓存以及为什么要关注它

想象一下，正在和某人交谈，但每次回应时都得从头开始，而不是记住之前说过的话。这会很慢且效率低下，对吧？在 Transformer 模型的世界里，也有类似的概念，这就是键值缓存（KV Cache）发挥作用的地方。从现在起，将这个概念简称为 KV 缓存。

在自回归模型的生成过程中，需要使用 KV 缓存进行优化。在自回归模型中，模型逐词预测文本。这个过程可能会很慢，因为模型一次只能生成一个词，而且每个新的预测都依赖于之前的上下文。也就是说，要预测生成的第 1000 个词，需要前 999 个词的信息，这表现为对这些词的表示进行一些矩阵乘法运算。而要预测第 1001 个词，同样需要前 999 个词的相同信息，再加上第 1000 个词的额外信息。这时，键值缓存就用于优化顺序生成过程，通过存储之前的计算结果，以便在后续的词生成中复用，这样就无需再次计算。

更具体地说，键值缓存就像是这些生成模型的一个内存库，模型会为之前处理过的词存储来自自注意力层的键值对。通过存储这些信息，模型可以避免重复计算，而是从缓存中检索之前词的键和值。请注意，缓存只能在推理时使用，训练时应禁用，否则可能会导致意外错误。

深入探究：缓存对象在注意力机制中的工作原理

当在输入中使用缓存对象时，注意力模块会执行几个关键步骤，以无缝整合过去和现在的信息。

注意力模块会将当前的键值对与缓存中存储的过去的键值对进行拼接。这会得到形状为 (新生成词的长度, 过去键值对的长度 + 新生成词的长度) 的注意力权重。本质上，过去和当前的键值对会被组合起来计算注意力分数，确保模型同时考虑之前的上下文和新的输入。拼接后的键值对用于计算注意力分数，得到形状为 (新生成词的长度, 过去键值对的长度 + 新生成词的长度) 的注意力权重。

因此，当迭代调用 forward() 而不是 generate() 方法时，必须确保注意力掩码的形状与过去和当前键值对的组合长度相匹配。注意力掩码的形状应为 (批次大小, 过去键值对的长度 + 新生成词的长度)。当调用 generate() 方法时，这通常会在内部处理。如果想使用 Cache 类实现自己的生成循环，请考虑这一点，并准备好注意力掩码以包含当前和过去词的值。

在编写自己的生成循环时，需要了解一个重要的概念，即 cache_position。如果想通过调用 forward() 重用一个已经填充好的 Cache 对象，必须传入一个有效的 cache_position，它将指示输入在序列中的位置。请注意，cache_position 不受填充的影响，并且每个词总是会增加一个位置。例如，如果键/值缓存包含 10 个词（无论其中有多少是填充词），下一个词的缓存位置应该是 torch.tensor([10])。

一个如何实现自己的生成循环的示例：

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

 model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained(model_id)

 past_key_values = DynamicCache()
 messages = [{"role": "user", "content": "Hello, what's your name."}]
 inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

 generated_ids = inputs.input_ids
 cache_position = torch.arange(inputs.input_ids.shape[1], dtype=torch.int64, device=model.device)
 max_new_tokens = 10

 for _ in range(max_new_tokens):
     outputs = model(**inputs, cache_position=cache_position, past_key_values=past_key_values, use_cache=True)
     # 贪心采样下一个词
     next_token_ids = outputs.logits[:, -1:].argmax(-1)
     generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)

     # 为下一个生成步骤准备输入，只保留未处理的词，在例子中只有一个新词
     # 并为新词扩展注意力掩码，如上所述
     attention_mask = inputs["attention_mask"]
     attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
     inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}
     cache_position = cache_position[-1:] + 1 # 为下一个词增加一个位置

 print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])

<|user|>
Hello, what's your name. 
<|assistant|>
My name is Sarah. 
<|

使用缓存进行生成

在 Transformers 中，支持多种缓存类型，以优化不同模型和任务的性能。默认情况下，所有模型都会使用缓存进行生成，[~DynamicCache] 类是大多数模型的默认缓存。它允许动态增加缓存大小，在生成过程中保存越来越多的键和值。如果出于某种原因不想使用缓存，可以在 generate() 方法中传入 use_cache=False。

参考下表，了解不同缓存类型之间的差异，并选择最适合用例的缓存类型。建议在调用模型之前初始化需要初始化的模型，并将其作为关键字参数传递给模型。在其他情况下，只需定义所需的 cache_implementation，会为处理其余的事情。

缓存类型	内存效率	支持 torch.compile()	建议初始化	延迟	长上下文生成
动态缓存（Dynamic Cache）	否	否	否	中等	否
静态缓存（Static Cache）	否	是	是	高	否
卸载缓存（Offloaded Cache）	是	否	否	低	是
卸载静态缓存（Offloaded Static Cache）	否	是	是	高	是
量化缓存（Quantized Cache）	是	否	否	低	是
滑动窗口缓存（Sliding Window Cache）	否	是	是	高	否
汇聚缓存（Sink Cache）	是	否	是	中等	是

这些缓存类可以在生成时通过 cache_implementation 参数进行设置。

量化缓存

键值缓存可能会占用大量内存，成为 长上下文生成的瓶颈，特别是对于大语言模型。在使用 generate() 时对缓存进行量化可以显著减少内存需求，但会牺牲一定的速度。

transformers 中的 KV 缓存量化在很大程度上受到论文 “KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache” 的启发，目前支持 [~QuantoQuantizedCache] 和 [~HQQQuantizedCache] 类。有关内部工作原理的更多信息，请参阅该论文。

要启用键值缓存的量化，需要在 generation_config 中指定 cache_implementation="quantized"。量化相关的参数应作为 dict 或 [~QuantizedCacheConfig] 类的实例传递给 generation_config。必须在 [~QuantizedCacheConfig] 中指定要使用的量化后端，默认是 quanto。

如果使用的是 quanto 后端，建议将缓存配置中的 axis - key/axis - value 参数设置为 0；如果使用的是 HQQ 后端，则设置为 1。对于其他配置值，除非内存不足，否则请使用默认值。在内存不足的情况下，可以考虑减小残差长度。

缓存量化在上下文长度较短且有足够的 GPU 显存可在不进行缓存量化的情况下运行时，可能会导致延迟增加。建议在内存效率和延迟之间寻求平衡。

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM

 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.float16, device_map="auto")
 inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

 out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"nbits": 4, "backend": "quanto"})
 print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's a great way to express myself. I like the way it makes me feel, the

卸载缓存

与 KV 缓存量化类似，[~OffloadedCache] 策略旨在减少 GPU 显存的使用。它通过将大多数层的 KV 缓存移动到 CPU 来实现这一点。当模型的 forward() 方法遍历各层时，该策略会将当前层的缓存保留在 GPU 上。同时，它会异步预取下一层的缓存，并将上一层的缓存发送回 CPU。与 KV 缓存量化不同，该策略总是产生与默认 KV 缓存实现相同的结果。因此，它可以作为默认实现的直接替代品或备用方案。

根据的模型和生成任务的特点（上下文大小、生成词的数量、束搜索的束数等），可能会注意到与默认 KV 缓存实现相比，生成吞吐量会有小幅下降。

要启用 KV 缓存卸载，在 generation_config 中传递 cache_implementation="offloaded"，或者直接在 generate() 调用中传递。使用 cache_implementation="offloaded_static" 可使用卸载的静态缓存。

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 ckpt = "microsoft/Phi-3-mini-4k-instruct"

 tokenizer = AutoTokenizer.from_pretrained(ckpt)
 model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16, device_map="auto")
 inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)

 out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
 print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

 out = model.generate(**inputs, do_sample=False, max_new_tokens=23)
 print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

缓存卸载需要 CUDA GPU，并且可能比动态 KV 缓存慢。如果遇到 CUDA 内存不足的错误，可以使用它。

以下示例展示了如何将 KV 缓存卸载作为备用策略使用：

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 def resilient_generate(model, *args, **kwargs):
     oom = False
     try:
         return model.generate(*args, **kwargs)
     except torch.cuda.OutOfMemoryError as e:
         print(e)
         print("retrying with cache_implementation='offloaded'")
         oom = True
     if oom:
         torch.cuda.empty_cache()
         kwargs["cache_implementation"] = "offloaded"
         return model.generate(*args, **kwargs)


 ckpt = "microsoft/Phi-3-mini-4k-instruct"
 tokenizer = AutoTokenizer.from_pretrained(ckpt)
 model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda:0")
 prompt = ["okay "*1000 + "Fun fact: The most"]
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 beams = { "num_beams": 40, "num_beam_groups": 40, "num_return_sequences": 40, "diversity_penalty": 1.0, "max_new_tokens": 23, "early_stopping": True, }
 out = resilient_generate(model, **inputs, **beams)
 responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

在具有 50 GB 显存的 GPU 上运行此代码时，会在成功生成 40 个束之前打印以下内容：

CUDA out of memory. Tried to allocate 4.83 GiB. GPU
retrying with cache_implementation='offloaded'

静态缓存

由于 “DynamicCache” 会在每个生成步骤中动态增长，因此它无法让利用 JIT 优化。[~StaticCache] 会为键和值预分配一个特定的最大大小，允许在不修改缓存大小的情况下生成到最大长度。

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM

 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.float16, device_map="auto")
 inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

 # 只需传递 cache_implementation="static"
 out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="static")
 tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name] and I am a [Your Position] at [Your Company]. I am writing"

卸载静态缓存

就像存在用于卸载 “DynamicCache” 的 [~OffloadedCache] 一样，也有卸载的静态缓存。它完全支持 JIT 优化。只需在 generation_config 中传递 cache_implementation="offloaded_static"，或者直接在 generate() 调用中传递。这将使用 [~OffloadedStaticCache] 实现。

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM

 tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
 model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", torch_dtype=torch.float16, device_map="auto")
 inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

 # 只需传递 cache_implementation="offloaded_static"
 out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="offloaded_static")
 tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

缓存卸载需要 CUDA GPU。

滑动窗口缓存

顾名思义，这种缓存类型会在之前的键和值上实现一个滑动窗口，只保留最后 sliding_window 个词。它应该与支持滑动窗口注意力的模型（如 Mistral

滑动窗口缓存

顾名思义，这种缓存类型会在之前的键和值上实现一个滑动窗口，只保留最后 sliding_window 个词。它应该与支持滑动窗口注意力的模型（如 Mistral 模型）配合使用。此外，与静态缓存类似，这种缓存对 JIT 友好，可以使用与静态缓存相同的编译技术。

请注意，只能将此缓存用于支持滑动窗口的模型，例如 Mistral 模型。

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, SinkCache

 tokenizer = AutoTokenizer.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B")
 model = AutoModelForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16, device_map="auto")
 inputs = tokenizer("Yesterday I was on a rock concert and.", return_tensors="pt").to(model.device)

 # 可以通过传入 cache_implementation 使用
 out = model.generate(**inputs, do_sample=False, max_new_tokens=30, cache_implementation="sliding_window")
 tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Yesterday I was on a rock concert and. I was so excited to see my favorite band perform live. I was so happy that I could hardly contain myself. I was jumping up and down and"

汇聚缓存

汇聚缓存（Sink Cache）是在论文 “Efficient Streaming Language Models with Attention Sinks” 中提出的。它允许在不进行任何微调的情况下生成长文本序列（根据论文，是“无限长度”）。这是通过巧妙处理之前的键和值来实现的，具体来说，它会保留序列中的几个初始词，称为“汇聚词（sink tokens）”。这是基于这样的观察：在生成过程中，这些初始词会吸引很大一部分注意力分数。“汇聚词”之后的词会以滑动窗口的方式丢弃，只保留最新的 window_size 个词。通过将这些初始词作为“注意力汇聚点”，模型即使在处理非常长的文本时也能保持稳定的性能，从而丢弃大部分之前的知识。

与其他缓存类不同，不能通过指定 cache_implementation 直接使用这种缓存。必须在调用 generate() 之前初始化缓存，如下所示：

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, SinkCache

 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.float16, device_map="auto")
 inputs = tokenizer("This is a long story about unicorns, fairies and magic.", return_tensors="pt").to(model.device)

 # 获取缓存，指定汇聚词的数量和窗口大小
 # 请注意，窗口大小已经包含了汇聚词，因此必须更大
 past_key_values = SinkCache(window_length=256, num_sink_tokens=4)
 out = model.generate(**inputs, do_sample=False, max_new_tokens=30, past_key_values=past_key_values)
 tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"This is a long story about unicorns, fairies and magic. It is a story about a young girl named Lily who discovers that she has the power to control the elements. She learns that she can"

编码器 - 解码器缓存

[~EncoderDecoderCache] 是一个包装器，旨在处理编码器 - 解码器模型的缓存需求。这种缓存类型专门用于管理自注意力和交叉注意力缓存，确保存储和检索这些复杂模型所需的过去键/值。编码器 - 解码器缓存的一个很棒的特性是，可以根据自己的用例为编码器和解码器设置不同的缓存类型。目前，这种缓存仅在 Whisper 模型中得到支持。

在使用方面，没有什么特别的操作，调用 generate() 或 forward() 会为处理一切。

特定模型的缓存类

有些模型需要以特定的方式存储之前的键、值或状态，上述缓存类无法使用。对于这种情况，有几个专门为特定模型设计的缓存类。这些模型只接受它们自己专用的缓存类，不支持使用其他任何缓存类型。一些例子包括适用于 Gemma2 系列模型的 [~HybridCache] 或适用于 Mamba 架构模型的 [~MambaCache]。

使用缓存进行迭代生成

已经了解了在生成时如何使用各种缓存类型。如果想在迭代生成的场景中使用缓存，例如在聊天机器人等应用中，交互涉及多个回合和持续的来回交流，该怎么办呢？使用缓存进行迭代生成可以让这些系统有效地处理持续的对话，而无需在每个步骤重新处理整个上下文。但在开始实现之前，应该了解一些技巧：

迭代生成的一般格式如下。首先，必须初始化一个想要的类型的空缓存，然后就可以开始迭代地输入新的提示。可以使用聊天模板来跟踪对话历史和进行格式化。

如果使用的是汇聚缓存 Sink Cache ，必须将输入裁剪到最大长度，因为汇聚缓存 Sink Cache 可以生成比其最大窗口大小更长的文本，但它要求第一个输入不超过最大缓存长度。

 import torch
 from transformers import AutoTokenizer,AutoModelForCausalLM
 from transformers.cache_utils import (
    DynamicCache,
    SinkCache,
    StaticCache,
    SlidingWindowCache,
    QuantoQuantizedCache,
    QuantizedCacheConfig,
 )

 model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map='auto')
 tokenizer = AutoTokenizer.from_pretrained(model_id)

 user_prompts = ["Hello, what's your name?", "Btw, yesterday I was on a rock concert."]

 past_key_values = DynamicCache()
 max_cache_length = past_key_values.get_max_cache_shape()

 messages = []
 for prompt in user_prompts:
     messages.append({"role": "user", "content": prompt})
     inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
     if isinstance(past_key_values, SinkCache):
         inputs = {k: v[:, -max_cache_length:] for k, v in inputs.items()}

     input_length = inputs["input_ids"].shape[1]

     outputs = model.generate(**inputs, do_sample=False, max_new_tokens=256, past_key_values=past_key_values)
     completion = tokenizer.decode(outputs[0, input_length: ], skip_special_tokens=True)
     messages.append({"role": "assistant", "content": completion})

print(messages)
[{'role': 'user', 'content': "Hello, what's your name?"}, {'role': 'assistant', 'content': "Hello, I'm AI."}, {'role': 'user', 'content': 'Btw, yesterday I was on a rock concert.'}, {'role': 'assistant', 'content': "I'm sorry to hear that you were on a rock concert yesterday. It sounds like a fun experience, but I'm not capable of experiencing music or concerts. However, I can provide you with some information about rock music and its history. Rock music emerged in the 1950s and 1960s in the United States and Britain, and it quickly gained popularity around the world. Some of the most famous rock bands of all time include The Beatles, The Rolling Stones, Led Zeppelin, and Pink Floyd. Rock music has a distinct sound and style, with elements of blues, country, and folk music. It often features guitar solos, heavy bass lines, and drums. Rock music has had a significant impact on popular culture, influencing genres such as punk rock, heavy metal, and alternative rock."}]

重用缓存以继续生成

有时，可能想先为某个前缀提示填充缓存对象的键/值，然后多次重用它，从该提示生成不同的序列。在这种情况下，可以构建一个 Cache 对象来保存指令提示，然后多次重用它，生成不同的文本序列。

 import copy
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache
 from accelerate.test_utils.testing import get_backend

 DEVICE, _, _ = get_backend() # 自动检测底层设备类型（CUDA、CPU、XPU、MPS 等）
 model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map=DEVICE)
 tokenizer = AutoTokenizer.from_pretrained(model_id)

 # 用足够大的最大长度初始化 StaticCache（以下示例为 1024 个词）
 # 如果更喜欢，也可以初始化一个 DynamicCache
 prompt_cache = StaticCache(config=model.config, max_batch_size=1, max_cache_len=1024, device=DEVICE, dtype=torch.bfloat16)

 INITIAL_PROMPT = "You are a helpful assistant. "
 inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(DEVICE)
 # 这是缓存的通用提示，需要在无梯度的情况下运行 forward 以便能够复制
 with torch.no_grad():
      prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

 prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
 responses = []
 for prompt in prompts:
     new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to(DEVICE)
     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]
     responses.append(response)

 print(responses)
['<s> You are a helpful assistant. Help me to write a blogpost about travelling.  I am excited to share my experiences with you.  I have been traveling for the past', '<s> You are a helpful assistant. What is the capital of France? \n\nAnswer: Paris is the capital of France.</s>']

旧版缓存格式

在引入 Cache 对象之前，大语言模型的缓存曾经是一个由张量组成的元组的元组。旧版格式的大小是动态的，会随着生成文本而增长，这与 DynamicCache 非常相似。如果的项目依赖于这种旧版格式，可以无缝地将其转换为 DynamicCache，反之亦然。

 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.float16, device_map="auto")
 inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

 # `return_dict_in_generate=True` 是返回缓存所必需的。`return_legacy_cache` 强制返回的缓存为旧版类型
 generation_outputs = model.generate(**inputs, return_dict_in_generate=True, return_legacy_cache=True, max_new_tokens=5)

 # 可以将旧版缓存转换为 DynamicCache，反之亦然。如果有自定义逻辑来以特定格式操作缓存，这会很有帮助。
 cache = DynamicCache.from_legacy_cache(generation_outputs.past_key_values)
 legacy_format_cache = cache.to_legacy_cache()

past_key_values

在大模型的文本生成计算过程中， past_key_values 用于保存之前计算过的键（key）和值（value），这正是 KV Cache 机制的具体体现。

原理层面

在基于 Transformer 架构的自回归语言模型中，每个时间步生成新的标记时，注意力机制都需要根据当前已生成的整个标记序列来计算注意力分数。在这个过程中，对于已经处理过的标记，其对应的键（key）和值（value）在后续的生成步骤中不会发生变化。因此，可以将这些已经计算好的键值对保存下来，避免重复计算。

past_key_values 就是用来存储这些已经计算过的键值对的对象。当模型处理新的标记时，会将新标记对应的键值对与 past_key_values 中存储的旧键值对进行合并，然后一起用于计算注意力分数，从而显著提高计算效率。

代码示例

一个使用 transformers 库进行文本生成，并利用 past_key_values 的示例代码：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# 准备输入文本
input_text = "Once upon a time"
input_ids = tokenizer.encode(input_text, return_tensors="pt")

# 初始化 past_key_values
past_key_values = None

# 模拟逐词生成过程
max_new_tokens = 10
generated_ids = input_ids
for _ in range(max_new_tokens):
    outputs = model(input_ids=generated_ids, past_key_values=past_key_values, use_cache=True)
    # 获取当前时间步的 logits
    logits = outputs.logits[:, -1, :]
    # 选择概率最大的标记作为下一个生成的标记
    next_token_id = torch.argmax(logits, dim=-1, keepdim=True)
    # 更新生成的标记序列
    generated_ids = torch.cat([generated_ids, next_token_id], dim=-1)
    # 更新 past_key_values
    past_key_values = outputs.past_key_values

# 解码生成的标记序列
output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
print(output_text)

在上述代码中：

1. 初始化 past_key_values：在开始生成之前，将 past_key_values 初始化为 None。
2. 模型推理：在每次调用模型时，将 past_key_values 作为参数传入，并设置 use_cache=True，表示启用缓存机制。
3. 更新 past_key_values：每次生成新的标记后，将模型输出的 past_key_values 更新为新的值，以便在下一个时间步使用。

实际作用

使用 past_key_values 保存之前计算过的键值对，避免了对已经处理过的标记的键值对进行重复计算，减少了计算量，从而显著提高了文本生成的速度。在一定程度上减少了不必要的计算，降低了内存的使用，特别是在生成长文本时，这种优势更加明显。