本文内容源自发表于计算机人工智能领域顶会 EMNLP 2023（The 2023 Conference on Empirical Methods in Natural Language Processing），由快手、腾讯等单位联合共同完成。
SAMP: A Model Inference Toolkit of Post-Training Quantization for Text Processing via Self-Adaptive Mixed-Precision
论文作者：
论文链接：
摘要
AI 进入大模型时代，模型的计算性能是一个长期被关注的问题。目前工业界推理引擎已经验证半精度浮点（FP16）和 8 位整型（INT8）计算可以大大缩减计算位宽，从而加快模型推理速度。但已知的 INT8 量化方法使用门槛较高，且精度损失大，会导致模型效果大打折扣。我们推出了一个自适应混合精度推理库（SAMP），通过混合精度架构自动控制浮点型运算和整型运算，使得模型在量化推理加速的同时，保证计算精度。
文本理解是自然语言处理（NLP）领域的基本任务之一，例如在信息检索、对话系统、情感识别、摘要生成、语言模型等多个应用场景中。基于 Transformer 的模型在一系列下游任务中不断刷新着最佳记录（例如 BERT、XLNet、Google T5 等）。然而在一些大型工业系统中，由于在推理阶段缺乏高 GPU 占用考虑和良好的内存管理，训练框架（例如 TensorFlow 或 PyTorch）并不是模型部署的最佳选择。
用于深度学习模型推理加速的传统工具（例如 NVIDIA TensorRT、TurboTransformers 和 LightSeq）主要针对固定大小或可变长度输入而设计。这些工具的优化概念主要考虑在线计算系统中的内存管理、算子融合和一些数据剪枝技术。
计算精度方面，大多采用单一精度计算（使用纯浮点数计算或纯整型计算），加速性能有限。其中比较亮眼的是，NVIDIA 开发的 FasterTransformer 对 Transformer 模型进行了定点加速（将所有 Transformer 层使用 INT8 量化），与浮点计算相比取得了优异的加速效果。但这种在所有 Transformer 层进行全 INT8 量化的方式，模型推理结果很难达到原有浮点计算的精度，在很多任务中精度损失大，难广泛应用。
为了解决这些问题，我们提出了一种推理工具包 SAMP，其中包含自适应混合精度编码器和一系列先进的融合策略。自适应混合精度编码器可以在大量的通用矩阵乘法（GEMM）运算和 Transformer 层中找到最佳浮点定点混合精度组合方式，使模型推理的性能最贴近用户需求（计算精度或推理效率）。最终，混合精度计算取得了比全定点计算更好的计算精度。融合策略对 embedding 算子和量化相关计算操作进行融合改进，使得 CUDA 内核调用减少一半。
与此同时，SAMP 是由 C++ 编程语言实现的端到端工具包（从 Tokenizer 到 Embedding、Encoder 及下游任务），具有出色的推理速度，也降低了训练后量化推理的工业应用门槛。表 1 显示了与同类系统相比的创新点。
以下是 SAMP 工作的主要贡献（核心亮点）：
1. 自适应：SAMP 在训练后量化推理方法中平衡计算精度和延迟性能。用户可以针对不同的任务选择合适精度和推理延迟的混合精度配置。SAMP 还可通过自适应分配方法推荐给用户最佳的量化组合模式。
2. 推理效率：在较宽的精度范围（浮点到定点）中，SAMP 显示出比其他推理工具包更好的推理加速。在中文语言理解测评基准（CLUE）分类任务数据集中，与 FasterTransformer 相比，SAMP 实现了高达 1.05-1.15 倍的加速。
3. 灵活性：SAMP 涵盖众多下游任务，如分类、序列标记、文本匹配等。Target 模块是可扩展的并且可以灵活定制。它对用户友好且对平台依赖性较低。SAMP 支持 C++ 和 Python API，仅需要 CUDA 11.0 或更高版本即可。另外，SAMP 也提供了许多模型转换工具，支持不同格式模型之间相互转换。
SAMP 端到端推理库由四大模块组成：分词器（Tokenizer）、词嵌入层（Embedding）、编码器（Encoder）、下游目标层（Downstream Target）。各模块有不同亮点：
-
Tokenizer：支持中英文的多粒度分词（如：字符级、词语级以及常见的BertTokenizer），该模块通过 C++ 多线程并行加速，计算速度很快；
-
Embedding：SAMP 将 token_embedding、position_embedding、segment_embedding 三个 kernel 融合为一个 kernel，减少了 CUDA 的调用次数。
-
Encoder：在 Transformer 编码器里，SAMP 会通过精度衰减感知算法推荐出一个能使模型推理计算效率和精度兼得的混合精度配置。
-
Downstream Target：SAMP 支持各种各样的 NLP 下游任务。
SAMP 提供两种混合精度推理模式：全量化模式（Fully-Quant mode）和FFN量化模式（Quant-FFN-Only）。
1. 全量化模式（Fully-Quant mode）
全量化模式下，计算 Transformer 的 CUDA kernel 内部通过量化、反量化以及算子融合等方法，将所有 kernel 间的数据流动转为 8 位整型，大大缩减了传输位宽，推理速度达到了最快。
2. FFN量化模式（Quant-FFN-Only）
Quant-FFN-Only 模式是仅对 Transformer 中的 Feed-Forward 子网络做 8 位整型量化，保留了 Multi-Head-Attention 网络的浮点数计算（FP32 或 FP16）。这种方式因为精度损失较为平滑，在精度与推理效率的置换过程中，更容易被 SAMP 推荐。
3. SAMP混合精度推理库的使用方法
当用户明确给出模型精度要求或时延要求时：SAMP会给出最接近用户需求的混合精度推理参数和对应模式，进行计算。
用户未明确给出模型精度要求或时延要求时：SAMP将通过下图算法1中的精度衰减感知算法，自动推荐一个混合精度模式，进行计算。
成果展示
1. 绝对速度
我们测试了 SAMP 推理库在“全精度浮点（FP32）”、“半精度浮点（FP16）”、“全定点（INT8）”三种精度下的模型计算效率，选取主流的 PyTorch 和最新版 FasterTransformer 进行比较。
我们选取了较为常见的 batch_size 和文本长度组合。实验结果来看，大多情况下 SAMP 的计算效率是比 Nvidia 官方推理库 FasterTransformer 更快。
2. 精度、加速效果兼得
另外，我们通过中文语言理解测评基准（CLUE），选取一些文本分类任务——AFQMC（蚂蚁金服问答库）、IFLY-TEK（长文本分类）和 TNEWS（新闻短文本分类）等，对 SAMP 做了“计算精度 vs 计算加速比”的置换实验。
相比 PyTorch-FP16 和 FasterTransformer-FP16/INT8，SAMP 提供了 SAMP-FP16、SAMP-Fully-Quant 和 SAMP-Quant-FFN-Only 三种模式，其中后两种模式实现了 Transformer 层维度的混合精度计算。
上表中下划线的结果表示精度衰减感知算法推荐给用户的混合精度结果。以 AFQMC 任务为例，相比 PyTorch-FP16，SAMP-Fully-Quant 模式的模型精度从 0.7337 降至 0.6874，推理速度提升达到 4.7751 倍，FasterTransformer-INT8 的模型精度从 0.7337 降至 0.5773，推理速度提升只达到 3.4990 倍。同样，在另外两个实验结果不难得出类似结论—— SAMP 在“计算精度 vs 计算加速比”置换实验中取得了不错的效果。
1. 解决了现有后量化（PTQ）推理工具在工业应用中精度损失大的问题；
2. 推动了后量化（PTQ）技术在 NLP 多个下游任务中大规模使用；
3. 推理库轻量、灵活，对用户友好，支持用户自定义任务目标。
原文始发于微信公众号（快手技术团队）：EMNLP 2023 | SAMP：基于自适应混合精度的训练后量化模型推理库
- 左青龙
- 微信扫一扫
-
- 右白虎
- 微信扫一扫
-
评论