如何提升Sentence Transformer推理阶段的GPU利用率以加速编码？

问题背景 #

使用Sentence Transformer计算1024维向量用于相似度检索，通过FastAPI搭建文本转向量API，计划处理1000万份文档。测试结果：

• CPU单句编码耗时约0.1秒
• GTX 1650单句编码耗时约0.01秒，GPU利用率50-60%
• Tesla T4单句编码耗时仍为0.01秒，但GPU利用率仅10-15%，算力严重闲置，增加worker数量无效果

1. 批量处理请求 #

单句推理无法充分发挥Tesla T4的并行算力，必须改用批量处理模式：

• 修改FastAPI接口，支持接收批量文本输入
• 调整model.encode的batch_size参数，根据T4显存（16GB）尝试64/128等数值

from fastapi import FastAPI
from typing import List
from sentence_transformers import SentenceTransformer

app = FastAPI()
model = SentenceTransformer("your-model-path", device="cuda")

@app.post("/encode-batch")
async def encode_batch(texts: List[str]):
    embeddings = model.encode(
        texts,
        device='cuda',
        normalize_embeddings=True,
        batch_size=64,  # 可根据显存占用调整
        show_progress_bar=False
    ).tolist()
    return {"embeddings": embeddings}

客户端调用时尽量批量发送文本，避免单句请求，这是提升GPU利用率最直接的手段。

2. 优化数据传输与推理参数 #

• 启用convert_to_tensor=True，减少CPU-GPU间的数据拷贝开销，最后再转CPU列表：

embeddings = model.encode(
    texts,
    device='cuda',
    normalize_embeddings=True,
    batch_size=64,
    show_progress_bar=False,
    convert_to_tensor=True
).cpu().tolist()

• 关闭进度条（show_progress_bar=False），避免不必要的CPU开销。

3. 调整FastAPI并发配置 #

FastAPI默认worker配置不适合GPU场景，需针对性调整：

• 使用uvicorn启动时，设置worker数量为2-4（T4的最优值，过多会导致显存竞争）
• 启用uvloop提升异步性能，限制并发数避免请求堆积

uvicorn main:app --host 0.0.0.0 --port 8000 --workers 2 --loop uvloop --limit-concurrency 100

4. 启用TensorRT加速（可选） #

针对Sentence Transformer模型，可通过Optimum库导出为TensorRT引擎，进一步提升批量推理速度：

from sentence_transformers import SentenceTransformer
from optimum.onnxruntime import ORTModelForFeatureExtraction
from optimum.onnxruntime.configuration import AutoQuantizationConfig

# 导出并量化模型
model = SentenceTransformer("your-model-path")
quant_config = AutoQuantizationConfig.avx512_vnni(is_static=False)
model.save_pretrained("trt-optimized-model")
ORTModelForFeatureExtraction.from_pretrained(
    "trt-optimized-model",
    export=True,
    quantization_config=quant_config
)

# 加载优化后的模型
optimized_model = SentenceTransformer("trt-optimized-model")

5. 监控与调优 #

使用watch -d -n 0.5 nvidia-smi命令监控GPU-Util指标：

• 若利用率仍低，逐步增大batch_size，直到显存占用接近T4的16GB上限
• 若出现显存不足（OOM），则降低batch_size或减少worker数量

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