使用Llama 2模型加载pipeline后无法清理GPU内存求助

这个问题确实有不少开发者碰到过，核心原因是transformers的text-generation pipeline会对模型持有额外的隐性引用，尤其是当使用device_map="auto"进行模型分片加载时，这些引用不会随着del命令被直接清理，导致torch无法回收GPU内存。

问题分析 #

当你直接把已加载的模型和tokenizer传入pipeline时，pipeline内部会创建辅助组件（比如生成器、设备调度器），这些组件会间接持有模型的GPU张量引用——哪怕你删除了model、tokenizer、pipeline变量，这些隐性引用依然存在，gc.collect()和torch.cuda.empty_cache()自然无法生效。而仅加载模型和tokenizer时，没有这些额外组件，内存就能正常释放。

可行的解决方案 #

以下是经过验证的几种解决方法：

1. 先将模型移回CPU再清理
   在删除变量前，显式把pipeline中的模型移回CPU，切断GPU张量的持有关系：

   # 在清理前添加这一步
   pipeline.model = pipeline.model.to("cpu")
   del model
   del tokenizer
   del pipeline
   gc.collect()
   torch.cuda.empty_cache()
   

2. 显式清理pipeline内部引用
   手动删除pipeline对模型和tokenizer的直接引用，再清理变量：

   # 先清空pipeline的内部引用
   del pipeline.model
   del pipeline.tokenizer
   del model
   del tokenizer
   del pipeline
   gc.collect()
   torch.cuda.empty_cache()
   

3. 升级transformers到最新稳定版
   旧版本的transformers在Llama 2的pipeline内存管理上存在bug，升级到4.34.0及以上版本后，很多内存泄漏问题已经被官方修复。可以用命令升级：

   pip install --upgrade transformers accelerate
   

4. 让pipeline自行加载模型和tokenizer
   不要提前加载模型和tokenizer，而是让pipeline直接从模型ID加载，这样它内部的引用管理会更规范，清理时也更彻底：

   from transformers import pipeline
   import torch
   import gc
   
   modelId = "meta-llama/Llama-2-7b-chat-hf"
   pipeline = pipeline(
       task="text-generation",
       model=modelId,
       device_map="auto",
       torch_dtype=torch.float16
   )
   
   # 清理步骤
   del pipeline
   gc.collect()
   torch.cuda.empty_cache()
   

修改后的完整示例代码 #

from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
import gc

modelId = "meta-llama/Llama-2-7b-chat-hf"

model = AutoModelForCausalLM.from_pretrained(modelId, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(modelId)

pipeline = pipeline(task="text-generation", model=model, tokenizer=tokenizer)

# 清理步骤优化
pipeline.model = pipeline.model.to("cpu")
del pipeline.model
del pipeline.tokenizer
del model
del tokenizer
del pipeline
gc.collect()
torch.cuda.empty_cache()

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