用LangChain调用本地离线Embedding模型时显存不释放的缓存清理策略

文章导读
用LangChain调用本地离线Embedding模型时,显存不释放通常不是LangChain本身的行为,而是底层模型实例和推理框架的生命周期没有被管理好。LangChain的Embedding封装只是发起调用,真正加载权重、分配显存的是Ollama、sentence-transformers、transformers或llama.cpp等后端。因此先确认显存被哪个进程占用,再选择对应的清理策略。
📋 目录
  1. 一 先确认:显存是被常驻进程占着,还是随调用增长
  2. 二 检查LangChain调用方式:是否每次调用都新建模型
  3. 三 按本地框架选择清理策略
  4. 四 验证清理是否生效
A A

用LangChain调用本地离线Embedding模型时,显存不释放通常不是LangChain本身的行为,而是底层模型实例和推理框架的生命周期没有被管理好。LangChain的Embedding封装只是发起调用,真正加载权重、分配显存的是Ollama、sentence-transformers、transformers或llama.cpp等后端。因此先确认显存被哪个进程占用,再选择对应的清理策略。

显存不释放一般有三种来源:模型实例被Python代码持有未释放、推理框架的显存缓存未清理、本地服务端(如Ollama)默认保持模型驻留。处理顺序是先看占用进程和显存趋势,再按框架分别处理:复用模型实例、主动释放引用并清理缓存、设置keep_alive或停止服务。清理动作会带来重新加载延迟,需要结合使用频率取舍。

先确认:显存是被常驻进程占着,还是随调用增长

打开终端执行 nvidia-smi,查看“Processes”区域里的进程名和PID。如果是python进程,可以记录PID,再在代码里打印当前进程ID做对比。如果显存总量只增长一次后稳定,通常是模型权重加载后驻留;如果每次调用Embedding都增长,优先怀疑每次调用都重新加载模型或累积了未释放的中间张量。

还可以用 watch -n 2 nvidia-smi 观察趋势。这一步是为了避免误杀:如果模型本来就还会继续用,保持驻留是正常的,真正需要处理的是只调用一次或低频调用后显存仍被长期占用的情况。

检查LangChain调用方式:是否每次调用都新建模型

很多“显存不释放”其实是代码在循环里反复创建Embedding对象。LangChain的OllamaEmbeddings、HuggingFaceEmbeddings等类,初始化时并不一定立即加载模型,但放到循环里反复实例化时,可能会出现重复加载或无法销毁的情况。建议把Embedding对象创建在循环外,全局复用。

一个通用骨架示例:

from langchain_community.embeddings import OllamaEmbeddings

# 在循环外创建一次
embeddings = OllamaEmbeddings(model='your-model')

def get_vec(text):
    # 循环内只调用,不再创建对象
    return embeddings.embed_query(text)

如果是HuggingFaceEmbeddings或自定义transformers封装,还需要在模型层面做类似复用。先排除这种低级问题,再谈清理策略。

用LangChain调用本地离线Embedding模型时显存不释放的缓存清理策略

按本地框架选择清理策略

如果确认模型实例是单例、调用也正常,但显存仍然不释放,就要看是哪种后端。

Ollama后端:调整keep_alive

Ollama默认会将模型在内存或显存中保留一段时间。LangChain的OllamaEmbeddings没有暴露卸载接口,但可以通过Ollama的API或命令行调整驻留时间。调用OllamaAPI时,设置keep_alive为0表示一次请求后立即卸载,或者用命令 ollama stop your-model 手动卸载。注意:如果应用仍然持有该连接,Ollama可能会重新加载。

示例:

# 请求完成后手动卸载
ollama stop your-model

也可以用API发送参数,但具体字段随Ollama版本变化,建议先查当前版本的接口响应。设置keep_alive=0会牺牲再次调用时的加载延迟,适合低频场景。

用LangChain调用本地离线Embedding模型时显存不释放的缓存清理策略

HuggingFace/transformers或sentence-transformers:释放引用并清理缓存

如果使用HuggingFaceEmbeddings直接加载本地模型,Python进程会持有模型对象。要真正确认释放,需要删除模型引用、清空推理缓存并回收CUDA缓存。示例如下:

import torch
import gc

# 假设 embeddings 内部持有模型对象
del embeddings
gc.collect()
torch.cuda.empty_cache()

但这里有一个边界:torch.cuda.empty_cache()只释放当前进程的CUDA缓存,已经返回给进程的显存不一定完全归还给操作系统;另外,如果代码中还有其他变量引用了模型,del不会生效。所以删除后最好打印引用计数,或者用弱引用检查是否还有外部持有。

验证清理是否生效

清理后,用nvidia-smi重新查看进程显存占用。如果占用下降,说明刚才的清理动作有效;如果占用不变,可能还有对象被引用,或者缓存是CUDA上下文的一部分,本身不会释放。这时可以用一段脚本在循环中调用embed_query并观察显存是否持续上涨。

验证清单:

  • 确认当前进程PID与nvidia-smi列出的一致。
  • 调用前记录显存占用baseline。
  • 执行清理动作后,等待几秒再查看显存。
  • 如果使用Ollama,检查是否有ollama serve进程仍在运行,并确认模型是否还在驻留列表。

最后提醒:上述清理策略属于生命周期管理,不是性能优化。把keep_alive调为0或频繁删除模型,会让下一次调用变慢。如果业务上要高频率使用Embedding,保留常驻反而是合理选择。显存是否“不释放”要和实际使用模式对比,不一定每次都需要清理。