LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale论文解读

Posted by lili on January 16, 2024

本文是论文LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale的解读。

目录

Abstract

大型语言模型已被广泛采用,但在推理时需要大量GPU内存。我们开发了一种用于transformers中前馈和注意力投影层的Int8矩阵乘法的过程,该过程将推理所需的内存减少了一半,同时保持了完整的精度性能。通过我们的方法,可以加载一个含有175B参数的16/32位检查点,将其转换为Int8,并立即使用,而不会出现性能降低。这得益于对在transformer语言模型中主导注意力和transformer预测性能的高度系统化的新兴特征属性的理解和处理。为了应对这些特征,我们开发了一个两部分的量化过程,称为LLM.int8()。首先,我们使用矢量化的量化,为矩阵乘法中的每个内积使用单独的归一化常数,以量化大多数特征。然而,对于突出的异常值,我们还包括了一种新的混合精度分解方案,该方案将异常值特征维度隔离到16位矩阵乘法中,同时仍然有99.9%以上的值在8位中相乘。使用LLM.int8(),我们通过经验证明,在具有175B参数的LLMs中执行推理是可能的,而不会出现性能降低。这一结果使得这样的模型变得更加易于使用,例如可以在具有消费级GPU的单个服务器上使用OPT-175B/BLOOM。我们开源了我们的软件。

1. Introduction

大型预训练语言模型在自然语言处理领域被广泛采用(Vaswani等,2017;Radford等,2019;Brown等,2020;Zhang等,2022),但在推理时需要大量内存。对于超过6.7B参数的大型transformer语言模型,前馈和注意力投影层以及它们的矩阵乘法操作占用了95%的参数和65-85%的计算(Ilharco等,2020)。减小参数大小的一种方法是将其量化为较少的比特并使用低比特精度的矩阵乘法。为实现这一目标,已经开发了transformers的8位量化方法(Chen等,2020;Lin等,2020;Zafrir等,2019;Shen等,2020)。虽然这些方法减少了内存使用,但通常会降低性能,通常需要在训练后进一步调整量化,并且仅用于具有不到350M参数的模型的研究。350M参数以下的无降级量化方法了解不足,而多十亿参数的量化仍然是一个未解之谜。在本文中,我们提出了第一个transformers的几十亿级别的Int8量化过程,而且不会产生任何性能下降。我们的过程使得可以加载一个具有175B参数的transformer,其权重为16或32位,将前馈和注意力投影层转换为8位,并立即将结果模型用于推理而不会出现性能下降。我们通过解决两个关键挑战来实现这一结果:在超过1B参数的规模上需要更高的量化精度,以及在从6.7B参数的规模开始,有必要显式表示稀疏但系统性的大幅度离群特征,这些特征一旦出现就会破坏所有transformer层的量化精度。这种精度损失在C4评估困惑度(第3节)以及在这些离群特征一旦出现时的零样本准确性中都有体现,如图1所示。

图1:OPT模型在WinoGrande、HellaSwag、PIQA和LAMBADA数据集上的平均零样本准确率。显示的是16位基线,最精确的先前8位量化方法作为基线,以及我们的新8位量化方法LLM.int8()。我们可以看到,一旦在6.7B参数的规模上发生系统性的离群值,常规的量化方法会失败,而LLM.int8()则保持了16位的准确性。

我们展示了我们方法的第一部分,即矢量量化,可以在达到2.7B参数的规模时保持性能。对于矢量量化,矩阵乘法可以看作是行向量和列向量的独立内积序列。因此,我们可以为每个内积使用单独的量化归一化常数来提高量化精度。在执行下一个操作之前,我们可以通过行和列归一化常数的外积进行反归一化,从而恢复矩阵乘法的输出。

要在不降低性能的情况下扩展到6.7B参数以上,关键是要理解在推理过程中隐藏状态的特征维度中极端离群值的出现。为此,我们提供了一个新的描述性分析,显示在大约25%的transformer层中首次出现了与其他维度相比大小高达20倍的大特征,然后随着将transformers扩展到6B参数,逐渐传播到其他层。在约6.7B参数左右,发生了一个相位转变,所有transformer层和75%的序列维度都受到极端大小特征的影响。这些离群值是高度系统化的:在6.7B规模上,每个序列有150,000个离群值,但它们只集中在整个transformer中的6个特征维度中。将这些离群特征维度设置为零会使top-1注意力softmax概率质量减少超过20%,并且尽管它们只占所有输入特征的约0.1%,但会使验证困惑度降低600-1000%。相比之下,删除相同数量的随机特征最多只会使概率减少0.3%,困惑度降低约0.1%。

为了支持具有如此极端离群值的有效量化,我们开发了混合精度分解,这是我们方法的第二部分。我们对离群特征维度执行16位矩阵乘法,对其他99.9%的维度执行8位矩阵乘法。我们将矢量量化和混合精度分解的组合命名为LLM.int8()。我们展示通过使用LLM.int8(),我们可以在具有175B参数的LLMs中执行推理而不降低性能。我们的方法不仅为我们提供了有关这些离群值对模型性能的影响的新见解,而且还首次使得能够在单个具有消费级GPU的服务器上使用非常大的模型,例如OPT-175B/BLOOM。尽管我们的工作侧重于使大型语言模型在不降级的情况下易于使用,但我们还在附录D中展示,我们保持了大型模型(例如BLOOM-176B)的端到端推理运行时性能,并为GPT-3模型(参数规模为6.7B或更大)提供了适度的矩阵乘法加速。我们开源了我们的软件,并发布了Hugging Face Transformers(Wolf等,2019)的集成,使我们的方法适用于所有具有线性层的托管Hugging Face Models。

图2:LLM.int8()的示意图。给定16位浮点输入$X_{f16}$和权重$W_{f16}$,特征和权重被分解为具有大幅度特征和其他值的子矩阵。离群特征矩阵在16位中相乘。所有其他值在8位中相乘。我们通过将$C_x$和$C_w$的行和列绝对最大值进行缩放,然后将输出量化为Int8来执行8位矢量乘法。Int32矩阵乘法输出$Out_{i32}$通过归一化常数$C_x \otimes C_w$的外积进行反量化。最后,离群和常规输出都在16位浮点输出中累积。

2 背景

在这项工作中,我们通过扩展transformer模型来推动量化技术的极限。我们关注两个问题:在什么规模和为什么量化技术会失败,以及这与量化精度有什么关系?为了回答这些问题,我们研究了高精度的非对称量化(零点量化)和对称量化(绝对最大量化)。虽然零点量化通过使用数据类型的完整位范围提供了高精度,但由于实际约束,它很少被使用。绝对最大量化是最常用的技术。

2.1 8位数据类型和量化

Absmax量化通过将输入缩放到8位范围[-127,127],方法是与$s_{xf16}$相乘,其中$s_{xf16}$是整个张量的绝对最大值除以127。这相当于除以无穷大范数并乘以127。因此,对于FP16输入矩阵$X_{f16} \in R ^{s \times h}$,Int8的absmax量化为:

其中,$ \lfloor \rceil $表示四舍五入到最近的整数。

零点量化通过使用归一化的动态范围$nd_x$进行缩放,然后通过零点$zp_x$进行移位,将输入分布转移到完整范围[-127,127]。通过这个仿射变换,任何输入张量都将使用数据类型的所有位,从而减小非对称分布的量化误差。例如,对于ReLU输出,在absmax量化中,[−127, 0) 中的所有值都未使用,而在零点量化中使用了完整的[-127,127] 范围。零点量化由以下方程表示:

【译注:感觉零点的计算有问题,应是$nd \times minX$】

为了在操作中使用零点量化,我们将张量$X_{i8}$和零点$zp_{x_{i16}}$同时输入到一个特殊的指令中,该指令在执行16位整数操作之前将$zp_{x_{i16}}$添加到$X_{i8}$的每个元素。例如,要根据两个零点量化的数字$A_{i8}$和$B_{i8}$以及它们的零点$zp_{a_{i16}}$和$zp_{b_{i16}}$做乘法,计算方法为:

如果设备(比如GPU或者TPU)没有类似的指令,则可以展开:

其中,$A_{i8}B_{i8}$是以Int8精度计算的,而其余部分则以Int16/32精度计算。因此,如果$multiply_{i16}$指令不可用,零点量化可能会较慢。在两种情况下,输出都作为32位整数$C_{i32}$累积。为了反量化$C_{i32}$,我们除以缩放常数$nd_{af16}$和$nd_{bf16}$。

16位浮点输入和输出的Int8矩阵乘法。给定隐藏状态$X_{f16} \in R^{s \times h}$ 和权重$W_{f16} \in R^{h \times o}$,其中序列维度为s,特征维度为h,输出维度为o,我们进行8位矩阵乘法,输入和输出为16位,具体如下:

其中,Q(·)是absmax或zeropoint量化,$c_{xf16}$和$c_{wf16}$是absmax的$s_x$和$s_w$或zeropoint量化的$nd_x$和$nd_w$的相应张量缩放常数。

上面的过程可以用下图来解释:

比如我们要对16位的输入X和参数W作矩阵乘法,我们对输入X按行做ABSMAX量化,因此这三行最大的是$C_X=[2,3,1]$。同样对W的列做ABSMAX量化,最大的值是$C_W=[1,2]$,然后用这些缩放因子对X和W进行int8量化得到$X_{I8},W_{I8}$。然后用上面的公式计算int8矩阵的乘法得到结果$Out_{I32}$,最后把它除以对应的$C_X,C_W$就得到16位的矩阵乘法结果。

3. 大规模的Int8矩阵乘法

使用每个张量单一缩放常数的量化方法的主要挑战是,一个离群值可能降低所有其他值的量化精度。因此,希望每个张量有多个缩放常数,如块内常数(Dettmers等,2022),以便将离群值的影响限制在每个块内。我们改进了量化的最常见的一种块方法,即行内量化(Khudia等,2021),通过使用矢量内量化,如下文更详细地描述。

为了处理在6.7B规模以上的所有Transformer层中出现的大量离群值特征,矢量内量化已经不再足够。为此,我们开发了混合精度分解,其中大量离群值特征维度的小部分(约0.1%)以16位精度表示,而其他99.9%的值以8位表示。由于大多数条目仍然以低精度表示,与16位相比,我们保留了约50%的内存减少。例如,对于BLOOM-176B,我们将模型的内存占用减少了1.96倍。矢量内量化和混合精度分解如图2所示。LLM.int8()方法是absmax矢量内量化和混合精度分解的组合。

3.1 矢量内量化

增加矩阵乘法的缩放常数数量的一种方式是将矩阵乘法视为一系列独立的内积。给定隐藏状态$X_{f16} \in R^{b\times h}$ 和权重矩阵$W_{f16} \in R^{h \times o}$,我们可以为$X_{f16}$的每一行分配一个不同的缩放常数$c_{xf16}$,并为$W_{f16}$的每一列分配$c_w$。为了反量化,我们通过$1/(c_{xf16} c_{wf16})$对每个内积结果进行反标准化。对于整个矩阵乘法,这相当于通过外积$c_{xf16} \otimes c_{wf16}$进行反标准化,其中$c_x \in R^s, c_w \in R^o$。因此,具有行和列常数的矩阵乘法的完整方程为:

我们把这个叫做基于行的矩阵乘法。 【译注:前面以及结合图片解释过了。】

3.2 LLM.int8()的核心:混合精度分解

在我们的分析中,我们证明了十亿级8位Transformer的一个显著问题是它们具有重要的大量特征(列),这对Transformer性能至关重要并需要高精度量化。然而,矢量内量化,我们最佳的量化技术,对隐藏状态的每一行进行量化,对于离群特征是无效的。幸运的是,我们看到这些离群特征在实践中既非常稀疏又有系统性,仅占所有特征维度的约0.1%,因此我们可以开发一种专注于这些特定维度的高精度乘法的新分解技术。

我们发现,对于给定的输入矩阵$X_{f16} \in R^{s \times h}$,这些离群值在几乎所有序列维度s上都是系统的,但仅限于特定的特征/隐藏维度h。因此,我们提出了矩阵乘法的混合精度分解,其中我们将离群特征维度分离到集合O = {i|i ∈ Z,0 ≤ i ≤ h}中,该集合包含所有至少有一个离群值的维度h,其幅度大于阈值α。在我们的工作中,我们发现α = 6.0足以将Transformer性能降低接近零。使用爱因斯坦符号表示法,其中所有指数都是上标,给定权重矩阵$W_{f16} \in R^{h \times o}$,矩阵乘法的混合精度分解定义如下:

其中,Sf16是用于Int8输入和权重矩阵Xi8和Wi8的去归一化项。这种8位和16位的分离允许对离群值进行高精度乘法,同时使用8位权重进行高效的内存矩阵乘法,其中超过99.9%的值为8位。由于直到13B参数的Transformer中异常特征维度的数量不大于7(|O| ≤ 7),因此此分解操作仅额外消耗约0.1%的内存。

【译注:这个过程如下图所示。通过离群(Outlier)检测,把输入X中包含异常值的列挑出来直接做16位的矩阵乘法,然后剩下的用前面解释的方法用int8乘法,最后把它们累加起来。】

3.3 实验设置

我们通过将几个公开可用的预训练语言模型的规模扩大到175B参数来衡量量化方法的鲁棒性。关键问题不是量化方法在特定模型上的表现如何,而是在不断扩展规模时该方法的性能趋势。我们使用两个实验设置。一个基于语言建模的困惑度,我们发现这是一个非常稳健的度量标准,对于量化降级非常敏感。我们使用这个设置来比较不同的量化基线。此外,我们评估了OPT模型在各种不同端任务上的零样本准确性降级,其中我们将我们的方法与16位基线进行比较。 对于语言建模设置,我们使用在fairseq(Ott等人,2019)中预训练的密集自回归Transformer,参数范围从125M到13B。这些Transformer已在Books(Zhu等人,2015),英语维基百科,CC-News(Nagel,2016),OpenWebText(Gokaslan和Cohen,2019),CC-Stories(Trinh和Le,2018)和英语CC100(Wenzek等人,2020)上进行了预训练。有关这些预训练模型的详细信息,请参阅Artetxe等人(2021)。 为了评估Int8量化后的语言建模降级,我们评估了C4语料库(Raffel等人,2019)验证数据上8位Transformer的困惑度,该语料库是Common Crawl语料库的子集。5我们在此评估中使用NVIDIA A40 GPU。 为了衡量零样本性能的降级,我们使用OPT模型(Zhang等人,2022),并在EleutherAI语言模型评估平台(Gao等人,2021)上评估这些模型。

3.4 主要结果

在C4语料库上评估的125M至13B Int8模型的主要语言建模困惑度结果如表1所示。我们看到absmax、row-wise和zeropoint量化在规模扩大时失败,超过2.7B参数的模型表现比较小的模型更差。相比之下,zeropoint量化在6.7B参数以上失效。我们的方法LLM.int8()是唯一保持困惑度的方法,因此LLM.int8()是唯一具有良好扩展趋势的方法。

表1:不同Transformer规模(从125M到13B参数)的C4验证困惑度的量化方法。我们看到absmax、row-wise、zeropoint和vector-wise量化随着规模扩大导致显著的性能降级,特别是在13B标记处,8位13B困惑度比8位6.7B困惑度更糟糕。如果使用LLM.int8(),我们在扩展规模时恢复了完整的困惑度。由于其非对称性,zeropoint量化显示出一定的优势。

当我们观察OPT模型在EleutherAI语言模型评估平台上的零样本性能的扩展趋势时,可以看到LLM.int8()在从125M扩展到175B参数时保持完整的16位性能。另一方面,基线8位absmax vector-wise量化在扩展时表现不佳,退化为随机性能。

虽然我们的主要重点是节省内存,但我们还测量了LLM.int8()的运行时间。与FP16基线相比,对于规模小于6.7B参数的模型,量化开销可能会减缓推理速度。然而,6.7B参数或更小的模型适合大多数GPU,并且在实践中对量化的需求较小。对于大矩阵乘法,等效于175B模型中的乘法,LLM.int8()的运行时间约快两倍。附录D提供了这些实验的更多细节。

4. 大规模Transformer中的突出特征

随着我们扩展transformer的规模,具有较大幅度的离群特征出现并强烈影响所有层及其量化。给定隐藏状态$X \in R^{s \times h}$,其中s是序列/token维度,h是隐藏/特征维度,我们将一个特征定义为一个特定的维度 hi。我们的分析观察了给定transformer的所有层中特定特征维度 hi。 我们发现,离群特征强烈影响transformer的注意力和整体预测性能。尽管在13B模型的2048个token序列中可能存在多达150,000个离群值,但这些离群特征是高度系统化的,仅表示至多7个唯一的特征维度$h_i$。该分析的见解对于开发混合精度分解至关重要。我们的分析解释了零点量化的优势以及在使用混合精度分解和小模型与大模型的量化性能中它们为何消失。

4.1 发现离群特征

定量分析新兴现象的困难有两方面。我们旨在选择用于分析的一小部分特征,以使结果易于理解而不过于复杂,同时还要捕捉重要的概率和结构化模式。我们使用经验方法找到这些约束条件。我们根据以下标准定义离群值:特征的幅度至少为6.0,影响至少25%的层,并影响至少6%的序列维度。更正式地,给定一个具有L层和隐藏状态$X_l \in R^{s \times h}$的transformer,其中l = 0…L,其中s是序列维度,h是特征维度,我们定义一个特征为任何隐藏状态$X_{l_i}$中的特定维度 $h_i$。我们跟踪具有至少一个幅度 ≥ 6 的值的维度 $h_i$,0 ≤ i ≤ h,并且仅在这些离群值在0…L所有层次中的至少25%的transformer层中的相同特征维度 $h_i$中出现,并且在所有隐藏状态$X_l$中的至少6%的所有序列维度 s 中出现时,我们才收集统计信息。由于特征离群值仅发生在注意力投影(键/查询/值/输出)和前馈网络扩展层(第一子层),我们在此分析中忽略了注意力函数和FFN收缩层(第二子层)。

图3:受大幅度离群特征影响的transformer各层百分比和所有序列维度,按(a) 模型大小或(b) C4困惑度表示。线条是(a)和(b)的4个和9个线性段的B样条插值。一旦发生阶段性转变,离群值将存在于所有层和大约75%的所有序列维度中。虽然(a)表明参数大小发生突然的阶段性转变,(b)表明随着困惑度降低,存在逐渐的指数性阶段性转变。在(a)中的明显转变与量化方法性能突然下降同时发生。

我们设定这些阈值的原因如下。我们发现,使用混合精度分解时,当将任何幅度大于等于6的特征视为离群特征时,困扰的困惑减少。对于受到离群值影响的层数,我们发现在大型模型中离群特征是系统化的:它们要么出现在大多数层中,要么根本不出现。另一方面,它们在小型模型中是概率性的:它们在每个序列的某些层次中有时出现。因此,我们设置受到离群特征影响的层数的阈值,以便在我们最小的125M参数模型中检测到单个离群特征。该阈值对应于在相同特征维度中至少25%的transformer层受到离群值影响。第二常见的离群特征仅出现在单个层次(2%的层次),表明这是一个合理的阈值。我们使用相同的程序找到在我们的125M模型中受离群特征影响的序列维度的阈值:离群特征至少出现在6%的序列维度中。

我们测试了最大达到13B参数的模型。为确保观察到的现象不是由软件中的错误引起的,我们评估了在三种不同软件框架中训练的transformers。我们评估了使用OpenAI软件的四个GPT-2模型,使用Fairseq(Ott et al., 2019)的五个Meta AI模型,以及使用Tensorflow-Mesh(Shazeer et al., 2018)的EleutherAI模型GPT-J。有关更多详细信息,请参阅附录C。我们还在两个不同的推理软件框架中执行了我们的分析:Fairseq和Hugging Face Transformers(Wolf et al., 2019)。

4.2 评估离群特征的影响

为了证明离群特征对注意力和预测性能至关重要,我们在将隐藏状态Xl传递到注意力投影层之前将离群特征设为零,然后将带有离群特征的常规softmax概率与没有离群特征的常规softmax概率进行比较。我们对所有层独立进行此操作,意味着我们将常规softmax概率值向前传递,以避免级联错误并隔离由于离群特征引起的效应。我们还报告了如果去除离群特征维度(将其设为零)并通过transformer传递这些修改后的隐藏状态时的困惑度降低。作为对照,我们对随机的非离群特征维度应用相同的过程,并注意到注意力和困惑度的降低。

图4:(a)中最大离群特征的幅度中位数表明离群大小发生了突然的变化。这似乎是突然出现后量化方法失败的主要原因。虽然离群特征维度的数量大致与模型大小成比例,(b)显示在所有分析的模型中,离群值的数量与困惑度严格单调递增。线条是9个线性段的B样条插值。

我们的主要定量结果可以总结为四个主要观点: (1) 通过参数数量衡量,跨transformer所有层出现大幅度特征的突现在6B和6.7B参数之间突然发生,如图3a所示,受影响层的百分比从65%增加到100%。受影响的序列维度数量迅速增加,从35%增加到75%。这种突然的变化与量化开始失败的点同时发生。 (2) 或者,通过困惑度衡量,可以看出transformer所有层中大幅度特征的突然出现可以看作是根据递减困惑度的指数函数平滑出现,如图3b所示。这表明出现并非突然发生,我们可能能够通过研究较小模型中指数趋势来在阶段性转变发生之前检测出突发特征。这也表明出现不仅仅是关于模型大小,还与困惑度有关,而困惑度与许多其他因素(例如所使用的训练数据量和数据质量)相关(Hoffmann等人,2022; Henighan等人,2020)。 (3) 离群特征的中位数离群值迅速增加,一旦离群特征在transformer的所有层中出现,如图4a所示。离群特征的大幅度和其非对称分布破坏了Int8量化的精度。这是量化方法在6.7B规模开始失败的核心原因 - 量化分布的范围太大,以至于大多数量化区间为空,小的量化值被量化为零,从根本上熄灭了信息。我们假设除了Int8推理外,对于6.7B规模以上的模型,正常的16位浮点训练也变得不稳定 - 如果乘以充满数量级为60的值的向量,就很容易超过最大的16位值65535。 (4) 离群特征的数量随着C4困惑度的降低而严格单调增加,如图4b所示,而与模型大小的关系是非单调的。这表明模型困惑度而不仅仅是模型大小决定了阶段性转变。我们假设模型大小只是许多要求出现的协变量中的一个重要因素。这些离群特征在阶段性转变之后高度系统化。例如,对于一个序列长度为2048的6.7B transformer,我们发现每个序列大约有150k个离群特征,但这些特征集中在仅有6个不同的隐藏维度中。 这些离群特征对transformer的性能至关重要。如果移除这些离群值,平均top-1 softmax概率将从约40%降低到约20%,验证困惑度增加600-1000%,尽管最多只有7个离群特征维度。当我们移除7个随机特征维度时,top-1概率只会降低0.02-0.3%,困惑度增加0.1%。这突显了这些特征维度的关键性质。对于这些离群特征的量化精度至关重要,因为即使是微小的错误也会极大地影响模型性能。

4.3 量化性能的解释

我们的分析表明,在大型Transformer中,特定特征维度中的离群值是普遍存在的,这些特征维度对于Transformer性能至关重要。由于逐行和逐向量量化分别缩放每个隐藏状态的序列维度 s(行),并且由于离群值出现在特征维度 h(列),因此这两种方法都不能有效处理这些离群值。这就是为什么绝对最大值量化方法在出现后迅速失败的原因。

然而,几乎所有的离群值都具有严格的不对称分布:它们要么完全是正的,要么完全是负的(请参见附录 C)。这使得零点量化对于这些离群值特别有效,因为零点量化是一种不对称的量化方法,可以将这些离群值缩放到完整的 [−127, 127] 范围内。这解释了我们在表1中的量化缩放基准测试中的强大性能。然而,在13B规模上,即使是零点量化也会由于积累的量化误差和离群值幅度的快速增长而失败,正如图4a所示。

如果我们使用我们完整的 LLM.int8() 方法进行混合精度分解,零点量化的优势就会消失,表明剩余的分解特征是对称的。然而,相对于逐行量化,逐向量量化仍然具有优势,表明模型权重的增强量化精度是需要的,以保持完整的精度预测性能。