使用Llama3调用model.generate()时如何清除残留KV缓存？

在使用HuggingFace Transformers调用Llama3的model.generate()时，开启use_cache=True会导致KV缓存持续累积——即使新对话没有历史上下文，旧缓存也会被保留，最终引发内存溢出。以下是几种直接有效的解决方法：

方法1：手动重置模型的past_key_values #

模型的KV缓存存储在model.past_key_values属性中，每次generate()调用完成后，直接将其设为None即可清除缓存：

# 执行generate调用
outputs = model.generate(input_ids=input_ids, use_cache=True, ...)
# 清除缓存
model.past_key_values = None

方法2：调用generate时显式指定past_key_values=None #

每次发起新对话的generate()调用时，显式传入past_key_values=None参数，强制模型忽略旧缓存，从头开始生成：

outputs = model.generate(
    input_ids=new_input_ids,
    use_cache=True,
    past_key_values=None,  # 强制不加载旧缓存
    max_new_tokens=512,
    ...
)

方法3：使用reset_cache()方法（Transformers 4.30+版本支持） #

从Transformers 4.30版本开始，模型类新增了专门的reset_cache()方法，调用后可直接清除所有KV缓存：

# 完成generate调用后执行
model.reset_cache()

方法4：通过对话Pipeline管理缓存（对话场景推荐） #

如果是对话场景，使用ConversationPipeline可以更便捷地管理会话和缓存：每次新对话创建全新的Conversation对象，或者调用clear_history()清空历史，Pipeline会自动处理缓存隔离：

from transformers import pipeline, Conversation

chat_pipeline = pipeline("conversational", model="meta-llama/Meta-Llama-3-8B-Instruct")

# 第一轮对话
conv1 = Conversation("请介绍一下Llama3")
result = chat_pipeline(conv1)

# 开启新对话，创建新的Conversation对象
conv2 = Conversation("请推荐一本Python入门书籍")
result = chat_pipeline(conv2)

# 或者清空现有会话的历史并复用对象
conv1.clear_history()
conv1.add_user_input("什么是大语言模型？")
result = chat_pipeline(conv1)

注意事项 #

• 确保新对话的输入prompt是完全独立的，不包含上一轮的历史上下文，配合缓存清除操作才能彻底避免累积。
• 若使用分布式训练或accelerate框架，需在所有进程上执行缓存清除操作，防止部分进程残留缓存导致内存泄漏。

内容的提问来源于stack exchange，提问作者Zheng Zhang