今天为大家带来的是2021年发表在IEEE International Conference for High Performance Computing,
Networking, Storage, and Analysis（SC）上的一篇对Transformer模型的Self-Attention架构的优化。本文针对Self-Attention架构引入了E.T，首先，本文介绍了一种新的自注意力架构，其中包括两个定制的自注意力算子与相应的序列长度感知优化，以及操作重新排序优化。其次，本文提出了一个注意力感知的修剪设计，明智地使用各种修剪算法，以减少更多的计算，从而实现显着缩短周转时间。对于剪枝算法，本文不仅改进了现有的剪枝算法，而且为Transformer模型定义了新的剪枝算法。综合考虑，通过评估E.T.在Transformer、BERTBASE和DistilBERT上的各种基准测试中，表明了优于主流项目上的性能，包括流行的Nvidia企业解决方案，即TensorRT和FasterTransformer。
1 介绍
通过设计新的Self-Attention架构和注意力感知修剪算法。E.T.将单个编码器加速了2.5倍；在Wikitext-2数据集上，模型大小减少了80%。新的自注意力计算比现有的TensorRT实现了2.9倍的加速。图1说明了Self-Attention在TensorRT和E.T.上的推理时间。
图1：WikiText-2数据集上的时间消耗
TensorRT是一个高性能神经网络推理框架，包括网络优化器和运行时引擎。给定一个训练好的神经网络，它通过融合层，选择最佳的内核实现，并根据网络结构重新组织计算来优化网络。
1）消除未使用的输出层，避免不必要的计算。
2）在可能情况下，融合卷积、偏置、ReLU层形成单一的层。
3）如果多个层采用相同的输入张量，或者同时使用类似的参数执行相同的操作，则它也融合这些层，这被称为水平层融合。
2 相关工作
2.1.新的Self-Attention算子
2.1.1．自定义自注意力算子
图 2：采用On-the-fly的Self-Attention架构
对于Self-Attention架构，本文在以下两方面进行了改进，用于加速Self-Attention训练推理。如图2所示，首先，将Scaling操作提前，用于实现纯FP16的注意力计算，保证纯FP16在计算期间不会遇到溢出问题；其次，采用外积减少Q和K进行GEMM时的访存量：包含加载Q的第一列和K的第一行进行计算（分块算法），减少了对Q和K的访存量。Q的Row i和Q_1*𝐾_1^T的Row i存储在共享内存，所需的总的内存大小如下：
同时，本文采用预计算线性变换算子，进一步加速模型训练推理时间。具体过程为将第七步的最终输出矩阵，重新变换计算得到Output输出矩阵，提高模型并行性。采用预计算线性变换算子与原设计前产生的结果完全相同。使用了预计算线性变换算子后的模型如下：
图 3：
采用预计算线性变换算子的Self-Attention架构
2.1.2．纯FP16注意力计算的重排序
在进行实验过程中，采用WiKiText-2数据集进行训练计算时，存在精度溢出问题，Self-Attention中的Q和K相乘的结果超出了FP16的范围[-65504,65504]，溢出情况如图4所示。为解决Q和K相乘时的精度溢出问题，本文将Scale操作提前来防止精度溢出。在这种情况下，需要使用混合的基于精度的张量操作，即Q_1*𝐾_1^T中的每一行都必须存储在FP32中，并转换成FP16再进行后续操作。与纯FP16相比，混合精度引入了两个开销：
2）将FP32转为FP16后再进行Mask和Softmax。
图 4：Q和K相乘的结果中精度溢出的部分
2.2.注意力感知模型修剪
2.2.1．Transformer模型的张量瓦片修剪
现有的剪枝算法包括行剪枝、列剪枝和不规则剪枝。尽管行/列/不规则修剪的权重矩阵可以转换为利用张量核的格式，但需要在预处理输入和后处理结果矩阵上支付较大的开销。本文介绍一种基于张量瓦片的修剪算法，该算法要么排除权重的整个张量瓦片，要么在权重矩阵内保持整个张量瓦片完整。显然，这种基于张量切片的修剪引入了更好的存储和控制流逻辑。图5为不同剪枝算法的剪枝效果，a是采用片和行的剪枝、b是不规则剪枝、c是列剪枝、d是基于片的剪枝。
图 5：不同剪枝算法的效果对比
2.2.2．自适应感知剪枝算法
图 6：重新加权训练、修剪和掩蔽再训练
2）判断当前矩阵是否落入预定义迭代中，若是，则划分、计算矩阵L2范数更新损失函数。
4）再进行划分和计算L2范数得到掩码矩阵，用掩码矩阵更新1中初始化的矩阵并再次训练，得到剪枝后训练的参数矩阵。
以Transformer模型训练为例，应用了该剪枝算法后的训练过程为：首先从头开始训练一个模型50个epoch，将得到的训练模型用作我们的预训练模型；其次，加载预训练的模型，运行50个epoch基于重新加权的训练；最后，对先前训练的模型应用修剪后再进行50个epoch的再训练。
3 实验
3.1.编码器延迟对比
图7：不同推理器延迟对比
图7显示了本文的编码器与目前编码器的延迟对比，包括TensorRT，FasterTransformer和PyTorch。当稀疏度低于40%时，进行注意力感知修剪。总的来说，可以观察到E.T.在所有稀疏情况下都优于PyTorch、TensorRT和FasterTransformer。值得注意的是，随着修剪率的增加，PyTorch、TensorRT和FasterTransformer的最大加速比分别提高至13.7倍、3.4倍和2.5倍
