最高加速9倍！字节跳动开源8比特混合精度Transformer引擎（2）

量化技术

int8 量化的加速收益主要来自如下几个方面：

1. GEMM 精度从 fp16 降低到 int8 后，计算时间缩短；

2. 自定义算子采用 int8 输入输出后，数据读写时间缩短；

3. 梯度采用 int8 存储后，多机之间通信时间缩短。

以 Transformer 模型为例，经过 LightSeq fp16 引擎加速后，自定义算子时间大大缩短，而 GEMM 时间占比提升到了 90% 左右，因此优化的重点转移到了 GEMM 提速。将 fp16 GEMM 替换为 int8 GEMM 不仅可以缩短 GEMM 时间，还可以减小前后算子的输入输出位宽，从而减小读写数据的时间。最后多机训练的瓶颈主要在梯度的通信，将梯度量化为 int8 精度可以大大加快分布式训练的速度。

量化原理

为了弥补量化带来的精度损失，通常需要用量化感知训练来模拟量化过程。如上图所示，量化感知训练就是将 float GEMM 的两个 float 输入分别做一遍量化和反量化（称之为伪量化结点），离散化成分段的浮点数输入，然后进行 float GEMM 运算。得到结果后再次进行量化与反量化，得到最终的浮点数结果。而量化的过程是不可导的，因此需要用 STE 方法来估计量化参数的梯度。之所以量化感知训练中需要插入伪量化结点，然后用 float GEMM 去模拟量化过程，是因为 TensorFlow 和 PyTorch 等训练框架不支持 int8 GEMM。

而 LightSeq 量化训练直接采用 int8 GEMM 来真实还原量化过程，因此相比传统的实现要更快，且更加节省显存。在推理的时候，同样采用离散化后的整数进行 int8 GEMM 运算，最后再反量化回浮点数结果。量化推理过程和量化训练完全一致，并且和传统的量化感知训练是完全等价的。

量化位置

整个量化 Transformer 的网络结构如上图所示，红色箭头表示需要加上量化和反量化结点的位置。

首先所有 int8 GEMM 的输入和输出都需要进行量化。由于 int8 GEMM 的 shape 限制，部分 GEMM（例如注意力分数的计算）仍然采用 float GEMM。此外第二层 FFN 的 GEMM 采用的是 int32 的输出，因为它的 GEMM 输入是 ReLU 激活函数的输出结果，只包含正数，非对称，因此如果采用 int8 输出的 GEMM，将无法反量化为正确的浮点数结果。

然后所有的模型权重 weight 都需要存储为 int8 类型，因此需要对 weight 做量化。而权重 bias 参数量较小，无需量化，保留 float 精度反而可以提升模型效果。

最后需要对 decoder 端的 cache 进行量化。因为在推理时，decoder 端的 cache 需要频繁进行读写，因此将 cache 量化为 int8 可以大大加快解码的速度。

量化策略

将一个浮点数矩阵量化为 int8 整数矩阵有很多方法，LightSeq 采用的是对称量化，即将正负数范围对称的浮点数区间等比例地映射到整数区间 [-127, 127] 上。

而实际上浮点数矩阵的数值范围通常并不对称，存在极少的离群值。如果直接按照离群值的范围来量化矩阵，会影响到量化后的精度，所以需要先对矩阵进行数值截断。

LightSeq 采用 PACT 方法进行截断[6]，将截断的范围当作模型可学习的参数，然后利用 STE 算法去估计参数的梯度，并进行反向传播优化。根据实践经验，权重 weight 的初始截断范围设为[-1, 1]，中间结果的初始截断范围设为[-16, 16]，可以在大部分任务上达到最好的效果。最后经过截断范围和其他模型参数的联合优化，量化模型的效果可以达到基本无损。

梯度通信量化

针对分布式训练场景，LightSeq 推出了梯度量化压缩技术。即对浮点精度的梯度进行 int8 量化，以减少梯度通信的时间消耗，从而加速训练，这就是梯度通信量化（GCQ）。

如上图所示，梯度通信量化的主要流程如下：

1. 计算每张卡上各自梯度的截断范围；

2. 对截断范围执行 all-reduce max 操作；

3. 每张卡使用统一的截断范围对各自梯度进行 int8 量化；

4. 对 int8 梯度执行 all-reduce sum 操作；

5. 每张卡对 all-reduce 后的梯度进行反量化，还原为浮点数梯度，并进行参数更新。

为了解决 int8 梯度在 all-reduce 过程中溢出的问题，LightSeq 首先将每张卡上的浮点数梯度除以卡数，再使用除之前的截断范围进行量化，最后进行 all-reduce 操作。这样每张卡上量化后的 int8 整数 all-reduce 完就不会溢出，但是单卡实际用于量化的比特数也因此而减少，所以目前方案在 2 机 8 卡效果几乎无损，但随着卡数的上涨，训练效果会有所下降。以 en2de 和 en2fr 翻译任务为例，在 4 机 8 卡上进行分布式量化训练，BLEU 值分别会下降 0.4 和 1.5 左右。未来 LightSeq 将会持续探索更好的方法来解决这一问题。

通用技术

除了上一章节中提到的量化技术以外，此次更新 LightSeq 还提出了几种通用的优化技术，不仅可以应用在量化模型中，也适用于其它所有精度模型的训练与推理。

算子融合

上图是 encoder 模块量化训练的计算图，LightSeq 将两次 GEMM 运算之间的所有操作融合成一个算子[7]，减少了 kernel 调用的次数，因此减少了总的计算时间。

图中黄色矩形表示 int8 GEMM，绿色矩形表示 float GEMM。这里采用 float GEMM 是由于 shape 的限制，不适合使用 int8 GEMM 加速。红色箭头表示流动数据的类型是 int8，绿色箭头表示第二层 FFN 的 GEMM 输出是 int32 数据类型。int8 GEMM 输入输出的量化与反量化操作都被融合到了前后 kernel 里，这不仅可以减少数据搬运，还可以减小显存占用。

在推理时，LightSeq 还针对 decoder 做了优化。如上图所示，在计算 self-attention 时，注意力得分的维度是(batch size, 1, sequence length)。因此在计算 value 乘积时，可以不采用 GEMM 运算，而直接手写加权求和的算子，从而将图中虚线框中的计算融合成一个 kernel。

自动显存管理

模型量化引入了更复杂的张量类型和张量依赖关系，这给显存管理带来新的挑战。为此，LightSeq 设计了新的显存管理机制。如上图所示，主要包括以下过程：

1. 训练启动前，根据每个算子的拓扑依赖关系，自动计算每个张量的生命周期及显存空间大小。其中，包含动态维度的张量按照此维度的最大量进行计算，例如机器翻译任务中的最大句长和最大 batch 句子数量。这些最大量在训练前已被指定；

2. 张量确定生命周期和大小后，分析显存复用关系。其中，无生命周期重合的张量可以共用一片显存空间，所有显存空间都是无数据类型的，可以被分配到任意数据类型的张量上；

3. 根据张量显存复用关系，申请多段显存空间，为每个张量分配实际的显存起止地址。

张量显存复用的分析，LightSeq 借鉴了论文 [3] 中提出的 Greedy by Size for Offset Calculation 方法，做了三个改进：

1. 支持了整个训练过程的显存复用（forward/backward）；

2. 不同数据类型能做到显存复用（int8/fp16/fp32）；

3. 在多段显存空间上容纳所有张量，而非一段非常大的显存空间，这样能有效提升显存利用率。

自动 GEMM 调优

LightSeq 的 int8 GEMM 采用了 NVIDIA 的 cuBLASLt 库，这也是目前 NVIDIA 显卡上最为高效的矩阵运算库。但是输入数据的 shape 或者显卡不同的话，GEMM 所采用的最优配置（例如数据排布、GEMM 算法等等）也可能不同，因此需要进行自动选取。LightSeq 采取的自动调优方案如下：

1. 在多种型号显卡上（例如 T4 和 A100）进行不同 shape 的 GEMM 最优配置搜索，并将结果保存到配置文件中，用户只需要下载即可；

2. 模型初始化时，加载对应型号显卡的配置文件，解析并保存到键值对为 (shape, 最优配置) 的字典中。如果没有对应型号显卡的配置文件，或者没有需要的 GEMM shape，那么用户可以选择自己搜索并保存，或者直接使用默认配置；

3. 模型前向或后向计算时，根据输入的 shape 在字典中寻找最优配置，然后进行 GEMM 计算。如果没有找到对应的 shape，那么直接采用默认的配置。

未来工作

未来 LightSeq 还将继续探索移动端的低精度量化、反向传播中梯度的量化、大模型量化等方向。

引用

[1] Wang, Xiaohui, et al. "LightSeq2: Accelerated training for transformer-based models on gpus." arXiv preprint arXiv:2110.05722 (2021).

[2] Micikevicius, Paulius, et al. "Mixed precision training." arXiv preprint arXiv:1710.03740 (2017).

[3] Pisarchyk, Yury, and Juhyun Lee. "Efficient memory management for deep neural net inference." arXiv preprint arXiv:2001.03288 (2020).

[4] Jacob, Benoit, et al. "Quantization and training of neural networks for efficient integer-arithmetic-only inference." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[5] Alistarh, Dan, et al. "QSGD: Communication-efficient SGD via gradient quantization and encoding." Advances in neural information processing systems 30 (2017).

[6] Choi, Jungwook, et al. "Pact: Parameterized clipping activation for quantized neural networks." arXiv preprint arXiv:1805.06085 (2018).

[7] Wang, Xiaohui, et al. "LightSeq: A high performance inference library for transformers." arXiv preprint arXiv:2010.13887 (2020).