运行 LLM 嵌入模型在 CPU 上速度很慢，在 GPU 上成本很高。我们将使用ONNX 模型量化将其速度提高 3 倍，看看不同的int8 格式如何影响新旧硬件上的性能，并在量化模型的基础上进一步进行ONNX 变压器优化。

无人谈论的向量搜索问题

要使用 LLM 嵌入执行语义搜索，您必须先计算这些嵌入。不幸的是，在市场上有许多向量搜索数据库的情况下，计算嵌入被认为是事后才考虑的问题，并且超出了范围。

嵌入推理是任何语义搜索系统的首要步骤。图片由作者提供。

一种新的开源搜索引擎Nixiesearch，它可以根据您的数据微调嵌入。由于我们在服务器端处理嵌入，因此我们并不惊讶地看到在 CPU 上运行嵌入对性能的巨大影响。

使用 e5-small-v2 嵌入模型的 Nixiesearch 索引过程的火焰图。图片由作者提供。

上面的火焰图显示，95% 的 CPU时间都花在了计算嵌入上。当然，您可以通过将索引切换到 GPU 来加快速度，但对于想要尝试 Nixiesearch 的人来说，这仍然是一个很大的遗憾。我们能否在仍使用 CPU 的情况下加快速度？

模型量化

目前所有的深度神经网络都只是矩阵运算的华丽组合。

注意层矩阵操作。图片来自 Vaswani 等人的《Attention Is All You Need》。

如上图所示，Transformer 网络的注意层只是矩阵之上的简单代数变换的组合。在经典实现中，这些矩阵包含 32 位浮点值。如果我们准备牺牲一点精度来获得更好的性能，将浮点大小从 32 字节减少到 8 字节，会怎么样？

32 位浮点矩阵乘法。图片由作者提供。

这种方法称为量化：降低用于存储模型权重和神经元激活值的矩阵的数值精度。

精度从 32 位提高到 8 位将使存储模型权重所需的 RAM 大小减少 4 倍，并有望使其在现代 CPU 上运行得更快。

但是，您不能将 32 位浮点数恰好放入 8 位整数而不会造成任何损失：每个数字只有 8 位存储空间，因此您只能编码 256 个不同的值！

训练好的 TinyBERT 中的权重分布。图片来自 J. Jin 等人的《KDLSQ-BERT：结合知识蒸馏和学习步长量化的量化 Bert》。

但神经网络中的权重和激活值并不是随机数！从上面的直方图可以看出，它们通常接近于零，并且呈正态分布。通过这一观察，我们可以将所有网络操作替换为量化感知操作——同时跟踪底层数值分布的零点和尺度。

ONNX 中基于运算符的量化。图片由作者提供。

这种方法称为运算符量化，被认为是解决问题的最直接方法。另一种选择是 QDQ 量化，您仍然使用相同的 32 位运算符，但在常规运算符之前注入量化-反量化运算符对 – 这通常要慢得多。

运算符与 QDQ 量化。来自 ONNX 文档 — 量化 ONNX 模型的图片。

在上图中，您可以看到在 QDQ 模式下向图中注入了额外的节点，这通常会导致更糟糕的性能：

额外的量化-反量化运算符不是免费的。
主要运算符像以前一样在 Float32数据上运行，因此您只能节省 RAM/VRAM 的使用，而不会节省性能。
QDQ 仅支持ONNX 运行时的静态量化– 有关详细信息，请参阅下一章。

为了简单起见，本文我们只针对运算符量化。

请注意，模型量化与所有主流向量搜索引擎（如 Elasticsearch、Qdrant、Vespa 和 Weaviate）支持的嵌入量化不同。量化模型仍像以前一样发出 Float32 嵌入– 它只是对权重和激活使用更紧凑的布局。

动态量化与静态量化

对于图中每个操作，计算零点和尺度量化参数有两个选项：

动态：每个操作符在运行时为每个批次重新计算这些参数。它更耗费资源，但也更精确——这些参数可能会在批次之间漂移。
静态：我们不会每次都为每个批次重新计算这些参数，而是进行离线校准——使用一个很小的数据集前馈网络，并根据观察到的分布静态记录量化参数。这没有额外的开销，但由于量化参数是静态的，因此在批次之间出现漂移的情况下，它们可能不完美。

静态量化的主要缺点是需要进行校准，这是一个额外的手动步骤。由于静态量化通常会导致精度下降，因此我们将在本文中重点介绍动态量化方法。

LLM 推理运行时

您在 HuggingFace Hub 上找到的模型文件仅包含您需要在模型权重之上执行的矩阵运算的定义（执行图）。

执行运行时是执行实际矩阵乘法的部分。图片由作者提供。

执行运行时会在您的硬件上解释并执行此图：

PyTorch 和 TensorFlow – 用于模型开发，以 Python 为中心，但也有其他语言（如 C/C++）的低级绑定。实际上，两者都针对 GPU 上的训练和批处理进行了更优化。
OpenVINO — 英特尔推出的以 CPU 为中心的运行时，仅限 Python/C++。
ONNX — 一个开放的多语言（Java/JS/WASM/C++/Python）和多后端（CPU/GPU/TPU）运行时。
TensorRT ——Nvidia 推出的 ONNX 兼容运行时，专门用于 GPU 执行。

Nixiesearch 与Apache Lucene绑定，作为 JVM 应用程序实现——并且其他开源搜索引擎选择使用 ONNX 作为神经网络的主要执行运行时。

将模型转换为 ONNX

您需要先进行转换才能在 ONNX 运行时内执行模型。

ONNX 模型转换流程。图片由作者提供。

为了实现这样的转换，我们开发了一个工具nixiesearch/onnx-convert，它是原始xenova/transformers.js项目中转换脚本的扩展版本。转换后的模型文件与任何 ONNX 风格的搜索引擎兼容，并且可以在 Elasticsearch Inference Processor、Vespa Embedder和直接在Nixiesearch中使用。

nixiesearch/onnx-convert转换工具。图片由作者提供。

ONNX 转换+量化过程有多个重要的可调参数，影响性能和质量：

底层量化格式：可以是有符号/无符号的 8 位整数和 16 位浮点数。这会对性能和精度产生什么影响？
ONNX Transformer 优化器：ONNX 可以将多个典型运算符融合到单个优化核心中，用于注意力块等。这真的很重要吗？它的级别如何影响最终推理延迟？

除了这些可变参数之外，我们将其他不太重要的参数固定为推荐的常量值：

ONNX opset ：ONNX 有多个版本，对数据类型和运算符的支持不同。我们选择了Python onnx 包支持的最新opset=17 。
每通道量化：图表是否应跟踪单个张量尺度和零点值，还是应按通道（张量中的切片）进行？关闭它可能会略微提高性能并降低精度 — 但由于需要使用额外的内存，我们选择默认启用它。
7 位比例：作为对数字溢出的额外保护，8 位值是否应该缩小一点？ONNX 文档指出，它可能会提高没有 AVX-VNNI 的旧硬件的精度。

我们将采用E5-v2系列嵌入模型，采用小型、基础和大型变体来查看每个变化对不同大小模型的影响。

QUint8/QInt8/Float16 和推理延迟

我们将e5-small-v2、e5-base-v2和e5–large-v2转换为QUint8、QInt8和Float16数值类型，仍然没有进行任何优化：

ONNX 转换工具。图片由作者提供。

并运行嵌入基准测试套件：

对于每种尺寸的每个模型，使用 JVM 中的 onnxruntime 计算嵌入延迟。
该套件运行在最新一代 AWS M7I.2xlarge实例上，该实例具有 8 个支持 AVX-VNNI 的 VCPU。

Float32 与其他格式之间的相对推理时间改进。值越高越好。图片来自作者。

上表显示，“对于 e5-small-v2 模型，4 个 token 的 QUInt8 格式比 Float32快 1.38 倍”。我们得到的结果相当令人惊讶：

好消息是，我们对基础模型和大型模型的推理速度提高了 3 倍以上！
在 VNNI 硬件上，QUInt8 和 QInt8 差别不大。
混合精度Float16 格式意外地比基线慢了 2 到 7 倍。这一令人惊讶的事实源于现代 CPU 缺乏对 FP16/BF16 的支持：只有支持 AMX 的最新 Intel Sapphire Rapids Xeon CPU才能原生处理这些数据类型。没有 AMX 的 CPU 每次发现 Float16 类型时都会执行向下转换-向上转换操作。

因此，你显然应该量化你的模型，但是如果你的硬件不支持 VNNI 怎么办？