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

发布人：CV研究院时间：2023-01-19 来源：工程师发布文章

SC22 接收论文：https://sc22.supercomputing.org/presentation/?id=pap211&sess=sess154
代码地址：https://github.com/bytedance/lightseq

如何继续提升速度？降低计算精度是比较直接的方法。2017 年以来，fp16 混合精度技术 [2] 获得了广泛应用。在对模型效果无损的前提下，将模型训练和推理的速度提升了 50% 以上。而为了维持模型效果，更低精度的方法（例如 int8）通常需要使用如下传统方案：

首先使用 fp16 混合精度将模型训练至收敛；
然后在模型计算密集型算子的权重、输入和输出位置处，插入伪量化结点，进行量化感知训练；
最后将带有伪量化结点的模型计算图转换到专用的 int8 推理引擎中，进行服务部署和模型推理。

虽然在多数任务上，上述方案可以实现模型效果无损，但还是存在以下问题：

使用方法复杂。例如要多一次量化感知训练 [4] 的过程，并且带有伪量化节点的计算图转换复杂。
训练速度慢。由于目前流行的深度学习框架不支持 int8 精度，所以量化感知训练需要插入 fp16 的伪量化结点来模拟 int8 量化，导致量化感知训练反而比 fp16 混合精度训练慢 2-3 倍。
推理部署难且加速比低。对比 fp32、fp16 等类型，int8 硬件和底层软件库优化相对滞后。例如在 NVIDIA GPU 上，int8 矩阵乘法加速受限于硬件架构和特定 shape，实际加速比远远低于理论值。

在下文中，如无特殊说明，量化都是指的 int8 精度的量化。

针对这些问题，字节跳动推出了全新版本的 LightSeq GPU 量化训练与推理引擎。支持 Transformer 系列模型的量化训练与推理，并做到了开箱即用，用户友好。LightSeq 快准狠地实现了 int8 精度的量化训练和推理：

快：A100 多卡训练最高加速 5.2 倍，T4 单卡推理最高加速 8.9 倍。
准：训练和推理效果基本无损。
狠：相同数据量下，显存占用最高减少 68%，模型存储空间减少 75%。

总体来说，LightSeq 新版量化训练与推理引擎具有如下几个优点：

1. 丰富的支持

支持完整的 Transformer 模块和多种解码算法，支持 Transformer、BERT、GPT、BART、ViT 等多种模型结构，支持 Fairseq、Hugging Face、NeurST 等多种训练框架接入量化训练、导出模型以及量化推理，提供了丰富的样例供用户参考。

2. 卓越的性能

相比于 fp16 精度的 LightSeq 推理引擎，int8 量化还可以进一步加速最高 70%，相比于 PyTorch 推理更是达到了最高 8.9 倍的加速比。同时显存占用相比 fp16 推理引擎降低了 30% 左右，模型存储空间只需要原来的四分之一。最后经过多个任务的验证，推理效果几乎无损。

3. 便捷的使用

LightSeq 已经针对多个训练库进行了量化支持，可以一键开启量化训练，然后轻松导出为 LightSeq 支持的模型格式，最后实现量化推理。除此之外，LightSeq 还支持训练后量化，无需额外训练即可体验量化推理。

使用方法

如上图所示，为了最大程度减小量化带来的损失，首先需要用 fp16 精度训练一个浮点数模型，将模型效果训到最好。然后开启量化进行 finetune，得到微调过的量化模型，此时模型效果已经基本恢复到浮点数模型的水平。接着将量化模型转换为 LightSeq 支持的 PB 或者 HDF5 模型格式，最后用 LightSeq 进行量化推理。

安装方法

LightSeq 安装非常简单，只需要一行命令即可：

pip install lightseq

量化训练

LightSeq 支持 Fairseq、Hugging Face、NeurST 等训练框架的量化接入，同时也可以自定义模型并开启量化训练。以 encoder 层为例，只需要先定义浮点数模型，然后开启量化即可：

from lightseq.training import LSTransformerEncoderLayerfrom lightseq.training.ops.pytorch.quantization import enable_quant
config = LSTransformerEncoderLayer.get_config(    model="bert-base",    max_batch_tokens=4096,    max_seq_len=512,    fp16=True,    local_rank=0,)layer = LSTransformerEncoderLayer(config)# 开启量化layer.apply(enable_quant)

量化推理

LightSeq 提供了便捷的 python 推理接口，只需要三行代码即可实现快速的量化推理：

import lightseq.inference as lsi
model = lsi.QuantTransformer(pb_path, batch_size)result = model.infer(input)

此外 LightSeq 还提供了 BERT、GPT、ViT 等模型的 python 接口，分别调用 QuantBert、QuantGpt 和 QuanVit 即可体验。

梯度通信量化

LightSeq 支持 Transformer 模型的梯度通信量化[5]，使用 Fairseq 或者 Hugging Face 即可轻松开启分布式量化训练，并同时支持浮点数模型和量化模型。在构建模型后，只需要为模型注册一个 communication hook 即可开启梯度通信量化，再开始训练过程。

from lightseq.training.gradient_comm_quantization import encode_and_decode, GCQStatefrom torch.nn.parallel import DistributedDataParallel 
# model could be from Fairseq or Hugging Face, wrapped by DDPmodel = DistributedDataParallel(model)state =  GCQState(process_group)# register hookmodel.register_comm_hook(state=state, hook=encode_and_decode)

性能测试

LightSeq 在多个任务上测试了量化训练、量化推理和梯度通信量化的速度，并且分析了显存占用情况和量化模型的效果。

量化训练速度

LightSeq 在 8 张 A100 显卡上进行了训练实验，主要对比对象是 Fairseq 的 Transformer、Hugging Face 的 BERT、GPT2 和 ViT。

可以看出，四种模型结构加速趋势都是类似的，加速比都会随着数据量的增大而减小，原因有三点：

随着数据量的增大，矩阵乘法 GEMM 的占比会明显增加，因此 PyTorch QAT 增加的额外的伪量化结点时间占比会逐渐减小，最后速度会和 PyTorch fp16 无限接近。
与此同时，随着 GEMM 占比升高，LightSeq fp16 自定义算子的提速效果也逐渐减小，因此时间上也会和 PyTorch fp16 无限接近。
由于 Ampere 架构显卡上 int8 GEMM 在 shape 较小时甚至不如 fp16 GEMM 快，在大 shape 下才能稍快一点，因此随着数据量增大，LightSeq int8 也会无限接近 LightSeq fp16 的速度。

量化推理速度

LightSeq 在单张 T4 显卡上进行了推理实验，主要对比对象是 Hugging Face 的 Transformer、BERT、GPT2 和 ViT。

可以看出，随着输入数据量的增大，LightSeq 与 PyTorch 的差距会逐渐减小，这也是 GEMM 占比升高造成的。比较 LightSeq fp16 和 LightSeq int8，可以看出随着数据量的增大，LightSeq int8 越来越快。这是因为在 T4 显卡上，int8 GEMM 的加速会随着 shape 的增大而有明显增加。因此在 T4 显卡上进行量化推理时，输入数据量越大，加速效果越好。

LightSeq 还针对机器翻译多个语向和多个测试集，测试了不同 batch size 下，LightSeq int8 推理相对于 LightSeq fp16 推理的加速比，实验同样是在单张 T4 显卡上进行的，采用的模型都是标准的 Transformer-Big。

可以得到和上文中相同的结论，随着 batch size 的增大，量化推理的加速比会逐渐升高。相比于 LightSeq fp16，最高还可以再加速近 70%，这极大地缩短了线上翻译模型的推理延时。

最后如上图所示，为了展示自动 GEMM 调优技术的效果，LightSeq 测试对比了 A100 显卡上 Transformer 和 BERT 模型 fp16、int8 调优前和 int8 调优后的延时。可以看出调优前某些 shape 的 int8 GEMM 速度甚至比 fp16 还要慢，而调优后全面超越了 fp16。

显存占用

LightSeq 分析了不同 batch size 下，量化模型相对于浮点数模型显存占用的加速比。可以看出随着 batch size 的增大，量化模型的显存占用优势更明显，最高可以减少 30% 左右。而 LightSeq fp16 引擎相对于 PyTorch 模型也极大程度减少了显存占用，因此 LightSeq int8 引擎最终能够减少最多 68% 左右的显存。

量化模型效果

针对机器翻译多个语向和多个测试集，LightSeq 测试了量化模型推理相对于浮点数模型 BLEU 的损失，采用的模型都是标准的 Transformer-Big。

在数据量较大的语向 en2zh 上，LightSeq int8 相对 BLEU 损失较大些，最大达到了 - 0.4。而在数据量较小的语向 en2es 上，LightSeq int8 不仅没有任何效果损失，反而比浮点数模型更好。总体而言，int8 量化模型的平均 BLEU 相比浮点数模型基本无损。在 GLUE 和 SQuAD 等多个任务上，LightSeq 也验证了量化模型的效果。

梯度通信量化