前言#

在构建企业级 RAG（检索增强生成）系统时，Rerank（重排序） 是提升准确率的关键一环。我们通常会选择强大的 CPU 服务器来部署这一服务，比如拥有恐怖核心数的 AMD EPYC 系列。

然而，许多人在部署 HuggingFace 的 text-embeddings-inference (TEI) CPU 版后会发现：为什么我的 CPU 利用率不高，推理速度也不尽如人意？ 感觉就像开着法拉利在送外卖。

根本原因在于：默认的模型格式，喂不饱你的 CPU。

FP32 vs INT8#

要理解性能瓶颈，我们需要理解 CPU 处理数据的两种方式：

1. FP32 (全精度浮点数) —— 精确但笨重 默认下载的模型权重是 FP32 格式。处理它就像让厨师用手术刀切菜，每一刀都精确到微米。

优点： 数值极其精确。
缺点： 数据量大，占用内存带宽，计算速度慢。对于 Rerank 这种只需要“比大小”排序的任务来说，这是一种极大的性能浪费。

2. INT8 (8位整数) —— 粗犷但极速 这就是我们今天要做的——量化 (Quantization)。我们将模型权重从复杂的浮点数“四舍五入”成简单的整数。这就像让厨师换上了大菜刀，虽然每一刀没那么精细，但切菜速度快了好几倍。

为什么速度快？因为 AVX-512！ 现代 CPU（如 AMD EPYC 9004系列或 Intel Xeon）都配备了强大的 SIMD（单指令多数据流） 指令集，比如 AVX-512。你可以把它想象成一条512车道的高速公路。

跑 FP32 时： 车身太宽，高速公路一次只能并排跑 16 辆车。
跑 INT8 时： 车身变窄，高速公路一次可以并排跑 64 辆车！

结论： 只要我们将模型转换为 INT8 格式，就能激活 CPU 的 AVX-512 VNNI 指令集，实现单次吞吐量 4 倍的理论提升，同时内存占用减少到原来的 1/4。而这一切，对于 Rerank 排序结果精度的影响微乎其微。

转换脚本#

使用 HuggingFace 的 optimum 库来完成这次转换。能自动识别你的 CPU 架构并应用最佳优化配置。

vim quantize_rerank.py

text

import os
import shutil
from optimum.onnxruntime import ORTQuantizer
from optimum.onnxruntime.configuration import AutoQuantizationConfig

# 1. 定义路径
# 这是你上一步生成的 FP32 ONNX 文件夹
input_model_dir = "./data/bge-reranker-v2-m3/onnx" 
# 这是我们要输出的 INT8 新文件夹
output_model_dir = "./data/bge-reranker-v2-m3/quantized"

print(f"🚀 正在准备量化...")
print(f"输入目录: {input_model_dir}")
print(f"输出目录: {output_model_dir}")

# 2. 定义量化配置 (针对 AMD EPYC 的 AVX-512 VNNI 优化)
# is_static=False 表示使用动态量化，这是 CPU 推理的最佳实践
qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)

try:
    # 3. 加载量化器
    quantizer = ORTQuantizer.from_pretrained(input_model_dir, file_name="model.onnx")

    # 4. 执行量化
    print("⏳ 正在量化模型 (FP32 -> INT8)...")
    quantizer.quantize(
        save_dir=output_model_dir,
        quantization_config=qconfig,
    )
    
    # 5. 关键步骤：把 tokenizer 和 config 文件也复制过去
    # 量化工具只生成 model.onnx，TEI 启动还需要配置文件
    print("📂 正在复制配置文件...")
    for filename in os.listdir(input_model_dir):
        if filename != "model.onnx" and not filename.endswith(".onnx"):
            src = os.path.join(input_model_dir, filename)
            dst = os.path.join(output_model_dir, filename)
            if os.path.isfile(src):
                shutil.copy2(src, dst)

    # 6. TEI 默认只认 "model.onnx"，但量化工具可能生成 "model_quantized.onnx"
    # 我们强制把它改名为 model.onnx
    quantized_file = os.path.join(output_model_dir, "model_quantized.onnx")
    target_file = os.path.join(output_model_dir, "model.onnx")
    
    if os.path.exists(quantized_file):
        os.rename(quantized_file, target_file)
        print("✅ 已重命名为 model.onnx 以适配 TEI")
    
    print(f"✅ 大功告成！新模型已保存在: {output_model_dir}")
    print("现在去修改 docker-compose.yml 吧！")

except Exception as e:
    print(f"❌ 发生错误: {e}")

只需要修改 docker-compose.yml，将挂载卷指向新的量化模型目录即可。TEI 的镜像无需更换，它会自动识别并加载 INT8 模型。将tei-rerank 的映射路径改成：这样即可使用量化模型

text

    volumes:
      # 👇 挂载刚才下载的 rerank 模型目录
      #- ./data/bge-reranker-v2-m3:/data/rerank
      - ./data/bge-reranker-v2-m3/quantized:/data/rerank

CPU 跑 Rerank 太慢？一个脚本开启 INT8 量化，性能大幅度提升！

前言#

FP32 vs INT8#

转换脚本#

评论