概述
当前大模型的基础架构正在向 Transformer 结构收敛1,Transformer架构自谷歌2017年提出后比较稳定,因此针对Transformer的计算设计专用的ASIC加速器很有必要。
尤其是“Attention is All you Need”==》“Money is All you Need”,哈哈哈。。。
推理加速、推理硬件加速的研究范畴
推理加速是指通过网络结构优化、模型剪枝、量化等手段加快模型收敛,减小模型推理时间
硬件推理加速指通过对模型网络及计算定制化设计硬件电路,从而加速模型计算过程;但由于原始网络规模较大,精度较高,直接设计电路需要的晶体管规模和芯片面积较大,因此需要对模型进行压缩,在精度不丢失的前提下高效地设计硬件电路。
二者虽然都是加速模型的推理过程,但研究重心不一样,前者关注模型全局计算过程的减少,后者关注模型局部计算过程的加速。
专用加速ASIC
范益波@知乎:Vision Transformer低功耗硬件加速器
Sohu
2024年,6 月 26 日,美国新兴的芯片创业公司 Etched 发布其首款 AI 芯片——Sohu,该芯片只能运行Transformer模型。
Etched官网
Etched官网:创始团队宣言
澎湃新闻报道
Transformer架构的计算特征
TransFormer的权重数据分为4类,包括:多头注意力输入线性层权重(Q, K, V)、多头注意力输出线性层权重(O)、前馈神经网络第1层线性层权重(FFN1)和第2层线性层权重(FFN2)。
计算行为
输入位置编码向量与词嵌入相加
矩阵的加法操作,矩阵元素是浮点数。
自注力矩阵的计算
Q
K
T
QK^T
QKT的矩阵乘计算,
d
k
\sqrt{d_k}
dk的根号运算,以及两个结果相除计算
再将结果进行softmax函数计算,包括乘法、求和、除法运算。
softmax操作
s o f t m a x ( Q K T d k ) softmax(\frac{QK^T}{\sqrt{d_{k}}}) softmax(dkQKT)
指数操作可能导致数值溢出问题
maximum subtraction trick[^5]可以把
转变为:
其中
x
m
a
x
x_{max}
xmax是
x
j
x_j
xj中的最大值
multi-Head Attention的conact和linear操作
残差求和操作
层归一化
平均数求和和除法操作
求方差操作
归一化处理(求差、除法等)
FFN的两级线性乘加运算
总体特征
矩阵乘占据着大部分的计算时间
如下图所示是LLama7B模型,batch_size=1,输入长度为128场景下各transformer操作在GPU上的耗时结果。
MHA具有较强的并行计算能力
大量的缓存以及访存操作
大模型的权重数据多,计算过程中产生大量的中间数据,许多非线性函数的计算数据间存在调用依赖,这些操作都需要硬件大量的缓存数据,因此LLM运算硬件需要延迟小、容量大的存储器件。
例如GPU平台的计算过程,将权重数据从显存(HBM)加载至on-chip的SRAM中,然后由SM读取并进行计算。计算结果再通过SRAM返回给显存。
序列长度增加会导致自注意力计算的占比增加
图片来源于文章2
模型压缩与优化
由于芯片面积和成本的限制,现阶段的大模型很难在一颗ASIC上部署,因此需要对模型进行压缩。
除了通用的模型量化和剪枝方法,对Transformer架构的模型有专用的优化方法。
网络剪枝
移除网络中权重接近零的连接剪掉,是网络稀疏化的一种手段。
数据量化
数据量化通过将高精度的
人工智能和机器学习入门
将数据精度从浮点型变成int,例如I-bert[6],这是一个int版本的BERT,文章[7]又基于I-bert做了专用硬件上的适配(3.4章节)。
GitHub: kssteven418 / I-BERT
张量拆分
算子融合
将必要的算子操作都融合在一起,灵活性低,但是性能更好,因为访存减少了。
例如将multi-head attention中的线性运算和非线性运算融合在一起。
硬件架构设计
稀疏注意力
稀疏注意力即让模型阈值低关联度的输入组合,并跳过注意力计算。该方法可以有效优化
O
(
n
2
)
O(n^2)
O(n2)的注意力算子,优化极长序列时的性能表现。
具体包括将计算中对最终结果影响不重要的数据查找,并近似为0,从而减少整体计算量。
-
A3
论文解析——A3: Accelerating Attention Mechanisms in Neural Networks with Approximation -
ELSA(ISCA’21)
-
Sanger(MICRO’21)
-
DOTA(ASPLOS’22)
将输入采用低秩分解,通过低秩的近似完成系数法处理器 -
近似计算乘法器(ISSCC’22)
【论文解析】A 28nm 27.5TOPS/W Approximate-Computing-Based Transformer Processor with Asymptotic Sparsity…
- SpAtten(HPCA’21)
流水化的全Attention计算硬件加速
将加速器实现加速颗粒度不仅仅局限于矩阵运算、或其他非线性运算,将attention完整的计算通过流水化的方式实现。
临时存储硬件设计优化
对于专用加速器设计,通常会把权重、KV cache、Q、K、V矩阵等原数据的值存储到临时存储中,此时该存储的设计带宽、延迟等指标影响了整体计算的性能。
双缓冲可以使访存操作和计算执行操作重叠执行
乒乓策略等
参考文献
AI云原生智能算力架构: AI智算及算力全产业链研究报告 2024 ↩︎
DOTA: Detect and Omit Weak Attentions for Scalable Transformer Acceleration ↩︎