论文解析——Full Stack Optimization of Transformer Inference: a Survey

作者及发刊详情

摘要

正文

主要工作贡献

这篇文章的贡献主要有两部分：

1. 分析Transformer的特征，调查高效transformer推理的方法
2. 通过应用方法学展现一个DNN加速器生成器Gemmini的case研究

1）分析和解析Transformer架构的运行时特性和瓶颈

2）Transformer推理的硬件架构

3）对特定Transformer架构的优化策略，比如剪枝和量化

4）Transformer架构下操作的映射和调度，以及相关挑战

5）通过自动化的神经架构搜索过程，设计和调整transformer架构，使其硬件更加高效

Transformer模型架构和性能瓶颈

transformer的基本架构

Transformer架构包含两个模块：MHA和FFN
Transformer架构的参数如下：

对Transformer架构的输入序列包含l个token，每个值都由一个d维度的向量表示，构成了$d\ast l$的矩阵。token可以是一个词或一个句子片段。

MHA的计算特征

MHA有三种不同的权重$W_Q$、$W_K$、$W_V$，具体执行流如图1所示，这些权重将会产生三种不同的激活，每种激活会被分成h个块（chunk，有隐藏维度d/h），因此这些块都被分成了h个不同的注意力头。q块和k块沿着隐藏层相乘，生成 $l\ast l$ 大小的激活矩阵，这些激活矩阵经过softmax操作，与v块相乘，得到attention头的激活，然后h个head组合生成结果$W_{out}$。每阶段的计算结果如Table2所示。最终线性层的输出经过残差链接和层归一化生成MHA模块的输出。

MHA总共有6个线性操作，其中4个是权重到激活的矩阵乘（$W_Q$、$W_K$、$W_V$、$W_{out}$），另外两个是激活到激活的矩阵乘（$query\ast key$、$attention.score\ast value$），本文将前者称为投影（projection），后者称为激活到激活矩阵乘。

FFN包含两个线性层操作，$d_{F}FN$通常是$d$的四倍，在两个线性层间有一个非线性层。

非线性操作的特征

非线性操作，包括Softmax, LayerNorm, 和 GELU需要片外计算的支持，虽然全部操作中占据了较小部分，但比矩阵乘更具挑战，如果处理不当将会产生额外的开销。

在有效利用临时内存和高效计算方面提出了挑战

需要在运行时传递所有的输入值，这些值都会保存在临时存储中。
比如softmax操作包括求指数操作、跨序列长维度的求和、每个指数向除以求和结果的归一化操作，需要解决指数溢出问题。
计算LayerNorm函数还需要跨隐藏维度多次传递整个输入值，先计算平均值，再计算标准差，然后再使用层归一化

encoder和decoder架构

encoder可以并行处理输入序列
encoder-only适合用于自然语言理解任务（sentiment analysis、sentence similarity analysis）
decoder一次只能推理出一个token，适合用于生成式任务

产生token的共同优化方法：在后续的迭代中，缓存和复用之前生成的token中间的K和V值

生成的token后继续传入到decoder的输入，复用可以节省计算时间。

模型的算力密度

矩阵乘在encoder-only和decoder-only中消耗了99%的FLOPS。

模型优化

硬件设计

参考文献

评

这是一篇关于Transformer推理的全栈技术综述