Full Stack Optimization of Transformer Inference
一、文章介绍
背景
Transformer模型被广泛应用于各种任务，如计算机视觉、自然语言处理、语音识别等，原因是它们的准确度很高。然而，这些模型的复杂性和规模不断增加，导致它们在推理阶段需要大量的计算资源和带宽，特别是在那些对延迟敏感的应用场景中，部署这些模型变得非常困难。
全栈优化的方法
为了解决这些问题，本文提出了一个“全栈”优化方法，旨在提升Transformer模型推理的效率。作者从硬件到软件进行了全面的分析，找出如何使现有的Transformer模型在计算和带宽上的需求最小化。
1. 硬件优化
作者详细分析了Transformer模型的架构对硬件的要求，尤其是非线性操作（如Layer Normalization、Softmax、GELU）对计算效率的影响。
通过调整硬件配置，改进如**矩阵乘法（matmuls）**等线性操作，减少非线性操作的计算开销。
1. Transformer架构对硬件的要求
Transformer模型内部有很多非线性操作，比如：
Layer Normalization：归一化
Softmax：用于归一化注意力机制中的权重分布，使得输出符合概率分布。
GELU（高斯误差线性单元）：一种激活函数
这些操作与线性操作不同，通常更耗时且计算复杂。这是因为非线性操作往往需要浮点计算，并且这些计算具有高带宽需求（也就是说，它们对数据的存储和传输要求很高）。
在硬件层面，特别是那些专门为深度学习设计的加速器上，这些操作并没有得到充分的优化支持，导致Transformer推理的效率不高。CNN的设计大部分依赖矩阵乘法，而Transformer模型需要更多的非线性操作，这带来了额外的计算负担。
2. 硬件配置的调整与优化
作者通过对硬件的配置进行调整来优化这些操作，特别是：
优化线性操作（矩阵乘法，matmuls）：在Transformer模型中，大量的计算是矩阵乘法。虽然这是线性操作，但如何有效地利用硬件的计算单元进行矩阵乘法，是提升推理效率的关键之一。
减少非线性操作的计算开销：通过专门为Transformer设计的加速器（比如Gemmini加速器的改进版本），可以减少对非线性操作（例如Softmax、LayerNorm、GELU）的浮点计算需求。作者提出一种方法，将这些操作转换为整数运算（integer-only BERT），通过整数运算来近似这些非线性操作，从而显著减少计算时间和带宽开销。
2. 神经架构搜索 (NAS)
这种方法的效果显著，最终实现了88.7倍的端到端推理速度提升，同时几乎没有性能下降。
采用了自动化的NAS，以进一步优化Transformer网络。通过这种方法，可以在硬件的基础上找到更高效、性能更佳的Transformer结构。
主要结论
Gemmini加速器优化：Gemmini最初是为卷积神经网络（CNN）设计的，但作者对它进行了改进，以支持Transformer模型的推理。通过调整硬件配置，如支持只用整数计算的BERT变体，性能提升了39.6倍。
操作融合的权衡：虽然在CNN中将操作融合（如BatchNorm与卷积融合）很直观，但在Transformer中，某些操作的融合反而会增加计算开销，因此需要特别的小心处理。
神经架构搜索：通过自动化的神经架构搜索，作者发现了一种可以优化硬件效率和模型性能的Transformer结构，能让EDP（能量-延迟乘积）减少10.6倍。
二. 硬件架构优化
1. 基线加速器
Gemmini加速器：文中使用了Gemmini生成的一个典型的CNN加速器，这个加速器最初是为ResNet50这样的卷积神经网络（CNN）工作负载优化的。其架构包括：
16×16的Systolic阵列：负责高效的矩阵乘法操作。
8-bit整数权重和输入存储在256KB的局部内存中。
32-bit部分和累加结果存储在64KB的累加器SRAM中。
外部缓存和DRAM：当网络层太大无法适配局部内存时，数据会被放入外部缓存和DRAM。
卷积操作优化：这个基线加速器中包含对ReLU和最大池化操作的原生支持，这些操作是卷积神经网络中常见的。然而，这个加速器没有为Transformer的特定操作（如LayerNorm、Softmax、GELU）提供支持，导致在运行BERT等Transformer工作负载时性能大大降低。
2. 性能瓶颈
当基线加速器用于Transformer推理时，硬件利用率非常低，特别是对功能单元的利用率小于1%。主要的原因是：
虽然矩阵乘法的利用率达到了74%，但非线性操作（如LayerNorm、Softmax等）需要由主处理器（CPU）来执行，因为加速器不支持这些操作。
这些操作往往依赖浮点运算，其能效远低于整数运算。同时，浮点运算需要进行数据的量化和反量化（De/Quantization），这些开销占据了总执行时间的96%。
在Transformer推理中，虽然大部分的浮点运算集中在矩阵乘法上，但实际的性能瓶颈是非线性操作。因此，优化这些操作成为关键。
3. 内存配置的调整
通过实验观察到，调整输入/权重局部内存和累加器的大小可以显著提升矩阵乘法（如查询×键的矩阵乘法）的性能：
增大累加器的容量，尤其适用于计算较大的矩阵乘法（如输出矩阵大小为 l×ll * ll×l 时），能够提高数据复用性并减少数据传输开销。
因此，作者将局部内存从256KB减小到64KB，同时将累加器从64KB增大到256KB。这些改变并没有增加总的SRAM容量或总面积，但大幅减少了矩阵乘法的总延迟（约36%）。
4. 硬件-软件协同设计
为了减少量化和反量化的开销，作者将基础的Transformer工作负载（原始的BERT模型）转化为一个整数版本的BERT（I-BERT），这使得所有的非线性操作可以用整数多项式近似来加速计算。
整数实现的GELU、LayerNorm和Softmax被添加到CNN加速器中。
通过引入新的“归一化单元”，加速器能够直接计算这些操作所需的数值（如求和、平方和、最大值等），避免了将计算任务传递给CPU执行。
这些矩阵乘法的结果会通过多个累加器读取步骤，计算所需的所有归约操作，最后再通过16个激活单元并行处理GELU和其他非线性操作。
通过引入这些优化特性，BERT推理性能提升了39.6倍。如图2所示，优化后，瓶颈再次回归到矩阵乘法，而非归一化或激活函数。
这种优化不仅消除了量化和反量化的必要，且通过将GELU和前置的矩阵乘法融合，使它们变成一个流水线操作，从而进一步提高了效率。
在硬件实现方面，新的硬件单元仅增加了加速器14%的总面积，且这些操作只增加了9.3%的功耗。
个人理解：
挟带线性运算的非线性运算：PE+功能单元
纯非线性运算：功能单元
归一化单元负责像LayerNorm、Softmax等操作的数值归约计算，如求和、平方和、最大值等，这些操作虽然非线性，但并没有改变它们的一些基础数值操作步骤。
激活单元负责处理激活函数（如GELU）的计算，这些非线性操作使用整数多项式近似来加速，避免复杂的浮点运算。
三． 调度优化
调度优化中的挑战
在第二节，讨论了Transformer模型中的非线性操作（例如LayerNorm和Softmax）对加速器设计的挑战。因为这些操作的计算特性，导致它们成为瓶颈，使得优化加速器调度变得至关重要。这里的重点是，如何在矩阵乘法和后续的非线性操作之间进行调度优化，以减少推理延迟。
非线性操作与高算术强度操作的融合
在DNN的调度中，一个常见的优化策略是将高算术强度的操作（如卷积、矩阵乘法）与低算术强度的操作（如归一化、激活函数）进行融合。比如：
- 在CNN中，卷积操作与ReLU、最大池化等操作的融合是很常见的。这种方法能够提高算力利用率，隐藏低效操作的延迟。
- 在量化网络的场景中，待归一化的部分和最终输出的比特宽度通常不一致。部分和比最终归一化输出的比特宽度要高，因此融合这些操作可以减少等待时间。
Transformer编码器中的融合挑战
类似的策略可以应用到Transformer编码器中，尤其是将LayerNorm和Softmax与它们前面的矩阵乘法进行融合。这样做的目标是通过融合操作，减少等待时间并提高计算效率。然而，执行这些融合策略可能会带来硬件和软件上的变化，包括：
- 在像Gemmini这样的DNN加速器中，可能需要额外的硬件支持来直接访问矩阵乘法后的部分和，供归一化操作使用。
- 矩阵乘法的执行调度也需要合适的约束条件，特别是需要最大化输出维度的tile（块）的大小，以确保行/列能在累加器中立即准备好，用于LayerNorm等归一化操作计算。
这种调度方法被称为融合优化调度，与传统的非融合调度相比，它试图通过融合多个操作来隐藏延迟。
图3：对比两种调度策略
图3通过对比非融合调度和融合优化调度，进一步分析了BERT模型的多头注意力机制模块（MHA）和前馈神经网络模块（FFN）在这两种调度下的性能差异。
1. MHA模块（图左）
图左展示了BERT基础模型中MHA的延迟分布。
- 非融合调度：这是一种传统调度方式，其中每个操作（如矩阵乘法、Softmax、LayerNorm等）都是独立执行的。由于Softmax是一个低算术强度的操作，它占据了大部分的延迟。Wout投影（Wout projection）和LayerNorm也是独立进行的，增加了整体的执行时间。
- 融合优化调度：在这种情况下，矩阵乘法（Q × K）与Softmax操作融合，使得Softmax的高延迟得以隐藏。Softmax操作的延迟占比显著减少，整体的MHA模块延迟减少了78%。这表明将Softmax与矩阵乘法进行融合，能够大大提升性能。
2. FFN模块（图右）
图右展示了BERT基础模型中前馈神经网络模块（FFN）的延迟分布。
- 非融合调度：这里，W2投影和LayerNorm是独立执行的，每个操作都有相对合理的延迟。
- 融合优化调度：当将W2投影与LayerNorm操作融合时，反而增加了总延迟，增加了27%。这表明，在FFN模块中，融合W2投影和LayerNorm并没有带来优化，反而因为某些原因（如资源争用）导致了性能下降。
在MHA模块中，融合矩阵乘法与Softmax操作能够有效地隐藏Softmax的延迟，减少总执行时间，提升性能。因此，这是一种非常有效的调度优化策略。
在FFN模块中，融合W2投影与LayerNorm操作并不适合。融合优化反而导致了额外的延迟，这表明对于不同的模块，调度策略需要灵活选择，不能一概而论。
四、实验
实验设置
基线架构：
使用了一个6层的Transformer模型，其他配置保持与BERT-Base模型一致。该模型基于WikiText-2数据集进行语言建模训练，使用了Adam优化器和线性学习率调度策略。具体参数如下：
训练50个epoch。
序列长度为512，batch size为16。
NAS训练：
作者采用了BigNAS策略，使用与基线模型相同的超参数来训练一个“超网”（supernet）。这个超网包含了不同的架构组合，搜索空间涵盖了以下维度：
Transformer的层数：3、4、5、6层。
注意力头的数量：4、6、8、10、12。
隐藏层维度：384到768。
前馈网络（FFN）维度：768到3072。
在超网训练完成后，作者使用了进化搜索（evolutionary search），通过40次迭代和40个个体的种群，结合0.2的变异概率，来从超网中搜索出最优子网（subnets）。每次迭代后，只有在**EDP（能量-延迟乘积）**和困惑度（perplexity）上表现最优的子网会被保留。
硬件成本评估：
为了快速评估每个子网在目标硬件上的延迟和能耗，作者使用了一个基于查找表的方法。这个查找表包含了通过Timeloop模拟得到的每个操作的延迟和能耗，最终通过将这些操作汇总来估计整个网络的端到端延迟和能耗。
搜索完成后，Pareto最优的子网会通过RTL仿真进一步精确估计延迟。对于能量测量，作者继续使用Timeloop的数据。
实验结果
在图4中，作者展示了不同Transformer架构在EDP、延迟和能耗方面的Pareto前沿结果，以及基线模型的性能作为对比。
1. 左图：EDP与困惑度的关系
- 该图展示了EDP（能量-延迟乘积）与模型困惑度的关系，横轴是归一化的EDP，纵轴是困惑度。
- 图中展示了通过**NAS（蓝色点）**找到的多个架构点，以及基线Transformer模型（×号）。
- 可以看到，通过NAS搜索找到的架构在EDP和困惑度之间的权衡上表现得更好。某些架构在较低的硬件成本（EDP较低）的同时，仍能保持较好的困惑度性能（类似甚至优于基线模型）。
2. 中间图：延迟与困惑度的关系
- 该图展示了延迟与模型困惑度的关系，横轴是推理延迟（以10⁹个周期为单位），纵轴是困惑度。
- 通过NAS搜索，作者找到了一些能够在容忍0.1点困惑度下降的情况下，延迟减少1.4倍的架构。如果允许困惑度下降1点，延迟可以减少2.4倍。
3. 右图：能耗与困惑度的关系
- 该图展示了能耗与困惑度的关系，横轴是能耗（以10⁻³焦耳为单位），纵轴是困惑度。
- 在能耗方面，通过NAS搜索到的架构也有显著改进。在容忍0.1点困惑度下降的情况下，能耗可以减少1.6倍，如果允许困惑度下降1点，能耗可以减少4.4倍。