[论文解读] I-BERT: Integer-only BERT Quantization

Sehoon Kim, Amir Gholami+3 more|arXiv (Cornell University)|Jan 5, 2021

Topic Modeling参考文献 86被引用 16

一句话总结

I-BERT 提出了一种针对 BERT 和 RoBERTa 模型的新型纯整数量化框架，支持仅使用整数运算的端到端推理，完全消除浮点运算。通过为 GELU、Softmax 和 LayerNorm 等非线性层引入轻量级多项式近似，I-BERT 在 GLUE 基准测试中实现了相当或略高的准确率，并在 T4 GPU 上实现了最高达 4.0× 的加速。

ABSTRACT

Transformer based models, like BERT and RoBERTa, have achieved state-of-the-art results in many Natural Language Processing tasks. However, their memory footprint, inference latency, and power consumption are prohibitive efficient inference at the edge, and even at the data center. While quantization can be a viable solution for this, previous work on quantizing Transformer based models use floating-point arithmetic during inference, which cannot efficiently utilize integer-only logical units such as the recent Turing Tensor Cores, or traditional integer-only ARM processors. In this work, we propose I-BERT, a novel quantization scheme for Transformer based models that quantizes the entire inference with integer-only arithmetic. Based on lightweight integer-only approximation methods for nonlinear operations, e.g., GELU, Softmax, and Layer Normalization, I-BERT performs an end-to-end integer-only BERT inference without any floating point calculation. We evaluate our approach on GLUE downstream tasks using RoBERTa-Base/Large. We show that for both cases, I-BERT achieves similar (and slightly higher) accuracy as compared to the full-precision baseline. Furthermore, our preliminary implementation of I-BERT shows a speedup of 2.4-4.0x for INT8 inference on a T4 GPU system as compared to FP32 inference. The framework has been developed in PyTorch and has been open-sourced.

研究动机与目标

• 解决在边缘设备和纯整数硬件上运行 BERT 等大型 NLP 模型时浮点推理效率低下的问题。
• 克服在不损失准确率的前提下，仅使用整数算术实现非线性运算（如 GELU、Softmax、LayerNorm）的挑战。
• 实现在无浮点单元的硬件（如 ARM Cortex-M 和 NVIDIA Turing Tensor Cores）上部署 BERT 类模型。
• 开发一种量化方案，在保持模型准确率的同时，最大化在纯整数加速器上的速度与效率。
• 证明在现代硬件上，纯整数推理在延迟和能效方面可优于混合精度推理。

提出的方法

• 提出 i-GELU，一种用于 GELU 的二阶多项式近似方法，可在最大误差为 1.8×10⁻² 的前提下实现高精度的纯整数计算。
• 设计一种基于轻量级多项式的纯整数 Softmax 近似方法，最大误差为 1.9×10⁻³。
• 通过已知的整数开方算法实现纯整数 LayerNorm，避免使用浮点运算。
• 将这些近似方法集成到完整的量化流水线中，其中所有参数和激活值均以 INT8 表示，矩阵乘法使用 INT32 累加。
• 采用混合量化策略：INT8 权重和激活值，INT32 累加，非线性运算全程使用整数运算，并通过后处理重量化保持精度。
• 确保计算图中所有运算始终保持在整数格式，避免任何向浮点数的类型转换。

实验结果

研究问题

• RQ1 是否可以仅使用整数算术准确近似 GELU、Softmax 和 LayerNorm 等非线性运算？
• RQ2 在使用纯整数推理时，为保持 BERT 类模型的准确率，可接受的近似误差水平是多少？
• RQ3 在专为整数运算优化的硬件（如 T4 GPU 和 ARM 处理器）上，纯整数推理是否能实现显著加速？
• RQ4 与全精度或混合精度基线相比，纯整数量化是否会导致准确率下降？
• RQ5 所提出的 i-GELU 近似方法在准确率和效率方面相较于其他纯整数 GELU 近似方法表现如何？

主要发现

• 在 Tesla T4 GPU 上，I-BERT 对 BERT-Large 模型使用 INT8 推理时，最大推理加速达到 4.0×，相比 FP32 推理。
• 在不同序列长度和批量大小下，BERT-Base 的平均延迟加速为 3.08×，BERT-Large 为 3.56×。
• I-BERT 维持或略微提升了 GLUE 平均得分——相比全精度基线，RoBERTa-Base 和 RoBERTa-Large 分别高出 0.3 和 0.5 分。
• 在四个下游任务上，i-GELU 近似方法平均比 h-GELU 近似高出 0.7 分，且与全精度 GELU 相比无准确率损失。
• 消融实验表明，i-GELU 在某些情况下甚至优于 h-GELU 和全精度 GELU，提示其优势可能不仅限于近似本身。
• 该框架已在 PyTorch 中完全开源，支持在纯整数硬件上部署，使无浮点单元的边缘设备也能实现高效推理。