使用 Optimum Intel 与 fastRAG 的 CPU 优化嵌入

Hugging Face 和 Intel 展示了一种在 Intel Xeon CPU 上显著加速嵌入模型的方法，使用 Optimum Intel 库和 fastRAG 框架。通过采用后训练静态量化至 int8，这些优化相较于基准 bf16 模型将延迟降低最高可达 4.5 倍，吞吐量提升最高可达 4 倍。

加速 RAG 的嵌入模型

嵌入模型是检索增强生成（Retrieval Augmented Generation，RAG）流水线的关键，承担三大主要功能：离线文档索引、实时查询编码以及检索文档的重新排序。虽然语义检索比稀疏检索（如 BM25）更好地捕获上下文，但其计算成本更高。

在 CPU 后端对这些模型进行优化至关重要：

  • 文档索引：提升吞吐量，以更快地索引大规模集合。
  • 查询编码：降低延迟，实现响应式实时检索。
  • 重新排序：快速处理候选集合，以满足对时间敏感的应用。

通过 Optimum Intel 与 IPEX 的技术优化

Optimum Intel 是一个开源库，旨在加速 Intel 硬件上的 Hugging Face 流水线。它利用 Intel® 高级向量扩展 512（AVX-512）、向量神经网络指令（VNNI）和 Intel® 高级矩阵扩展（AMX）来优化深度学习工作负载。

量化工作流

优化过程侧重于将精度从 fp32 降至 int8，包括以下步骤：

  1. 安装：使用 optimum[neural-compressor] 和 intel-extension-for-transformers。
  2. 后训练静态量化：该过程使用校准集（例如 qasper 数据集的子集）来确定权重和激活的动态范围，以在转为 int8 时最小化精度损失。
  3. 推理：通过 IPEXModel 加载量化模型，利用 Intel Extension for PyTorch (IPEX) 实现优化的运行时执行。

BGE 模型的性能基准

评估使用了 BGE（北京人工智能研究院）模型，分别为小型（45M）、基础（110M）和大型（355M）参数规模，在第 4 代 Intel Xeon 8480+ CPU 上进行。

精度保持

将模型量化至 int8 在 MTEB（Massive Text Embedding Benchmark）任务中相较于 fp32 模型的精度下降极小：

  • 重新排序：错误率低于 1%。
  • 检索：错误率低于 1.55%。

延迟与吞吐量

相较于基准 PyTorch bf16 实现，int8 量化模型展现出显著提升：

  • 延迟：最高加速 4.5 倍。小型和基础模型的延迟低于 10ms，且大型模型保持在 20ms 以下。
  • 吞吐量：提升最高达 4 倍。所有模型规模在文档长度为 256 token、批量大小为 128 时达到峰值吞吐量。

与 fastRAG 的集成

fastRAG 是 Intel Labs 的一个研究框架，用于优化的 RAG 流水线，兼容 Haystack。优化后的嵌入模型通过以下两个模块集成到 fastRAG 中：

  • QuantizedBiEncoderRetriever：用于对密集索引进行文档索引和检索。
  • QuantizedBiEncoderRanker：用于对检索到的文档进行重新排序，以提升相关性。

这些模块使开发者能够在基于 Haystack 的流水线中轻松将标准嵌入模型替换为 Intel 优化版本，从而提升索引和检索阶段的效率。

Sources