我要投稿

来喽，手把手优化向量模型的吞吐量(2)实战

发布日期：2024-08-20 05:15:43 浏览次数： 2708

作者：哎呀AIYA

微信搜一搜，关注“哎呀AIYA”

上篇我们介绍了，在RAG系统中向量模型使用的情况，本篇文章我们来介绍如何提升向量模型的吞吐量。下图是优化后的bge模型的吞吐量对比：

使用 Optimum Intel 和 IPEX 优化嵌入模型

Optimum Intel 是一个开源库，其针对英特尔硬件对使用 Hugging Face 库构建的端到端流水线进行加速和优化。Optimum Intel 实现了多种模型加速技术，如低比特量化、模型权重修剪、蒸馏以及运行时优化。

Optimum Intel 在优化时充分利用了英特尔® 先进矢量扩展 512 (英特尔® AVX-512) 、矢量神经网络指令 (Vector Neural Network Instructions，VNNI) 以及英特尔® 高级矩阵扩展 (英特尔® AMX) 等特性以加速模型的运行。具体来说，每个 CPU 核中都内置了 BFloat16 ( bf16 ) 和 int8 GEMM 加速器，以加速深度学习训练和推理工作负载。除了针对各种常见运算的优化之外，PyTorch 2.0 和 Intel Extension for PyTorch (IPEX) 中还充分利用了 AMX 以加速推理。

使用 Optimum Intel 可以轻松优化预训练模型的推理任务。本文将介绍广为人知的 MTEB 排行榜上取得了亮眼的排名的 BGE 嵌入模型。

BGE 技术细节

双编码器模型基于 Transformer 编码器架构，其训练目标是最大化两个语义相似的文本的嵌入向量之间的相似度，常见的指标是余弦相似度。举个常见的例子，我们可以使用 BERT 模型作为基础预训练模型，并对其进行微调以生成嵌入模型从而为文档生成嵌入向量。有多种方法可用于根据模型输出构造出文本的嵌入向量，例如，可以直接取 [CLS] 词元的嵌入向量，也可以对所有输入token的嵌入向量取平均值。

双编码器模型是个相对比较简单的嵌入编码架构，其仅针对单个文档上下文进行编码，因此它们无法对诸如查询 - 文档及文档 - 文档这样的交叉上下文进行编码。然而，最先进的双编码器嵌入模型已能表现出相当有竞争力的性能，再加上因其架构简单带来的极快的速度，因此该架构的模型成为了目前的标配。

这里，我们主要关注 3 个 BGE 模型: small、base 以及 large，它们的参数量分别为 45M、110M 以及 355M，嵌入向量维度分别为 384、768 以及 1024。

请注意，下文展示的优化过程是通用的，你可以将它们应用于任何其他嵌入模型 (包括双编码器模型、交叉编码器模型等)。

模型量化分步指南

下面，我们展示如何提高嵌入模型在 CPU 上的性能，我们的优化重点是降低延迟 (batch size 为 1) 以及提高吞吐量 (以每秒编码的文档数来衡量)。我们用 optimum-intel 和 INC (Intel Neural Compressor) 对模型进行量化，并用 IPEX 来优化模型在 Intel 的硬件上的运行时间。

第 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 数据集的一个子集作为校准集。

# 数据集地址https://huggingface.co/datasets/allenai/qasper

第 3 步: 加载模型，运行推理

仅需运行以下命令，即可加载量化模型:

from optimum.intel import IPEXModelmodel = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")

随后，我们使用 transformers 的 API 将句子编码为向量:

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] tokenembeddings = outputs[0][:, 0]

使用 MTEB 进行模型评估

将模型的权重量化到较低的精度会导致准确度的损失，因为在权重从 fp32 转换到 int8 的过程中会损失精度。所以，我们在如下两个 MTEB 任务上对量化模型与原始模型进行比较以验证量化模型的准确度到底如何:

检索 - 对语料库进行编码，并生成索引库，然后在索引库中搜索给定查询，以找出与给定查询相似的文本并排序。
重排 - 对检索结果进行重排，以细化与给定查询的相关性排名。

下表展示了每个任务在多个数据集上的平均准确度 (其中，MAP 用于重排，NDCG@10 用于检索)，表中 int8 表示量化模型， fp32 表示原始模型 (原始模型结果取自官方 MTEB 排行榜)。与原始模型相比，量化模型在重排任务上的准确度损失低于 1%，在检索任务中的准确度损失低于 1.55%。