体系结构论文（五十五）：Full Stack Optimization of Transformer Inference

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较低）的同时，仍能保持较好的困惑度性能（类似甚至优于基线模型）。