混合量化:一种用于减少模型尺寸且精度影响最小的全新量化方法

量化感知训练和推理的统一套件

在多 GPU 上运行大型模型可能有助于减少延迟,但会显著增加部署成本,尤其是在模型尺寸不断增大的情况下。为了缓解这个问题,我们采用了模型压缩技术并引入了一种新的方法,该方法以最小的精度影响对 Transformer 网络进行量化。我们的技术通过在更少或相同数量的 GPU 上使用定制的推理内核,实现了与 FP16 模型相似或更好的性能。

我们的方案灵活,因为用户可以实验任何量化配置,例如用于量化精度的目标位数以及模型在训练期间被量化的调度。此外,我们将 FP16 和量化精度结合在一起,形成了一种混合精度机制,以平滑从高精度到低精度的过渡。最后,我们使用参数的二阶梯度(特征值)来调整训练期间的量化调度。

量化方法

应用量化主要有两种方法:对已训练模型进行离线量化和量化感知训练 (QAT),它在训练期间减少数据精度。与前一种方案不同的是,QAT 在训练优化过程中考虑了精度损失的影响,对模型进行训练。这将显著提高量化模型的精度。MoQ 构建在 QAT 方法之上,不同之处在于我们使用混合精度来训练模型,使其朝着目标量化方向发展,同时还定义了一个用于降低精度的调度。

所有现有的 QAT 方法都从训练开始到结束,以特定精度(位数)对模型进行量化。然而,即使使用相对较高的量化精度(8 位),也会出现模型精度的下降,这对于一些下游任务可能是不可接受的。例如,Q8BERT 工作尝试对 BERT 网络进行 QAT,这在某些任务中可以获得良好的精度,但在其他任务(如 SQuAD)中 F1 分数会下降 0.8%。其他技术,例如 Q-BERT,在量化参数矩阵时使用分组量化,分组大小较大 (128),以获得更高的精度,但它们仍然不如基线方法。

在这里,我们介绍 MoQ 作为一种灵活的线性量化解决方案,允许用户在模型训练过程中定义调度。类似于迭代剪枝以注入稀疏性,我们从更高的精度(16 位量化或 FP16)开始量化,并逐渐降低量化位数或 FP16 部分的混合精度比例,直到达到目标精度(8 位)。为了控制精度转换,我们定义了一个名为量化周期的超参数,它指示何时应该降低精度。我们观察到,通过使用这种调度,我们可以获得最接近基线的精度。请注意,为了达到特定精度,我们需要定义起始位数和周期,这样在训练样本数量内,模型最终将使用目标位数进行量化。有关更多信息,请参阅量化教程。

为了动态调整量化精度,我们使用特征值作为指标,它显示了训练对精度变化的敏感程度。特征值以前曾用于 Q-BERT 中的量化,以选择网络不同部分的精度位数。为了将此与 MoQ 相结合,我们将特征值根据其绝对值聚类到几个区域中,并相应地调整每个区域的量化周期,特征值的大小越大,因子越大,精度下降越慢。

MoQ (8bit)

图 1. 使用不同层的特征值对 GLUE 任务之一 (QNLI) 进行量化调度。不同的颜色表示 Bert-Base 的 0 到 11 层。

图 1 显示了将特征值与 MoQ 结合使用的结果,针对一个 12 层的 Bert Base 模型。正如我们所见,前几层 (0-4) 比最后几层更容易受到精度降低的影响,因为它们的量化周期比其他层大一个数量级。从该图中可以获得的另一个观察结果是,相邻层以相同的方式降低精度。例如,图 1 左侧图表中的第 9、10 和 11 层,以及右侧图表中的第 0 和 4 层以及第 1 和 3 层,获得了相似的调度。这是由于这些层的特征值在整个训练过程中相似。

MoQ (混合精度) 图 2:使用目标量化周期为 4 位的 QNLI 混合精度量化。

图 2 显示了另一种混合精度量化,它将目标位数设置为 4,但量化周期通过每一层的特征值不断更新。正如我们所见,所有层的最终量化位数都不同。前几层仍然可以量化为 8 位,因为训练样本不足以降低量化位数。另一方面,最后几层继续降低精度。我们最终将整个网络的平均精度降低到 6 位,同时保持模型精度(精度下降 0.3%)。

MoQ (混合精度) 图 3:使用 MoQ 对 Bert SQuAD plus 进行混合精度量化。

作为另一个示例,我们使用基于特征值的 MoQ 对 Bert-Large 进行 SQuAD 微调量化。图 3 显示了我们在每一层微调结束时获得的位数。这里,我们看到与 GLUE 任务上的 BertBase 相比,精度谱略有不同。正如该图所示,我们可以比中间层更积极地降低前几层的精度。此外,最后几层可以容忍非常低的精度,类似于开始几层。这种量化方式最终实现了 90.56 的 F1 分数,与基线非常接近。

量化推理内核

通过使用其他量化方法,在模型量化后,只有当存在对基于整数的操作的硬件支持时,它才能提高性能。因此,所有 GeMM 操作的输入和输出都需要进行量化。然而,由于输入的范围可能会因请求而异,因此在推理时找到每个输入的范围很具有挑战性。另一方面,对所有输入使用静态范围可能会影响推理精度。

为了缓解这个问题,我们引入了推理自定义内核,既不需要硬件支持,也不需要输入量化。这些内核读取量化参数,并在运行时进行反量化,并使用 GPU 核心中的浮点运算单元进行 GeMM 操作。使用这些内核的主要好处是,它们减少了加载模型所需的内存占用,因此我们可以在更少的 GPU 上运行推理,同时通过节省运行 GPU 推理所需的内存带宽来提高性能。

关于量化实现,我们使用不同的算法来根据数据范围和舍入策略对值进行量化。我们支持对称和非对称量化,这两种方案使用得最多。我们对 QAT 应用了这两种技术,并获得了非常相似的结果,但由于对称方法更易于实现,因此我们基于对称方法实现我们的推理内核。关于舍入,我们支持 随机舍入 作为除普通舍入以外的另一种选择。我们已经看到,对于将精度降低到 4 位或更低,随机舍入更有帮助,因为它在训练期间具有无偏的随机行为。

易用性

为了通过 Deepspeed 启用量化,我们只需要通过 JSON 配置文件传递调度即可。为了添加量化的影响,我们在优化器更新参数之前对参数进行量化和反量化。因此,我们不会在模型方面进行任何更改来对模型进行量化。相反,我们通过降低以 FP16 格式保存的数据精度来模拟量化的影响。通过使用这种类型的实现,我们可以使用训练特征(例如步骤数、参数的特征值和原始 FP16 数据格式)灵活地更改精度。正如本博文所示,我们可以通过在整个训练过程中自适应地更改量化调度来提高量化模型的质量。有关如何使用 MoQ 方案的更多信息,请查看我们的 量化教程

提高量化精度。

为了展示我们的量化方案如何保持精度,我们在几个任务和网络上对 MoQ 进行了实验:Bert-Base 上的 GLUE 任务和 Bert-Large 上的 SQuAD。表 1 显示了没有量化 (w/o Quant) 的基线、没有使用任何训练期间调度(Basic Quant)的简单量化以及我们的 MoQ 方案的精度结果。在不使用任何调度的情况下,8 位量化的精度通常低于基线,并且在这个工作负载中,它会导致精度下降 1.02 个点 (ACC)。相比之下,MoQ 使 8 位量化能够获得与 FP16 基线相当的精度,甚至具有更高的 ACC,这证明了我们的量化方法的有效性。

任务 STSB MRPC COLA WNLI SST2 RTE QNLI QQP MNLI SQuAD ACC+
w/o QAT(FP16) 88.71 88.12 56.78 56.34 91.74 65.3 90.96 90.67 84.04 90.56 0
Basic QAT 88.9 88.35 52.78 55.3 91.5 64.2 90.92 90.59 84.01 90.39 -0.87
MoQ 88.93 89 59.33 56.34 92.09 67.15 90.63 90.94 84.55 90.71 0.75

更新: