本文是一篇综述性的博客，探讨总结当下常用的大型 transformer 效率优化方案。
大型 Transformer 模型如今已经成为主流，为各种任务创造了 SOTA 结果。诚然这些模型很强大，但训练和使用起来代价非常昂贵。在时间和内存方面存在有极高的推理成本。概括来说，使用大型 Transformer 模型进行推理的难点，除了模型的规模不断扩大外，还有两个不可忽略的地方：
内存消耗大：推理时，需要把模型参数和中间状态都保存到内存中。例如：KV 存储机制下的缓存中的内容在解码期间需要存储在内存中，举例来说，对于 batch size 为 512，上下文长度为 2048 的设置来说，KV 缓存里需要的空间规模为 3TB，这是模型大小的 3 倍；注意力机制的推理成本和输入序列的长度呈正相关；
低并行性：推理生成过程以自回归的方式执行，使解码过程难以并行。
在这篇文章中，领导 OpenAI 应用研究的 Lilian Weng 写了一篇博客，文中介绍了几种提高 transformer 推理效率的方法。一些是通用的网络压缩方法，而另一些则应用于特定的 transformer 架构。
Lilian Weng 现为 OpenAI 应用人工智能研究负责人，主要从事机器学习、深度学习等研究 。她本科毕业于香港大学，硕士毕业于北京大学信息系统与计算机科学系，之后前往印度安纳大学布鲁顿分校攻读博士。
模型综述
通常将以下内容视为模型推理优化的目标：
使用更少的 GPU 设备和更少的 GPU 内存，减少模型的内存占用；
减少所需的 FLOP，降低计算复杂度；
减少推理延迟，运行得更快。
可以使用几种方法来降低推理过程在内存中的成本，并且加快速度。
在多 GPU 上应用各种并行机制来实现对模型的扩展。模型组件和数据的智能并行使得运行具有万亿级参数的大模型成为可能；
将暂时未使用的数据卸载到 CPU，并在以后需要时读回。这样做对内存比较友好，但会导致更高的延迟；
智能批处理策略；例如 EffectiveTransformer 将连续的序列打包在一起，以删除单个批次中的 padding；
神经网络压缩技术，例如剪枝、量化、蒸馏。就参数数量或位宽而言，小尺寸的模型应该需要少量的内存，也就运行得更快；
特定于目标模型架构的改进。许多架构上的变化，尤其是注意力层的变化，有助于提高 transformer 的解码速度。
本篇文章的重点是网络压缩技术和 transformer 模型在特定体系结构下的改进。
量化策略
在深度神经网络上应用量化策略有两种常见的方法：
训练后量化（PTQ）：首先需要模型训练至收敛，然后将其权重的精度降低。与训练过程相比，量化操作起来往往代价小得多；
量化感知训练 (QAT)：在预训练或进一步微调期间应用量化。QAT 能够获得更好的性能，但需要额外的计算资源，还需要使用具有代表性的训练数据。
值得注意的是，理论上的最优量化策略与实际在硬件内核上的表现存在着客观的差距。由于 GPU 内核对某些类型的矩阵乘法（例如 INT4 x FP16）缺乏支持，并非下面所有的方法都会加速实际的推理过程。
Transformer 量化挑战
许多关于 Transformer 模型量化的研究都有相同的观察结果：训练后将参数简单地量化为低精度（例如 8 位）会导致性能显着下降，这主要是由于普通的激活函数量化策略无法覆盖全部的取值区间。
图 1. 只将模型权重量化为 8 位，激活函数使用完整的精度的时候能取得较好的效果（W8A32）；激活函数量化为 8 位时，无论权重是否为低精度（W8A8 和 W32A8）效果都不如 W8A32。
Bondarenko 等人在一个小型 BERT 模型中观察到，由于输出张量中存在强异常值，FFN 的输入和输出具有非常不同的取值区间。因此，FFN 残差和的逐个张量的量化可能会导致显著的误差。
随着模型参数规模继续增长到数十亿的级别，高量级的离群特征开始在所有 transformer 层中出现，导致简单的低位量化效果不佳。Dettmers 等人观察到大于 6.7B 参数的 OPT 模型就会出现这种现象。模型大了，有极端离群值的网络层也会变多，这些离群值特征对模型的性能有很大的影响。在几个维度上的激活函数异常值的规模就可以比其他大部分数值大 100 倍左右。
图 2. 不同规模的 OPT 模型在四个语言任务（WinoGrande、HellaSwag、PIQA、LAMBADA）上的平均零样本准确率。
混合精度量化
解决上述量化挑战的最直接方法是以不同的精度对权重和激活函数进行量化。
GOBO 模型是首批将训练后量化应用于 transformer 的模型之一（即小型 BERT 模型）。GOBO 假设每一层的模型权重服从高斯分布，因此可以通过跟踪每层的均值和标准差来检测异常值。异常值特征保持原始形式，而其他值被分到多个 bin 中，并且仅存储相应的权重索引和质心值。
基于对 BERT 中只有某些激活层（例如 FFN 之后的残差连接）导致性能大幅下降现象的观察，Bondarenko 等人通过在有问题的激活函数上使用 16 位量化而在其他激活函数上使用 8 位来采用混合精度量化。
LLM.int8 () 中的混合精度量化是通过两个混合精度分解实现的：
因为矩阵乘法包含一组行和列向量之间的独立内积，所以可以对每个内积进行独立量化。每一行和每一列都按最大值进行缩放，然后量化为 INT8；
异常值激活特征（例如比其他维度大 20 倍）仍保留在 FP16 中，但它们只占总权重的极小部分，不过需要经验性地识别离群值。
图 3.LLM.int8（）两种混合精度分解方法。
细粒度量化
图 4. 不同粒度量化对比。d 是模型大小 / 隐空间维度，h 是一个 MHSA（多头自注意力）组件中的头数。
简单地量化一层中的整个权重矩阵（逐个张量或逐个层量化）是最容易实现的，但量化粒度往往不尽如人意。
Q-BERT 将分组量化应用于微调的 BERT 模型，将 MHSA（多头自注意力）中每个头的单个矩阵 W 视为一个组，然后应用基于 Hessian 矩阵的混合精度量化。
Per-embedding group (PEG) 激活函数量化的设计动机是观察到离群值仅出现在少数几个维度中。对每个嵌入层都量化的代价非常昂贵，相比之下，PEG 量化将激活张量沿嵌入维度分成几个大小均匀的组，其中同一组中的元素共享量化参数。为确保所有异常值都分组在一起，PEG 应用了一种基于取值范围的嵌入维度排列算法，其中维度按其取值范围排序。
ZeroQuant 与 Q-BERT 一样都对权重使用分组量化，然后还对激活函数使用了 token-wise 量化策略。为了避免代价昂贵的量化和反量化计算，ZeroQuant 构建了独特的内核来将量化操作与其之前的运算符融合。
使用二阶信息量化
Q-BERT 针对混合精度量化开发了 Hessian AWare 量化 (HAWQ)。其动机是，具有更高 Hessian 谱的参数对量化更敏感，因此需要更高的精度。这种方法本质上是一种识别异常值的方法。
从另一个角度来看，量化问题是一个优化问题。给定一个权重矩阵 W 和一个输入矩阵 X ，想要找到一个量化的权重矩阵 W^ 来最小化如下所示的 MSE 损失：
GPTQ 将权重矩阵 W 视为行向量 w 的集合，并对每一行独立量化。GPTQ 使用贪心策略来选择需要量化的权重，并迭代地进行量化，来最小化量化误差。更新被选定的权重会生成 Hessian 矩阵形式的闭合解。GPTQ 可以将 OPT-175B 中的权重位宽减少到 3 或 4 位，还不会造成太大的性能损失，但它仅适用于模型权重而不适用于激活函数。
异常值平滑
众所周知，Transformer 模型中激活函数比权重更难量化。SmoothQuant 提出了一种智能解决方案，通过数学等效变换将异常值特征从激活函数平滑到权重，然后对权重和激活函数进行量化 (W8A8)。正因为如此，SmoothQuant 具有比混合精度量化更好的硬件效率。
图 5. SmoothQuant 将尺度方差从激活函数迁移到离线权重，以降低激活函数量化的难度。由此产生的新权重和激活矩阵都易于量化。
基于每个通道的平滑因子 s，SmoothQuant 根据以下公式缩放权重：
根据平滑因子
可以很容易地在离线状态下融合到前一层的参数中。超参数 α 控制从激活函数迁移到权重的程度。该研究发现 α=0.5 是实验中许多 LLM 的最佳取值。对于激活异常值较大的模型，可以将 α 调大。
量化感知训练 (QAT)
量化感知训练将量化操作融合到预训练或微调过程中。这种方法会直接学习低位表示的模型权重，并以额外的训练时间和计算为代价获得更好的性能。
最直接的方法是在与预训练数据集相同或代表预训练数据集的训练数据集上量化后微调模型。训练目标可以与预训练目标相同（例如通用语言模型训练中的 NLL/MLM）或特定于的下游任务（例如用于分类的交叉熵）。
另一种方法是将全精度模型视为教师模型，将低精度模型视为学生模型，然后使用蒸馏损失优化低精度模型。蒸馏通常不需要使用原始数据集。
剪枝
网络剪枝是在保留模型容量的情况下，通过修剪不重要的模型权重或连接来减小模型大小。剪枝可能需要也可能不需要重新训练。剪枝可以是非结构化的也可以是结构化的。
非结构化剪枝允许丢弃任何权重或连接，因此它不保留原始网络架构。非结构化剪枝通常对硬件要求比较苛刻，并且不会加速实际的推理过程；
结构化剪枝不改变权重矩阵本身的稀疏程度，可能需要遵循某些模式限制才能使用硬件内核支持的内容。本文专注于那些能实现 transformer 模型的高稀疏性的结构化剪枝。
构建剪枝网络的常规工作流程包含三个步骤：
1. 训练密集型的神经网络直到收敛；2. 修剪网络以去除不需要的结构；3. （可选择）重新训练网络，让新权重保持之前的训练效果。
通过剪枝在密集模型中发现稀疏结构，同时稀疏网络仍然可以保持相似性能的灵感是由彩票假设激发的：这是一个随机初始化的密集前馈网络，它包含一个子网络池。其中只有一个子集（稀疏网络）是中奖彩票（winning tickets），这个中奖彩票在独立训练时可以达到最佳性能。
如何剪枝
Magnitude pruning 是最简单但同时又非常有效的剪枝方法 - 只裁剪那些绝对值最小的权重。事实上，一些研究发现，简单的量级剪枝方法可以获得与复杂剪枝方法相当或更好的结果，例如变分 dropout 和 l_0 正则化。Magnitude pruning 很容易应用于大型模型，并在相当大的超参数范围内实现相当一致的性能。
Zhu & Gupta 发现，大型稀疏模型能够比小型但密集的模型获得更好的性能。他们提出了 Gradual Magnitude Pruning (GMP) 算法，该算法在训练过程中逐渐增加网络的稀疏性。在每个训练步骤中，具有最小绝对值的权重被屏蔽为零以达到所需的稀疏度并且屏蔽的权重在反向传播期间不会得到梯度更新。所需的稀疏度随着训练步骤的增加而增加。GMP 过程对学习率步长策略很敏感，学习率步长应高于密集网络训练中所使用的，但不能太高以防止收敛。
迭代剪枝多次迭代上述三个步骤中的第 2 步（剪枝）和第 3 步（重新训练），每次只有一小部分权重被剪枝，并且在每次迭代中重新训练模型。不断重复该过程，直到达到所需的稀疏度级别。
如何再训练
再训练可以通过使用相同的预训练数据或其他特定于任务的数据集进行简单的微调来实现。
Lottery Ticket Hypothesis 提出了一种权重 rewinding 再训练方法：剪枝后，将未剪枝的权重重新初始化回训练初期的原始值，然后以相同的学习率时间表进行再训练。
学习率 rewinding 仅将学习率重置回其早期值，而保持未剪枝的权重自最后一个训练阶段结束以来不变。研究者观察到 (1) 使用权重 rewinding 的再训练结果优于通过跨网络和数据集进行微调的再训练，以及 (2) 在所有测试场景中学习率 rewinding 与权重 rewinding 的效果持平甚至更优。
稀疏化
稀疏化是扩大模型容量同时保持模型推理计算效率的有效方法。本文考虑两种类型的 transformer 稀疏性：
稀疏化的全连接层，包括自注意力层和 FFN 层；
稀疏模型架构，即 MoE 组件的合并操作。
通过剪枝实现的 N:M 稀疏化
N:M 稀疏化是一种结构化的稀疏化模式，适用于现代 GPU 硬件优化，其中每 M 个连续元素中的 N 个元素为零。例如，英伟达 A100 GPU 的稀疏张量核心支持 2:4 稀疏度以加快推理速度。
图 6. 2:4 结构化稀疏矩阵及其压缩表示。
为了使密集型神经网络的稀疏化遵循 N:M 结构化稀疏模式，英伟达建议使用三步操作来训练剪枝后的网络：训练 –> 剪枝以满足 2:4 稀疏性 –> 重新训练。
(1) 对矩阵中的列进行排列可以在剪枝过程中提供更多可能，以保持参数的数量或满足特殊限制，如 N:M 稀疏性。只要两个矩阵对应的轴按相同的顺序排列，矩阵乘法的结果就不会改变。例如，(1) 在自注意力模块中，如果 query 的嵌入矩阵 Q 的轴 1 和 key 嵌入矩阵 K^⊤的轴 0 采用相同的排列顺序，则 QK^⊤的矩阵乘法最终结果保持不变。
图 7. Q（轴 1）和 K^⊤（轴 0）上相同排列，自注意力模块的结果不变。
(2) 在包含两个 MLP 层和一个 ReLU 非线性层的 FFN 层内，可以将第一个线性权重矩阵 W_1 沿轴 1 排列，然后第二个线性权重矩阵 W_2 沿轴 0 按相同顺序排列。
图 8. W_1（轴 1）和 W_2（轴 0）上有着相同的排列，可以保持 FFN 层的输出不变。为简单起见，图示省略了偏差项，但也应对它们应用相同的排列。
