利用 PyTorch 原生栈在服务器 CPU 上高效进行小型 NLP 模型的 INT8 推理

# 在服务器CPU上使用PyTorch原生栈实现小型NLP模型的高效INT8推理
来源: https://arxiv.org/html/2608.18182

###### 摘要
在大型语言模型时代，小型NLP模型，特别是BERT系列编码器，在分类、排序和检索等工业工作负载中仍然非常重要。在服务器CPU上，INT8量化提供了一种具有吸引力的延迟-吞吐量-成本权衡方案，但用户越来越期望这种加速能直接在原生PyTorch栈中可用。我们将SmoothQuant集成到TorchAO中，并通过TorchInductor的图级融合以及跨oneDNN、AVX512_VNNI和基于AMX实现的高效INT8 GEMM内核选择，针对Intel Xeon CPU优化了由此产生的推理路径。在BERT、DistilBERT和XLM-RoBERTa的基准测试中，该方法最高可实现5.8倍的端到端吞吐量加速，且相对于FP32基线几乎没有——在某些情况下甚至没有可测量的——精度损失。我们还通过详细的性能分析（使用roofline模型）验证了我们的工作。该实现已上游至PyTorch和TorchAO，支持使用原生PyTorch工具进行开箱即用的部署。

*关键词*NLP ⋅ INT8量化 ⋅ CPU推理 ⋅ PyTorch ⋅ TorchInductor ⋅ Roofline模型 ⋅ Intel Xeon

## 1 引言
尽管大型语言模型吸引了近期的大部分关注，但基于BERT系列的小型NLP模型仍广泛部署于分类、排序和检索等生产任务（5 (https://arxiv.org/html/2608.18182#bib.bib2)；21 (https://arxiv.org/html/2608.18182#bib.bib3)；3 (https://arxiv.org/html/2608.18182#bib.bib22)）。由于其有利的成本概况和服务器CPU基础设施的普遍性，CPU推理在这些部署中继续发挥核心作用。在PyTorch生态系统中，对使用原生软件栈进行高效推理有着强烈的需求。在CPU上直接执行FP32通常并非最优，因为Transformer推理主要由矩阵乘法（GEMM）主导，这些操作计算和内存成本高昂。INT8量化是一种广泛采用的技术，可在保持可接受精度的同时减少模型占用空间和执行时间（12 (https://arxiv.org/html/2608.18182#bib.bib4)；15 (https://arxiv.org/html/2608.18182#bib.bib5)）。同时，现代Intel Xeon CPU提供了诸如VNNI和AMX等指令集支持，可以显著加速INT8 GEMM工作负载（9 (https://arxiv.org/html/2608.18182#bib.bib6)）。

在这项工作中，我们将SmoothQuant（一种旨在缓解Transformer中激活异常值的训练后量化方法）集成到原生PyTorch栈中，并优化其在现代Xeon CPU上的性能。我们的贡献包括：

- • 将端到端的SmoothQuant工作流集成到TorchAO的前端（18 (https://arxiv.org/html/2608.18182#bib.bib9)）并进行了改进以支持后续的图优化。
- • 在TorchInductor中添加了针对INT8 GEMM算子的图融合传递，消除了权重布局转换和后操作处理的运行时开销。
- • 为INT8 GEMM启用了内核选择机制，以在基于oneDNN的内核（16 (https://arxiv.org/html/2608.18182#bib.bib10)）和TorchInductor CPP后端中基于模板的GEMM内核之间进行选择，以实现峰值性能。
- • 在TorchInductor的CPP后端中添加了AVX512_VNNI微内核，以覆盖第三代Xeon CPU，同时为更新的平台提供基于AMX的内核。
- • 使用roofline模型和基准测试进行了全面的性能分析，以展示所提出工作流和优化的有效性。
- • 在BERT、DistilBERT和XLM-RoBERTa上，针对SQuAD（20 (https://arxiv.org/html/2608.18182#bib.bib20)）和MultiNLI（22 (https://arxiv.org/html/2608.18182#bib.bib21)）进行了基准测试，结果显示精度损失可忽略不计，并且与FP32基线相比吞吐量显著提升。

这项工作加强了服务器CPU上小型NLP模型的PyTorch原生推理路径，使用户能够以最小的努力部署SmoothQuant优化的INT8推理，而无需离开PyTorch生态系统或使用第三方工具。

## 2 相关工作
量化长期以来一直被研究作为提高推理效率的一种方式。仅整数量化因其工作（如12 (https://arxiv.org/html/2608.18182#bib.bib4)）而普及，随后的综述总结了训练后量化和量化感知训练的更广泛领域（15 (https://arxiv.org/html/2608.18182#bib.bib5)）。对于Transformer模型，有几种方法明确针对激活异常值或模型敏感性引起的量化困难。LLM.int8()（4 (https://arxiv.org/html/2608.18182#bib.bib11)）为大型语言模型中的异常通道引入了混合精度处理。GPTQ（6 (https://arxiv.org/html/2608.18182#bib.bib12)）专注于精确的仅权重训练后量化，而AWQ（14 (https://arxiv.org/html/2608.18182#bib.bib13)）使用激活感知的权重量化。与这些方法相比，SmoothQuant（24 (https://arxiv.org/html/2608.18182#bib.bib1)）将激活难度转移到权重上，以便在常见硬件上高效应用标准INT8算术。

多个软件栈支持CPU上的量化推理。Intel Neural Compressor（11 (https://arxiv.org/html/2608.18182#bib.bib14)）提供了广泛的量化支持，但通常依赖PyTorch本身或其他第三方工具进行部署。Hugging Face Optimum Intel（8 (https://arxiv.org/html/2608.18182#bib.bib15)）将面向Intel的优化集成到Hugging Face生态系统中，而ONNX Runtime（17 (https://arxiv.org/html/2608.18182#bib.bib16)）通过导出的ONNX图提供INT8推理。SGLang（25 (https://arxiv.org/html/2608.18182#bib.bib17)）和vLLM（13 (https://arxiv.org/html/2608.18182#bib.bib18)）是重要的大型语言模型服务系统，但它们侧重于解码器风格的大型语言模型服务，而非BERT系列编码器推理，并且不针对原生PyTorch中的SmoothQuant。

Intel Extension for PyTorch（IPEX）（10 (https://arxiv.org/html/2608.18182#bib.bib19)）以前在oneDNN之上提供了高效的SmoothQuant工作流，但该路径依赖于基于旧版torch.jit的优化；其维护此后已转向原生PyTorch 2.x编译和上游集成。我们的工作与这些方法不同，它提供了一个完全原生的PyTorch解决方案，将SmoothQuant、TorchAO、TorchInductor图融合和ISA感知的CPU内核选择集成在一个工作流中。这种集成允许在PyTorch生态系统内无缝部署优化的INT8推理。

## 3 背景与挑战
### 3.1 SmoothQuant与量化模式
激活异常值在Transformer模型中很常见，可能导致在直接进行INT8量化时产生较大的量化误差。SmoothQuant（24 (https://arxiv.org/html/2608.18182#bib.bib1)）通过重新分配激活和权重之间的比例因子来解决此问题：激活张量除以一个按通道的平滑因子，而相应的权重张量乘以相同的因子。这种转换在减少激活异常值的同时，在数学上保留了线性层。平滑因子是使用一个小的校准数据集离线得出的。平滑之后，模型更容易进行标准的INT8量化（12 (https://arxiv.org/html/2608.18182#bib.bib4)；15 (https://arxiv.org/html/2608.18182#bib.bib5)）。

由此产生的量化流程可以是静态的或动态的，具体取决于激活量化参数的确定时间。在静态量化中，比例因子和零点在校准期间获得，然后在图中固定。在动态量化中，它们在运行时计算，并且可能因输入而异。静态量化通常提供较低的运行时开销，而动态量化有时可以更好地保留精度，但代价是额外的执行工作。

### 3.2 PyTorch与TorchAO
PyTorch在研究和生产中被广泛使用（19 (https://arxiv.org/html/2608.18182#bib.bib7)）。PyTorch 2引入了TorchInductor（其API torch.compile可能更为人所知），它支持图捕获、优化和后端下沉，以提高执行效率（2 (https://arxiv.org/html/2608.18182#bib.bib8)）。在CPU上，TorchInductor的CPP后端是实现原生服务器级推理性能的关键路径。即时模式Inductor（AOTI）是TorchInductor的一个专用版本，它将模型编译为共享二进制文件，消除了Python开销，并使其更容易集成到生产系统中。TorchAO是PyTorch原生的量化库，为PTQ和QAT提供用户友好的API（18 (https://arxiv.org/html/2608.18182#bib.bib9)）。它通过模块化设计支持多种低精度数据类型和量化方法。

### 3.3 挑战
SmoothQuant的集成在正确性方面是直接的，但在CPU上实现高性能需要额外的工程工作。TorchAO以eager模式量化模型，没有后端感知的下沉或图级优化。这使得生成的模型正确，但可能不够快，无法满足生产CPU推理的要求。为了使用原生PyTorch实现峰值性能，量化模型必须使用TorchInductor编译，这引入了几个工程挑战：

- • SmoothQuant工作流与现有的TorchAO量化流共享公共代码，其中包含不利于高性能推理的低效操作。
- • 缩放的INT8 GEMM算子在建模和原型设计方面很方便，但性能有限。
- • 利用TorchInductor的内核选择机制是一个关键策略，以便不同的形状可以选择最佳的后端实现。
- • 支持仅AVX512_VNNI的旧平台（如Ice Lake）以及具有AMX支持的新型Xeon系列。

## 4 设计与实现
### 4.1 总体设计
我们的工作流构建在TorchAO的量化路径和TorchInductor的图优化及CPP后端之上。要将SmoothQuant应用于预训练模型，用户调用TorchAO的高级API，该API：

- • 通过插入观察器收集激活统计信息来为校准准备模型；
- • 通过在准备好的模型上运行校准数据来搜索平滑因子；
- • 用存储量化权重、比例因子和平滑因子的自定义张量子类替换`torch.nn.Linear`权重；以及
- • 通过`torch._int_mm`将高精度线性操作替换为INT8 × INT8 → INT32 GEMM，然后进行输出缩放。

量化后，线性层以INT8执行，而其余算子保持较高精度。然后，该模型通过TorchInductor编译，TorchInductor应用图优化并将结果下沉到CPP后端。我们添加了图传递，将围绕`torch._int_mm`的缩放INT8 GEMM子图融合到支持oneDNN的算子`torch.ops.onednn.qlinear_pointwise`中。我们还扩展了TorchInductor的基于模板的GEMM路径，加入了AMX和AVX512_VNNI微内核，以便下沉时可以在基于oneDNN的内核和生成的模板内核之间进行选择。

图1：基于TorchAO和TorchInductor构建的SmoothQuant工作流。TorchAO准备和量化模型，TorchInductor融合INT8子图，CPP后端在oneDNN和基于模板的GEMM内核之间进行选择。生成C++代码并进行编译以供部署。

### 4.2 改进通用量化流的实现
SmoothQuant的集成与现有的TorchAO量化流共享许多公共代码，我们进行了几项改进以提高运行时性能：

- • 应用平滑因子时，我们预先计算因子的倒数，以避免运行时的重复除法。
- • 通过忽略一些不必要的张量扩展操作，我们避免了张量拷贝。
- • 我们忽略虚拟零点以节省计算。

### 4.3 克服`torch._int_mm`的局限性
我们使用oneDNN实现了`torch._int_mm`的CPU版本，以支持TorchAO中SmoothQuant流程等情况下的INT8 GEMM。然而，该原语本身存在三个重要局限性：

- • **运行时权重布局转换**：在Xeon上进行高效INT8 GEMM通常需要为VNNI和AMX指令定制的阻塞权重布局。在运行时将权重从普通布局转换为阻塞布局代价高昂。
- • **缺少比例因子和后操作处理**：该算子计算INT8 × INT8 → INT32 GEMM而没有epilogue，因此比例因子和其他后操作需要单独处理。
- • **oneDNN设置开销**：oneDNN需要辅助数据结构，其每次调用的构造可能成为不小的开销。

我们通过扩展TorchInductor中的图融合能力来解决这些问题。匹配以下模式及其变体：
`input_reshape -> int_mm -> dtype_convert -> mul -> mul -> output_reshape -> add_bias`
其中`int_mm`是INT8 GEMM算子，第一个`mul`应用激活缩放，第二个`mul`应用权重缩放，`add_bias`添加偏置项。匹配的子图然后被单个`torch.ops.onednn.qlinear_pointwise`算子替换。这个融合算子接受阻塞布局的权重，在内核内应用激活和权重缩放作为后操作，并在预热后缓存oneDNN元数据。该内核还可以通过TorchInductor现有的融合能力与其他算子（如GeLU（7 (https://arxiv.org/html/2608.18182#bib.bib25)））进一步融合。同时在图中添加了相应的权重布局转换算子，将权重从普通布局转换为阻塞布局，该算子会在TorchInductor的常量折叠传递中进行评估和折叠。

### 4.4 面向最大性能的基于模板的GEMM内核
尽管`qlinear_pointwise`在许多情况下表现良好，但它并非在所有矩阵形状下都达到最优。因此，我们扩展了TorchInductor的基于模板的GEMM内核作为替代下沉目标。具体而言，我们使用AMX和AVX512_VNNI intrinsics实现了INT8 GEMM微内核，并将它们集成到CPP后端的生成内核路径中。在下沉期间，TorchInductor的内核选择机制会生成候选内核、编译它们、进行微基准测试，并选择最快的实现供将来执行。这个一次性的选择过程使系统能够在oneDNN的通用性和模板生成内核的形状特化性之间取得平衡。

### 4.5 支持Ice Lake：s8s8到u8s8转换
TorchAO将激活和权重都量化为有符号INT8（s8s8）。然而，第三代Xeon（Ice Lake）不提供原生的s8 × s8 GEMM指令；AVX512_VNNI要求激活为无符号INT8（u8），即u8 × s8。对于`Y = X · W^T`的GEMM，因此我们将激活张量`X`移动一个常数：
`X_{u8} = uint8(int32(X_{s8}) + 128)`。
为了保持数学等价性，该移动被视为一个零点，并在GEMM结果中进行补偿。对于权重矩阵`W ∈ Z^{N×K}`，补偿项为：
`c = (∑_{k=1}^K W_{:,k}) · 128`。
而GEMM结果计算为：
`Y = X_{u8} · W^T - c`。
由于移动和权重都是常数，权重转换和补偿向量在推理前预计算。尽管...