使用Hugging Face Pipeline时的缓存彻底清理问题求助

针对你遇到的Pipeline运行后GPU内存上升且无法通过gc.collect()和torch.cuda.empty_cache()清理的问题，可尝试以下具体方案：

1. 禁用或清理KV缓存 #

文本生成任务中，模型默认会保留KV缓存以加速后续生成，但这会持续占用显存。你可以：

• 初始化Pipeline时直接禁用缓存：

  pipe= pipeline(
      "text-generation",    
      model=model_1,        
      tokenizer=tokenizer_1,   
      torch_dtype=torch.bfloat16,  
      max_new_tokens=500,  
      min_new_tokens=-1, 
      do_sample=True,          
      top_k=50,      
      top_p=1, 
      repetition_penalty = 1.15, 
      temperature=0.01,
      use_cache=False  # 禁用KV缓存
  )
  

• 若需要缓存加速单次生成，运行后手动清理：

  outcome = pipe(input)
  # 清理模型层的KV缓存
  for layer in model_1.layers:
      if hasattr(layer, 'past_key_values'):
          layer.past_key_values = None
  

2. 彻底释放对象引用后再清理显存 #

之前的清理方法无效，大概率是因为相关变量的引用未被彻底删除。按以下顺序执行：

# 运行Pipeline并处理结果
outcome = pipe(input)
# 处理完结果后，删除所有相关变量
del outcome, pipe
# 若不再需要模型和tokenizer，也一并删除
del model_1, tokenizer_1
# 强制触发垃圾回收
import gc
gc.collect()
# 清空CUDA未被引用的显存
torch.cuda.empty_cache()

3. 用上下文管理器隔离计算过程 #

使用torch.no_grad()上下文管理器，减少计算中生成的梯度张量占用，同时确保退出时自动释放部分资源：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

with torch.no_grad():
    tokenizer_1 = AutoTokenizer.from_pretrained(model_1)
    model_1 = AutoModelForCausalLM.from_pretrained(model_1, device_map=device_1, torch_dtype=torch.bfloat16)
    pipe= pipeline(
        "text-generation",    
        model=model_1,        
        tokenizer=tokenizer_1,   
        torch_dtype=torch.bfloat16,  
        max_new_tokens=500,  
        min_new_tokens=-1, 
        do_sample=True,          
        top_k=50,      
        top_p=1, 
        repetition_penalty = 1.15, 
        temperature=0.01
    )
    outcome= pipe(input)
    # 在这里处理outcome

# 退出上下文后执行清理
del tokenizer_1, model_1, pipe, outcome
gc.collect()
torch.cuda.empty_cache()

4. 升级Transformers版本 #

部分旧版本存在内存泄漏bug，升级到最新稳定版可解决：

pip install --upgrade transformers accelerate

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