使用 🤗 Optimum Intel 和 fastRAG 优化 CPU Embedding

发布于 2024 年 3 月 15 日

在 GitHub 上更新

访客

访客

访客

访客

Embedding 模型可用于许多应用，例如检索、重排、聚类和分类。近年来，研究界在 Embedding 模型方面取得了重大进展，从而显著增强了所有基于语义表示的应用。诸如 BGE、GTE 和 E5 等模型在 MTEB 基准测试中名列前茅，在某些情况下甚至优于专有的 Embedding 服务。Hugging Face 的模型中心提供了各种大小的模型，从轻量级（1-3.5 亿参数）到 7B 参数模型（例如 Salesforce/SFR-Embedding-Mistral）。这些基于编码器架构的轻量级模型是在 CPU 后端上运行基于语义搜索的应用（例如检索增强生成，RAG）的理想优化和利用对象。

在这篇博客中，我们将展示如何在基于 Xeon 的 CPU 上实现显著的性能提升，并展示使用 fastRAG 将优化后的模型集成到现有 RAG 管道中是多么容易。

使用 Embedding 模型进行信息检索

Embedding 模型将文本数据编码为密集向量，捕捉语义和上下文含义。这通过更具上下文的方式表示词和文档关系，从而实现准确的信息检索。通常，语义相似度将通过 Embedding 向量之间的余弦相似度来衡量。

是否应该始终使用密集向量进行信息检索？两种主要方法各有优缺点。

稀疏检索匹配 n-gram、短语或元数据，以高效且大规模地搜索大型集合。然而，由于查询和文档之间的词汇差距，它可能会错过相关文档。
语义检索将文本编码为密集向量，比词袋模型更好地捕捉上下文和含义。即使存在词汇不匹配，它也能检索到语义相关的文档。然而，与 BM25 等词汇匹配方法相比，它的计算密集、延迟更高，并且需要复杂的编码模型。

Embedding 模型和 RAG

Embedding 模型在 RAG 应用中具有多种关键作用：

离线过程：在索引/更新检索文档库（索引）时，将文档编码为密集向量。
查询编码：在查询时，它们将输入查询编码为用于检索的密集向量表示。
重排：在初次检索后，它们可以通过将检索到的文档编码为密集向量并与查询向量进行比较来对它们进行重排。这允许对最初缺少密集表示的文档进行重排。

优化 RAG 管道中的 Embedding 模型组件对于获得更高效的体验非常重要，特别是：

文档索引/更新：更高的吞吐量允许在初始设置或定期更新期间更快速地编码和索引大型文档集合。
查询编码：较低的查询编码延迟对于响应式实时检索至关重要。更高的吞吐量支持高效地编码许多并发查询，从而实现可扩展性。
重排检索到的文档：初次检索后，Embedding 模型需要快速编码检索到的候选项以进行重排。较低的延迟允许对时间敏感型应用快速重排文档。更高的吞吐量支持并行重排更大的候选集，以实现更全面的重排。

使用 Optimum Intel 和 IPEX 优化 Embedding 模型

Optimum Intel 是一个开源库，可加速在英特尔硬件上使用 Hugging Face 库构建的端到端管道。Optimum Intel 包含多种加速模型的技术，如低比特量化、模型权重剪枝、蒸馏和加速运行时。

Optimum Intel 中包含的运行时和优化利用了英特尔 CPU 上的英特尔® 高级矢量扩展 512（英特尔® AVX-512）、矢量神经网络指令（VNNI）和英特尔® 高级矩阵扩展（英特尔® AMX）来加速模型。具体而言，它在每个核心中都内置了 BFloat16 (bf16) 和 int8 GEMM 加速器，以加速深度学习训练和推理工作负载。AMX 加速推理在 PyTorch 2.0 和 Intel Extension for PyTorch (IPEX) 中引入，此外还有针对各种常见算子的其他优化。

使用 Optimum Intel 可以轻松优化预训练模型；许多简单的示例可以在这里找到。

示例：优化 BGE Embedding 模型

在这篇博客中，我们重点关注北京人工智能研究院的研究人员最近发布的 Embedding 模型，因为他们的模型在广泛采用的 MTEB 排行榜上显示出有竞争力的结果。

BGE 技术细节

双编码器模型是基于 Transformer 的编码器，经过训练以最小化两个语义相似文本向量之间的相似度度量，例如余弦相似度。例如，流行的 Embedding 模型使用 BERT 模型作为基础预训练模型，并对其进行微调以用于 Embedding 文档。表示编码文本的向量由模型输出创建；例如，它可以是 [CLS] 标记向量或所有标记向量的平均值。

与更复杂的 Embedding 架构不同，双编码器仅编码单个文档，因此它们缺乏编码实体（如查询-文档和文档-文档）之间的上下文交互。然而，由于其简单的架构，最先进的双编码器 Embedding 模型表现出有竞争力的性能并且速度极快。

我们专注于 3 个 BGE 模型：small、base 和 large，分别包含 4500 万、1.1 亿和 3.55 亿参数，编码为 384/768/1024 大小的 Embedding 向量。

我们注意到，我们下面展示的优化过程是通用的，可以应用于其他 Embedding 模型（包括双编码器、交叉编码器等）。

分步指南：通过量化进行优化

我们提供了一个分步指南，用于提高 Embedding 模型的性能，重点是减少延迟（批量大小为 1）和增加吞吐量（以每秒编码的文档数衡量）。此方法利用 optimum-intel 和 Intel Neural Compressor 对模型进行量化，并使用 IPEX 在基于英特尔的硬件上进行优化运行时。

第 1 步：安装软件包

要安装 optimum-intel 和 intel-extension-for-transformers，请运行以下命令：

pip install -U optimum[neural-compressor] intel-extension-for-transformers

第 2 步：训练后静态量化

训练后静态量化需要一个校准集来确定权重和激活的动态范围。校准是通过在模型上运行一组有代表性的数据样本，收集统计信息，然后根据收集到的信息对模型进行量化，以最小化精度损失。

以下代码片段显示了量化的代码：

def quantize(model_name: str, output_path: str, calibration_set: "datasets.Dataset"):
    model = AutoModel.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    def preprocess_function(examples):
        return tokenizer(examples["text"], padding="max_length", max_length=512, truncation=True)

    vectorized_ds = calibration_set.map(preprocess_function, num_proc=10)
    vectorized_ds = vectorized_ds.remove_columns(["text"])

    quantizer = INCQuantizer.from_pretrained(model)
    quantization_config = PostTrainingQuantConfig(approach="static", backend="ipex", domain="nlp")
    quantizer.quantize(
        quantization_config=quantization_config,
        calibration_dataset=vectorized_ds,
        save_directory=output_path,
        batch_size=1,
    )
    tokenizer.save_pretrained(output_path)

在我们的校准过程中，我们使用了 qasper 数据集的一个子集。

第 3 步：加载和运行推理

加载量化模型只需运行以下命令：

from optimum.intel import IPEXModel

model = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")

将句子编码为向量可以像我们习惯使用 Transformers 库一样完成。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")
inputs = tokenizer(sentences, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    # get the [CLS] token
    embeddings = outputs[0][:, 0]

我们在下面的评估部分提供了关于如何配置 CPU 后端设置的其他重要细节（正确的机器设置）。

使用 MTEB 进行模型评估

将模型权重以较低精度量化会引入精度损失，因为我们从 `fp32` 权重转换为 `int8` 时会丢失精度。因此，我们旨在通过将优化后的模型与原始模型在两项 MTEB 任务上进行比较来验证其准确性：

检索 (Retrieval) - 对语料库进行编码，并通过搜索索引给定查询来创建排名列表。
重排 (Reranking) - 重新排列检索结果以提高与查询的相关性。

下表显示了每种任务类型的平均准确率（在多个数据集上）（重排为 MAP，检索为 NDCG@10），其中 `int8` 是我们的量化模型，`fp32` 是原始模型（结果来自官方 MTEB 排行榜）。量化模型在重排任务中的错误率与原始模型相比低于 1%，在检索任务中低于 1.55%。

重排

检索


BGE-small
BGE-base
BGE-large

int8	fp32	差异
0.5826	0.5836	-0.17%
0.5886	0.5886	0%
0.5985	0.6003	-0.3%

int8	fp32	差异
0.5138	0.5168	-0.58%
0.5242	0.5325	-1.55%
0.5346	0.5429	-1.53%

速度和延迟

我们将我们模型的性能与另外两种常见的模型使用方法进行比较：

使用 PyTorch 和 Huggingface 的 Transformers 库，精度为 `bf16`。
使用 Intel extension for PyTorch (IPEX) 运行时，精度为 `bf16`，并使用 torchscript 对模型进行跟踪。

实验设置说明：

硬件 (CPU)：第四代英特尔至强 8480+，2 个插槽，每个插槽 56 个核心。
Pytorch 模型在 1 个 CPU 插槽上使用 56 个核心进行评估。
IPEX/Optimum 设置使用 ipexrun 在 1 个 CPU 插槽上进行评估，核心数从 22 到 56 不等。
在所有运行中，都安装了 TCMalloc 并将其定义为环境变量。

我们如何进行评估？

我们创建了一个脚本，使用模型的词汇表生成随机示例。我们加载了原始模型和优化模型，并比较了在我们上面提到的两种场景下编码这些示例所需的时间：批量大小为 1 时的延迟，以及使用批处理示例编码时的吞吐量。

基准 PyTorch 和 Hugging Face

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained("BAAI/bge-small-en-v1.5")

@torch.inference_mode()
def encode_text():
    outputs = model(inputs)

with torch.cpu.amp.autocast(dtype=torch.bfloat16):
    encode_text()

IPEX torchscript 和 `bf16`

import torch
from transformers import AutoModel
import intel_extension_for_pytorch as ipex


model = AutoModel.from_pretrained("BAAI/bge-small-en-v1.5")
model = ipex.optimize(model, dtype=torch.bfloat16)

vocab_size = model.config.vocab_size
batch_size = 1
seq_length = 512
d = torch.randint(vocab_size, size=[batch_size, seq_length])
model = torch.jit.trace(model, (d,), check_trace=False, strict=False)
model = torch.jit.freeze(model)

@torch.inference_mode()
def encode_text():
    outputs = model(inputs)

with torch.cpu.amp.autocast(dtype=torch.bfloat16):
    encode_text()

Optimum Intel 与 IPEX 和 `int8` 模型

import torch
from optimum.intel import IPEXModel

model = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")

@torch.inference_mode()
def encode_text():
    outputs = model(inputs)

encode_text()

延迟性能

在本次评估中，我们旨在衡量模型的响应速度。这是一个在 RAG 管道中编码查询的用例示例。在本次评估中，我们将批处理大小设置为 1，并测量不同文档长度的延迟。

我们可以看到，量化模型在整体上具有最佳的延迟，对于 small 和 base 模型，延迟低于 10 毫秒，对于 large 模型，延迟低于 20 毫秒。与原始模型相比，量化模型的延迟速度提升高达 4.5 倍。

latency
图 1. BGE 模型的延迟。

吞吐量性能

在我们的吞吐量评估中，我们旨在寻找以每秒文档数为单位的峰值编码性能。我们将文本长度设置为 256 个标记，因为这是 RAG 管道中平均文档长度的一个很好的估计，并使用不同的批处理大小（4、8、16、32、64、128、256）进行评估。

结果显示，与其他模型相比，量化模型达到了更高的吞吐量值，并在批处理大小为 128 时达到峰值吞吐量。总体而言，对于所有模型大小，量化模型在各种批处理大小下，与基准 `bf16` 模型相比，性能提升高达 4 倍。

throughput small
图 2. BGE small 的吞吐量。

throughput base
图 3. BGE base 的吞吐量。

throughput large
图 4. BGE large 的吞吐量。

使用 fastRAG 优化 Embedding 模型

作为一个例子，我们将演示如何将优化后的检索/重排模型集成到 fastRAG 中（它也可以轻松地集成到其他 RAG 框架中，如 Langchain 和 LlamaIndex）。

fastRAG 是由英特尔实验室开发的一个研究框架，用于高效和优化的检索增强生成管道，结合了最先进的 LLM 和信息检索。fastRAG 与 Haystack 完全兼容，并包括新颖高效的 RAG 模块，以便在英特尔硬件上进行高效部署。要开始使用 fastRAG，我们邀请读者查看安装说明并使用我们的指南开始使用 fastRAG。

我们将优化的双编码器 Embedding 模型集成在两个模块中：

QuantizedBiEncoderRetriever – 用于从密集索引中索引和检索文档。
QuantizedBiEncoderRanker – 用于使用 Embedding 模型作为复杂检索管道的一部分，对文档列表进行重排。

使用优化的 Retriever 进行快速索引

让我们通过使用一个利用优化 Embedding 模型的密集检索器来创建一个密集索引。

首先，创建一个文档存储库：

from haystack.document_store import InMemoryDocumentStore

document_store = InMemoryDocumentStore(use_gpu=False, use_bm25=False, embedding_dim=384, return_embedding=True)

然后，向其中添加一些文档：

from haystack.schema import Document

# example documents to index
examples = [
   "There is a blue house on Oxford Street.",
   "Paris is the capital of France.",
   "The first commit in fastRAG was in 2022"  
]

documents = []
for i, d in enumerate(examples):
    documents.append(Document(content=d, id=i))
document_store.write_documents(documents)

加载一个带有优化双编码器 Embedding 模型的 Retriever，并对文档存储库中的所有文档进行编码：

from fastrag.retrievers import QuantizedBiEncoderRetriever

model_id = "Intel/bge-small-en-v1.5-rag-int8-static"
retriever = QuantizedBiEncoderRetriever(document_store=document_store, embedding_model=model_id)
document_store.update_embeddings(retriever=retriever)

使用优化的 Ranker 进行重排

下面是一个将优化模型加载到 ranker 节点中的示例，该节点会根据查询对从索引中检索到的所有文档进行编码和重排：

from haystack import Pipeline
from fastrag.rankers import QuantizedBiEncoderRanker

ranker = QuantizedBiEncoderRanker("Intel/bge-large-en-v1.5-rag-int8-static")

p = Pipeline()
p.add_node(component=retriever, name="retriever", inputs=["Query"])
p.add_node(component=ranker, name="ranker", inputs=["retriever"])
results = p.run(query="What is the capital of France?")

# print the documents retrieved
print(results)

完成！创建的管道可用于从文档存储中检索文档，并使用（另一个）Embedding 模型对检索到的文档进行重排，以重新排序文档。此 Jupyter Notebook 中提供了一个更完整的示例。

有关更多与 RAG 相关的方法、模型和示例，我们邀请读者探索 fastRAG/examples 笔记本。

更多博客文章

在至强处理器上借助 🤗 Optimum Intel 实现 SetFit 的极速推理

作者： 2024年4月3日访客 • 11

在至强（Xeon）平台上借助 🤗 Optimum Intel 加速 StarCoder：Q8/Q4 和推测解码

作者： 2024年1月30日访客 • 9

社区

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录发表评论