Full Stack Optimization of Transformer Inference: a Survey

摘要

分析和研究现有transformer架构中的瓶颈及其与以往卷积模型的异同；
transformer架构对硬件的影响，包括非线性操作、Softmax、GELU，以及线性操作对硬件设计的影响；
优化固定transformer架构的方法；
为transformer模型找到正确的映射和调度操作的挑战；
通过使用神经体系结构搜索来调整架构来优化transformer模型的方法。

Introduction

transformer 与CNN的区别

与传统的cnn聚焦设计相比，transformer 主要由矩阵乘法（矩阵模数）和内存密集型的非线性操作组成。
transformer 模型的计算图和数据流比cnn更复杂，有更多的操作节点，更多的数据流分割和连接。

本文贡献

分析transformer 的运行时特性，并探讨不同的高效transformer 推理方法；
将所调查的方法应用于全堆栈深度神经网络（DNN）加速器发生器Gemmini 上进行案例研究。

这项工作的长期目标是描述跨硬件和软件堆栈的不同因素，以优化transformer 推断。

Transformer 结构和瓶颈

多头注意力

[图片上传失败...(image-ebe0e9-1681201013747)]

一个MHA模块由6个线性操作组成，其中4个是相同的权重到激活矩阵(projections)（即𝑊𝑄、𝑊𝐾、𝑊𝑉、𝑊out投影），其余两个是激活到激活((act-to-act)矩阵（query × key and attention score × value）。这两种矩阵乘法有不同的行为。

FFN

[图片上传失败...(image-3053f7-1681201013747)]

FFN模块是一个由两个线性层组成的相对简单的块。输入序列首先通过权重矩阵𝑊1的第一个线性层从隐藏维𝑑投影到更高的FFN维 $𝑑_{FFN}$ 。随后，投影序列通过权重矩阵𝑊2的第二线性层投影序列投影回原始维数𝑑。一般来说，尺寸 $𝑑_{FFN}$ 比𝑑大4×，导致𝑊1和𝑊2的纵横率为4：1（例如，在BERT-Base中）。在这两个线性层之间是一个非线性层。

Encoder

Decoder

Model analysis

在短序列中，FLOPs和MOPs主要由MHA和FFN投影计算贡献；对于长序列，act-to-act矩阵乘法贡献了更多的FLOPs占比，softmax贡献了更多的MOPs占比。

Softmax operation

指数操作

指数部分容易数值溢出，通常使用减法的技巧

序列的长度方向进行求和
通过除总和进行归一化
$Softmax(x_i) =\frac{e^{x_i} }{\sum_je^{x_j}}$

LayerNorm

计算均值
计算标准差
计算标准化
$p_{out} = \frac{p_{in}-u_t}{\delta_t}*\gamma_e + \beta_e$
非线性算子最大的挑战在于算子融合，LayerNorm 需要实时计算均值和标准差，为了将此操作与前面的matmul操作融合，在写入结果之前，必须在计算均值和方差的维度上累积整个输出矩阵。这导致了不规则的平铺尺寸和较低的数据重用。因此，在将这些操作与之前的层融合与使用更好的平铺维度来最大化重用之间存在一个重要的权衡。第五部分介绍这个平衡方式。