使用FastAPI+SentenceTransformers推理时遭遇CUDA内存不足问题

问题背景 #

当前使用FastAPI搭配Gunicorn/Uvicorn作为服务器，每个Gunicorn Worker（共10个）启动时都会独立加载基于GPU的SentenceTransformer模型，导致GPU显存被多个重复模型实例耗尽，触发torch.cuda.OutOfMemoryError。

1. 减少Gunicorn Worker数量 #

每个Worker会加载完整的模型到GPU，10个Worker意味着10份模型权重占用显存。根据GPU总显存（15.74GiB）和单模型显存占用量调整Worker数量，建议留出足够余量给推理过程：
修改Gunicorn配置文件：

TIMEOUT 0
GRACEFUL_TIMEOUT 120
KEEP_ALIVE 5
WORKERS 4  # 示例值，根据模型实际显存占用调整，比如单模型占2GiB时，4个Worker仅占用8GiB左右

2. 预加载模型实现Worker共享显存 #

利用Gunicorn的preload_app机制，在Worker fork前提前加载模型，让所有Worker共享同一个GPU模型实例，避免重复加载：

调整代码结构（确保模型在Worker启动前加载）：

from sentence_transformers import SentenceTransformer
from fastapi import FastAPI
import uvicorn

# 全局加载模型，在Worker fork前完成初始化
encoding_model = SentenceTransformer(model_name, device='cuda')

app = FastAPI()

@app.get("/search/")
def encode(query):
    return encoding_model.encode(query).tolist()

def main():
    uvicorn.run(app, host="127.0.0.1", port=8000)

if __name__ == "__main__":
    main()

修改Gunicorn配置，启用预加载：

TIMEOUT 0
GRACEFUL_TIMEOUT 120
KEEP_ALIVE 5
WORKERS 10  # 可保持原有数量，因为模型仅加载一次
PRELOAD_APP = true

修正Docker启动命令（确保用Gunicorn管理Uvicorn Worker，使配置生效）：

CMD ["gunicorn", "-c", "gunicorn.conf.py", "app.main:app", "-k", "uvicorn.workers.UvicornWorker"]

3. 优化PyTorch显存分配 #

根据错误提示，设置环境变量减少显存碎片：
在Dockerfile中添加环境变量：

ENV PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

或在启动脚本中临时设置：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 && gunicorn -c gunicorn.conf.py app.main:app -k uvicorn.workers.UvicornWorker

4. 手动清理推理后显存 #

在接口处理完成后清理未使用的显存，避免累积占用：

@app.get("/search/")
def encode(query):
    import torch
    try:
        return encoding_model.encode(query).tolist()
    finally:
        torch.cuda.empty_cache()

注意：仅在显存紧张时使用，频繁调用会影响推理性能。

5. 模型量化减少显存占用 #

使用8位量化加载模型，大幅降低显存占用（需安装bitsandbytes库）：

encoding_model = SentenceTransformer(
    model_name, 
    device='cuda', 
    model_kwargs={"load_in_8bit": True}
)

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