【翻译】大型 Transformer 模型推理优化

翻译原文：Large Transformer Model Inference Optimization | Lil'Log  原文作者：Lilian Weng

目录

• 方法概述
• 蒸馏 Distillation
• 量化 Quantization
  • Transformer 量化的挑战
  • 训练后量化 (PTQ)
    • 混合精度量化 Mixed-precision quantization
    • 细粒度量化
    • 量化的二阶信息
    • 异常值平滑
  • 量化感知训练 (QAT)
• 修剪 Pruning
  • 如何修剪？
  • 如何重新训练？
• 稀疏性 Sparsity
  • 通过修剪实现 N:M 稀疏性
  • 稀疏化 Transformer
  • 混合专家 MoE Mixture-of-Experts
    • 路由策略改进
    • 内核改进
• 架构优化
  • 稀疏注意模式 Sparse Attention
  • 循环 Recurrence
  • 内存节省设计
  • 自适应注意 Adaptive Attention
• 引用
• 参考文献

[Updated on 2023-01-24: add a small section on Distillation.]

大型Transformer模型如今已成为主流，可为各种任务创建SoTA结果。它们功能强大，但训练和使用成本非常高。极高的推理成本（无论是时间成本还是内存成本）是采用强大的 Transformer 大规模解决现实世界任务的一大瓶颈。

为什么大型 Transformer 模型很难运行推理？除了 SoTA 模型的规模不断扩大外，还有两个主要因素导致了推理挑战（Pope 等人，2022 年）：

1. 内存占用大。推理时，模型参数和中间状态都需要存储在内存中。例如，

• KV cache应在解码时存储在内存中；例如，对于批处理大小为 512 且上下文长度为 2048 的情况，KV cache总计 3TB，即模型大小的 3 倍(!). 
• 注意力机制的推理成本与输入序列长度成二次方关系。

2. 可并行性低。推理生成以自回归方式执行，使解码过程难以并行化。

在这篇文章中，我们将研究几种使 Transformer 推理更高效的方法。一些是通用的网络压缩方法，而另一些则特定于 Transformer 架构。

方法概述

我们通常将以下内容视为模型推理优化的目标：

• 通过使用更少的 GPU 设备和更少的 GPU 内存来减少模型的内存占用；
• 通过降低所需的 FLOPs 数量来降低所需的计算复杂度；
• 减少推理延迟并加快运行速度。

可以使用多种方法使推理在内存中更便宜或/和在时间上更快。

1. 应用各种并行性以在大量 GPU 上扩展模型。模型组件和数据的智能并行性使运行数万亿个参数的模型成为可能。
2. 内存卸载将暂时未使用的数据卸载到 CPU 并在以后需要时读回它们。这有助于减少内存使用，但会导致更高的延迟。
3. 智能批处理策略；例如，EffectiveTransformer 将连续序列打包在一起以在一个批次中去除填充padding.
4. 网络压缩技术，例如修剪、量化、蒸馏。就参数数量或位宽而言，较小尺寸的模型需要的内存更少，运行速度更快。
5. 针对目标模型架构的改进。许多架构变化（尤其是针对注意力层的架构变化）有助于提高 Transformer 解码速度。

查看上一篇关于大型模型训练的文章，了解不同类型的训练并行性和节省内存的设计，包括 CPU 内存卸载。这篇文章重点介绍网络压缩技术和针对 Transformer 模型的架构特定改进。

蒸馏

知识蒸馏 (KD；Hinton 等人 2015 年，Gou 等人 2020 年) 是一种构建更小、更便宜的模型（“学生模型”）的直接方法，通过将技能从预先训练的昂贵模型（“教师模型”）转移到学生身上来加快推理速度。对于如何构建学生架构没有太多限制，除了与教师匹配的输出空间以构建适当的学习目标。

给定一个数据集，通过蒸馏损失训练学生模型来模仿老师的输出。通常神经网络有一个softmax层；例如，LLM输出token的概率分布。我们将softmax之前的logits层分别表示为Zt和Zs，分别用于教师和学生模型。蒸馏损失用较高的温度T最小化两个softmax输出之间的差异。当已知基本事实ground truth标签y时，我们可以将其与基本事实和学生的soft logits之间的监督学习目标相结合，例如使用交叉熵。

其中是一个超参数，用于平衡软学习目标和硬学习目标。一个常见选择是 KL 散度/交叉熵。

一个成功的早期试验是 DistilBERT（Sanh 等人，2019 年），它能够将 BERT 的参数减少 40%，同时在微调的下游任务上保持 BERT 97% 的性能，并且运行速度提高 71%。预训练 DistilBERT 的损失是软蒸馏损失、监督训练损失（即 BERT 情况下的掩蔽语言建模损失）和特殊余弦嵌入损失的组合，用于对齐教师和学生之间的隐藏状态向量。

蒸馏可以轻松地与量化、修剪或稀疏化技术相结合，其中教师模型是原始的全精度密集模型，而学生被量化、修剪pruned或修剪trimmed以具有更高的稀疏度。

量化

在深度神经网络上应用量化有两种常见方法：

1. 训练后量化 (PTQ)：首先训练模型以使其收敛，然后将其权重转换为较低精度而无需更多训练。与训练相比，它通常实施起来相当便宜。
2. 量化感知训练 (QAT)：量化应用于预训练或进一步微调。QAT 能够获得更好的性能，但需要额外的计算资源和对代表性训练数据的访问。

我们应该意识到理论最佳量化策略与硬件内核支持之间的差距。由于 GPU 内核缺乏对某些类型的矩阵乘法（例如 INT4 x FP16）的支持，因此并非以下所有方法都能加速实际推理。

Transformer 量化的挑战

许多关于 Transformer 模型量化的研究都有相同的观察结果：简单的低精度（例如 8 位）训练后量化会导致性能显着下降，这主要是由于激活的动态范围高，而简单的激活量化策略无法保持容量。

Bondarenko 等人 (2021) 在小型 BERT 模型中观察到，由于输出张量tensor中存在强异常值，FFN 的输入和输出具有非常不同的动态范围。因此，对 FFN 残差和进行每个张量量化可能会导致显着误差。

随着模型大小继续增长到数十亿个参数，所有 Transformer 层中开始出现高强度的异常值特征，导致简单的低位量化失败。Dettmers 等人 (2022) 在参数大于 6.7B 的 OPT 模型中观察到了这种现象。较大的模型具有更多具有极端异常值的层，这些异常值特征对模型性能有显著影响。少数维度上的激活异常值的规模可能比大多数其他值大约 100 倍。

训练后量化 (PTQ)

混合精度量化

解决上述量化挑战的最直接方法是针对权重和激活以不同的精度实施量化。

GOBO (Zadeh 等人，2020) 是首批在 Transformer（即小型 BERT 模型）上应用训练后量化的模型之一。它假设每层的模型权重遵循高斯分布，因此通过跟踪每层的平均值和标准差来检测异常值。异常值特征保持原始形式，而其他值则被分成多个 bin，并且仅存储权重和质心值的相应 bin 索引。

根据 BERT 中只有某些激活层（例如 FFN 后的残差连接）会导致性能大幅下降的观察结果，Bondarenko 等人 (2021) 采用混合精度量化，对有问题的激活使用 16 位量化，而对其他激活使用 8 位量化。

LLM.int8() (Dettmers 等人，2022) 中的混合精度量化是通过两个混合精度分解实现的：

1. 由于矩阵乘法包含行向量和列向量之间的一组独立内积，我们可以对每个内积施加独立量化：每行和每列都按绝对最大值缩放，然后量化为 INT8。
2. 异常值激活特征（例如比其他维度大 20 倍）保留在 FP16 中，但它们仅占总权重的一小部分。如何识别异常值是经验性的。

单层量化整个权重矩阵（“逐张量”或“逐层”量化）最容易实现，但不会导致量化的良好粒度。

Q-BERT（Shen, Dong & Ye, et al. 2020）将分组量化应用于微调的 BERT 模型，将 MHSA（多头自注意力）中每个头的单个矩阵W视为一个组，然后应用基于 Hessian 的混合精度量化。

逐嵌入组 (PEG) 激活量化的动机是观察到异常值仅出现在少数d（隐藏状态/模型大小）维度中（Bondarenko et al. 2021）。逐嵌入的计算成本相当高。相比之下，PEG 量化将激活张量沿嵌入维度分成几个大小均匀的组，其中同一组中的元素共享量化参数。为了确保所有异常值都分组在一起，它们应用了基于范围的确定性嵌入维度排列，其中维度按其值范围排序。

ZeroQuant (Yao et al. 2022) 使用分组量化作为权重，与 Q-BERT 相同，并使用标记量化作为激活。为了避免昂贵的量化和反量化计算，ZeroQuant 构建了定制内核，将量化操作与其先前的运算符融合在一起。

量化的二阶信息

Q-BERT (Shen, Dong & Ye, et al. 2020) 为其混合精度量化开发了 Hessian AWare Quantization (HAWQ)。其动机是具有较高 Hessian 谱（即较大的顶部特征值）的参数对量化更敏感，因此需要更高的精度。它本质上是一种识别异常值的方法。

从另一个角度来看，量化问题是一个优化问题。给定一个权重矩阵W和一个输入矩阵X，我们希望找到一个量化权重矩阵来最小化 MSE：

GPTQ（Frantar 等人，2022 年）将权重矩阵W视为行向量的集合w并对每行单独应用量化。GPTQ 迭代量化更多贪婪选择的权重，以最小化量化误差。所选权重的更新具有闭式公式，利用 Hessian 矩阵。如果感兴趣，请阅读论文和 OBQ（最佳脑量化；Frantar 和 Alistarh，2022 年）方法中的更多详细信息。GPTQ 可以将 OPT-175B 中权重的位宽减少到 3 或 4 位，而不会造成太多性能损失，但它仅适用于模型权重，而不适用于激活。

离群值平滑

众所周知，在 Transformer 模型中，激活比权重更难量化。SmoothQuant (Xiao & Lin 2022) 提出了一种智能解决方案，通过数学等效变换将离群值特征从激活平滑到权重，然后对权重和激活进行量化 (W8A8)。因此，SmoothQuant 比混合精度量化具有更好的硬件效率。

平滑因子可以很容易地离线融合到先前层的参数中。超参数控制我们将量化难度从激活迁移到权重的程度：。本文发现
是实验中许多 LLM 的最佳点。对于激活中具有更显著异常值的模型，
可以调整为更大。

量化感知训练 (QAT)

量化感知训练将量化操作融合到预训练或微调过程中。它直接学习低位表示中的模型权重，并以额外的训练时间和计算为代价获得更好的性能。

最直接的方法是在与预训练数据集相同或具有代表性的训练数据集上对量化后的模型进行微调。训练目标可以与预训练的目标相同（例如，一般语言模型训练中的 NLL/MLM），也可以特定于我们关心的下游任务（例如，分类的交叉熵）。

另一种方法是将全精度模型视为老师，将低精度模型视为学生，然后使用蒸馏损失优化低精度模型。蒸馏通常不需要使用原始数据集；例如，维基百科数据集是一个不错的选择，甚至随机标记也可以带来不错的性能提升。逐层知识蒸馏 (LKD；Yao 等人，2022) 方法逐层量化网络，并使用其原始的未量化版本作为教师模型。给定相同的输入，LKD 最小化与层权重相乘和量化层权重相乘之间的 MSE。

修剪Pruning

网络修剪是通过修剪不重要的模型权重或连接来减小模型大小，同时保留模型容量。它可能需要或不需要重新训练。修剪可以是非结构化的，也可以是结构化的。

• 非结构化修剪可以删除任何权重或连接，因此它不保留原始网络架构。非结构化修剪通常不适用于现代硬件，并且不会导致实际的推理加速。
• 结构化修剪旨在保持密集矩阵乘法形式，其中一些元素为零。它们可能需要遵循某些模式限制才能与硬件内核支持的内容一起工作。在这里，我们专注于结构化修剪，以实现 Transformer 模型的高稀疏性。

构建修剪网络的常规工作流程有三个步骤：

1. 训练密集网络直至收敛；
2. 修剪网络以删除不需要的结构；
3. 可选择重新训练网络以使用新权重恢复性能。

通过网络修剪在密集模型中发现稀疏结构，同时稀疏网络仍能保持相似性能的想法是由彩票假设 (LTH) 激发的：随机初始化的密集前馈网络包含一个子网络池，其中只有一个子集（稀疏网络）是“中奖彩票”，在单独训练时可以实现最佳性能。

如何修剪？

幅度修剪是最简单但非常有效的修剪方法 - 修剪绝对值最小的权重。事实上，一些研究（Gale 等人，2019 年）发现，简单的幅度修剪方法可以实现与复杂修剪方法相当或更好的结果，例如变分丢弃（Molchanov 等人，2017 年）和正则化（Louizos 等人，2017 年）。幅度修剪易于应用于大型模型，并在各种超参数中实现相当一致的性能。

Zhu 和 Gupta（2017 年）发现，大型稀疏模型能够比小型但密集的模型实现更好的性能。他们提出了渐进幅度修剪 (GMP) 算法，该算法在训练过程中逐渐增加网络的稀疏性。在每个训练步骤中，绝对值最小的权重被掩蔽为零，以达到所需的稀疏度，并且掩蔽的权重在反向传播期间不会获得梯度更新。所需的稀疏度会随着训练步骤的增加而增加。GMP 的过程对学习率计划很敏感，该计划应高于密集网络训练中使用的计划，但不能太高，以免阻止收敛。

迭代剪枝（Renda 等人，2020 年）多次迭代步骤 2（剪枝）和步骤 3（重新训练）：每次迭代中仅剪枝一小部分权重并重新训练模型。该过程重复进行，直到达到所需的稀疏度。

如何重新训练？

重新训练步骤可以是使用相同的预训练数据或其他特定于任务的数据集进行简单的微调。

彩票假说提出了一种权重倒退再训练技术：修剪后，未修剪的权重在训练早期重新初始化为原始值，然后使用相同的学习率计划进行再训练。

学习率倒退（Renda 等人，2020 年）仅将学习率重置回其早期值，而未修剪的权重自上一个训练阶段结束以来保持不变。他们观察到 (1) 使用权重倒退的再训练优于在网络和数据集上使用微调的再训练，(2) 在所有测试场景中，学习率倒退都与权重倒退相当或优于权重倒退。

稀疏性

稀疏性是一种有效的方法，可以扩大模型容量，同时保持模型推理的计算效率。在这里，我们考虑了 Transformer 的两种稀疏性：

• 稀疏化的密集层，包括自注意力层和 FFN 层。
• 稀疏模型架构；即通过合并混合专家 (MoE) 组件。

通过修剪实现 N:M 稀疏性

N:M 稀疏性是一种结构化的稀疏性模式，可以很好地与现代 GPU 硬件优化配合使用，其中每个M连续元素N中都有零。例如，Nvidia A100 GPU 的稀疏张量核心支持 2:4 稀疏性，以实现更快的推理 (Nvidia 2020)。

为了将密集神经网络稀疏化以遵循 N:M 结构化稀疏模式，Nvidia (2020) 建议使用三步常规工作流程来训练修剪网络：训练 –> 修剪以满足 2:4 稀疏性 –> 重新训练。

排列列可以在修剪过程中提供更多选项，以保持较大幅度的参数或满足 N:M 稀疏性等特殊限制 (Pool & Yu 2021)。只要两个矩阵的成对轴以相同的顺序排列，矩阵乘法的结果就不会改变。例如，

（1）在自注意力模块中，如果对查询嵌入矩阵  的轴 1 和键嵌入矩阵   的轴 0 应用相同的排列顺序，则 的矩阵乘法的最终结果将保持不变。

（2）在包含两个 MLP 层和一个 ReLU 非线性层的 FFN 层中，我们可以以相同的顺序沿轴 1 排列第一个线性权重矩阵 ，并沿轴 0 排列第二个线性权重矩阵 。

为了强制 N:M 结构化稀疏性，我们将一个矩阵的列拆分为多个列幻灯片（称为“条纹”），我们可以很容易地观察到，每个条纹内的列顺序和条纹顺序对 N:M 稀疏性限制都没有影响。

Pool & Yu (2021) 提出了一种迭代贪婪算法来寻找最大化 N:M 稀疏性权重幅度的最佳排列。所有通道对都经过推测交换，并且只采用导致幅度最大增加的交换，生成新的排列并结束一次迭代。贪婪算法可能只能找到局部最小值，因此他们引入了两种逃避局部最小值的技术：

1. 有界回归：在实践中，两个随机通道被交换，最多交换固定次数。解决方案搜索仅限于一次通道交换的深度，以保持搜索空间宽而浅。
2. 狭窄、深入的搜索：选择多个条纹并同时优化它们。

与以默认通道顺序修剪网络相比，如果在修剪之前对网络进行置换，则可以实现更好的性能。

为了从头开始训练具有 N:M 稀疏性的模型，Zhou & Ma 等人 (2021) 扩展了 STE（直通估计器；Bengio 等人，2013 年），该算法通常用于模型量化中的反向传播更新，可用于幅度修剪和稀疏参数更新。

STE 计算密集参数相对于修剪网络的梯度，并将其应用于密集网络
作为近似值：

扩展版本 SR-STE（稀疏精炼 STE）通过以下方式更新密集权重：

其中 是 的掩码矩阵，而 是元素乘法。SR-STE 的提出是为了通过 (1) 限制 中修剪的权重值和 (2) 提升 中未修剪的权重来防止二元掩码发生较大变化。

与 STE 或 SR-STE 不同，Top-KAST（Jayakumar 等人，2021 年）方法可以在前向和后向传递的整个训练过程中保持恒定的稀疏性，但不需要具有密集参数或密集梯度的前向传递。

在一个训练步骤中，Top-KAST 处理如下：

1. 稀疏前向传递：选择一个参数子集  ，包含每个层按大小排列的顶部参数K，限制为权重的顶部比例D。如果参数化 不在（有效权重）中，则在时间 t 时将参数清零。

     2. 稀疏后向传递：然后将梯度应用于更大的参数子集其中B包含 -比例的权重和 。更新更大比例的权重可以更有效地探索不同的剪枝掩码，从​​而更有可能导致顶部D-比例活动权重的排列。

训练分为两个阶段，B\A集合中的附加坐标控制引入多少探索。预计探索量将通过训练过程逐渐减少，最终使掩码稳定下来。

为了防止富人越来越富的现象，Top-KAST 通过 L2 正则化损失惩罚主动权重的大小，以鼓励更多地探索新项目。在更新过程中，B\A中的参数受到的惩罚比A更大，以获得更高的选择标准，从而稳定掩码。

稀疏化 Transformer

Scaling Transformer（Jaszczur 等人，2021 年）将 Transformer 架构中的自注意力层和 FFN 层都稀疏化，使单示例推理的速度提高了 37 倍。

稀疏 FFN 层：每个 FFN 层包含 2 个 MLP 和一个 ReLU。由于 ReLU 会引入大量零，因此它们在激活上实施固定结构，以在一个N元素块中强制仅存在 1 个非零值。稀疏模式是动态的，每个标记都不同。

其中    中的每个激活对应于W1 中的一列和 W2中的一行。控制器实现为低秩瓶颈密集层，   和  。它使用arg max进行推理以选择哪些列应为非零，并在训练期间使用 Gumbel-softmax 技巧（Jang 等人，2016 年）。因为我们可以在加载 FFN 权重矩阵之前计算   ，所以我们知道哪些列将被清零，因此选择不将它们加载到内存中以加快推理速度。

稀疏 QKV（注意）层：在注意层中，维度  被划分为S模块，每个模块的大小为 。为了确保每个细分都可以访问嵌入的任何部分，Scaling Transformer 引入了一个乘法层（即乘法层逐元素地乘以来自多个神经网络层的输入），它可以表示任意排列，但包含的参数比密集层少。

给定一个输入向量，乘法层输出：

乘法层的输出是一个大小为   的张量。然后它由二维卷积层处理，其中length 和 S被视为图像的高度和宽度。这样的卷积层进一步减少了注意层的参数数量和计算时间。

为了更好地处理长序列，Scaling Transformer 进一步配备了 Reformer（Kitaev 等人，2020）的 LSH（局部敏感散列）注意力和 FFN 块循环，从而产生了 Terraformer。

专家混合模型

专家混合模型 (MoE) 依赖于一组“专家”网络，每个示例仅激活网络子集以获得预测。这个想法起源于 1990 年代（Jacobs 等人，1991 年），与集成方法密切相关。有关如何将 MoE 模块合并到 transformer 中的详细信息，请查看我之前关于大型模型训练技术的帖子和 Fedus 等人 2022 年发表的关于 MoE 的调查论文。

使用 MoE 架构，解码时仅使用部分参数，因此节省了推理成本。每个专家的容量可以通过超参数容量因子C进行调整，专家容量定义为：

其中，每个 token 都会选择顶级-k专家。较大的C会导致更高的专家容量和更高的性能，但计算成本更高。当C>1时，会添加松弛容量；否则，当C<1时，路由网络需要忽略一些 token.

路由策略改进

MoE 层有一个路由网络，用于为每个输入 token 分配一个专家子集。原始 MoE 模型中的路由策略是将每个 token 以不同的方式路由到首选专家，因为它们以自然顺序出现。如果将 token 路由到已达到其容量的专家，则该 token 将被标记为“溢出”并被跳过。

V-MoE（Vision MoE；Riquelme 等人，2021 年）将 MoE 层添加到 ViT（Vision Transformer）中。它与以前的 SoTA 的性能相匹配，但只需要一半的推理计算。V-MoE 可以扩展到 15B 个参数。他们的实验使用了k=2、32 位专家和每 2 位专家的放置位置（这意味着 MoE 被放置在其他每个层中）。

由于每个专家的容量有限，如果一些重要且信息丰富的 token 在预定义的序列顺序中出现得太晚（例如句子中单词的顺序或图像块的顺序），则可能必须丢弃它们。为了避免原始路由方案中的这种缺陷，V-MoE 采用 BPR（批量优先级路由）首先将具有高优先级分数的令牌分配给专家。BPR 在专家分配之前计算每个token的优先级分数（最高k路由器分数的最大值或总和），并相应地改变token的顺序。这保证了专家容量缓冲区将首先用关键token填满。

当  时，BPR 的效果比 vanilla 路由好得多，此时模型开始丢弃大量 token。它使模型即使在非常低的容量下也能与密集网络竞争。

在研究如何解释图像类别与专家关联时，他们观察到早期的 MoE 层更通用，而后期的 MoE 层可以专门用于一些图像类别。

任务 MoE（任务级混合专家；Kudugunta 等人，2021 年）考虑了任务信息，并在任务级别而不是单词或 token 级别路由 token 进行机器翻译。他们以 MNMT（多语言神经机器翻译）为例，并根据目标语言或语言对对翻译任务进行分组。

token 级别路由是动态的，每个 token 的路由决策都是不相交的。因此，在推理时，服务器需要预加载所有专家。相比之下，在给定固定任务的情况下，任务级路由是静态的，因此一个任务的推理服务器只需要预加载k专家（假设是顶级k路由）。根据他们的实验，与密集模型基线相比，任务 MoE 可以实现与 token MoE 类似的性能提升，峰值吞吐量高出 2.6 倍，解码器大小减少 1.6%。

任务级 MoE 本质上是根据预定义的启发式方法对任务分布进行分类，并将这些人类知识纳入路由器。当不存在此类启发式方法时（例如，考虑一般的句子延续任务），如何利用任务 MoE 并不简单。

PR-MoE（金字塔残差 MoE；Rajbhandari 等人，2022 年）每个 token 传递一个固定的 MLP 和一个选定的专家。由于观察到后面层的 MoE 更有益，PR-MoE 在后面层采用更多导出。DeepSpeed 库实现了灵活的多专家、多数据并行性，从而能够跨层使用不同数量的专家训练 PR-MoE.

[未完，待更新]