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 也验证了量化模型的效果。
梯度通信量化
由于在多机多卡场景下通信瓶颈更加明显，所以梯度通信量化主要应用在分布式训练场景。因此 LightSeq 在 2 机 8 卡的 A100 上进行了分布式训练的速度测试。
可以看出，梯度通信量化的训练加速效果整体上随着输入数据的增大而减弱。这主要是因为随着输入数据的增大，计算时间占比升高，梯度通信时间占比减少，梯度量化的收益也随之减小。
LightSeq 还额外增加了不同数量网卡（NIC）下的训练速度测试。可以看到使用梯度通信量化的分布式训练速度相比原始的 LightSeq fp16 有大幅度提升。
*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。
发布文章
