基于Transformer架构的大模型推理硬件加速器设计

概述

当前大模型的基础架构正在向 Transformer 结构收敛¹，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)。

在这里插入图片描述

计算行为

输入位置编码向量与词嵌入相加

矩阵的加法操作，矩阵元素是浮点数。

自注力矩阵的计算

$QK^T$ 的矩阵乘计算， $\sqrt{d_k}$ 的根号运算，以及两个结果相除计算
再将结果进行softmax函数计算，包括乘法、求和、除法运算。

softmax操作

$softmax(\frac{QK^T}{\sqrt{d_{k}}})$

指数操作可能导致数值溢出问题

maximum subtraction trick[^5]可以把
在这里插入图片描述
转变为：

其中 $x_{max}$ 是 $x_j$ 中的最大值

multi-Head Attention的conact和linear操作

残差求和操作

层归一化

平均数求和和除法操作
求方差操作
归一化处理（求差、除法等）

FFN的两级线性乘加运算

总体特征

矩阵乘占据着大部分的计算时间

如下图所示是LLama7B模型，batch_size=1，输入长度为128场景下各transformer操作在GPU上的耗时结果。
在这里插入图片描述

MHA具有较强的并行计算能力

大量的缓存以及访存操作

大模型的权重数据多，计算过程中产生大量的中间数据，许多非线性函数的计算数据间存在调用依赖，这些操作都需要硬件大量的缓存数据，因此LLM运算硬件需要延迟小、容量大的存储器件。
例如GPU平台的计算过程，将权重数据从显存（HBM）加载至on-chip的SRAM中，然后由SM读取并进行计算。计算结果再通过SRAM返回给显存。

序列长度增加会导致自注意力计算的占比增加

图片来源于文章²
在这里插入图片描述

模型压缩与优化

由于芯片面积和成本的限制，现阶段的大模型很难在一颗ASIC上部署，因此需要对模型进行压缩。
除了通用的模型量化和剪枝方法，对Transformer架构的模型有专用的优化方法。

网络剪枝

移除网络中权重接近零的连接剪掉，是网络稀疏化的一种手段。

数据量化

数据量化通过将高精度的
人工智能和机器学习入门

将数据精度从浮点型变成int，例如I-bert[^{6]，这是一个int版本的BERT，文章[}7]又基于I-bert做了专用硬件上的适配（3.4章节）。
GitHub: kssteven418 / I-BERT

张量拆分

算子融合

将必要的算子操作都融合在一起，灵活性低，但是性能更好，因为访存减少了。
例如将multi-head attention中的线性运算和非线性运算融合在一起。

硬件架构设计

稀疏注意力

稀疏注意力即让模型阈值低关联度的输入组合，并跳过注意力计算。该方法可以有效优化 $O(n^2)$ 的注意力算子³，优化极长序列时的性能表现。
具体包括将计算中对最终结果影响不重要的数据查找，并近似为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）

MIT工作，网络中靠前的层级连剪枝，靠后的层跨层预测为0，减小计算

基于随机投影的注意力预测

清华大学秦育彬³

注意力层优化

Eager Prediction

预测Softmax的输出，

流水化的全Attention计算硬件加速

将加速器实现加速颗粒度不仅仅局限于矩阵运算、或其他非线性运算，将attention完整的计算通过流水化的方式实现。

硬件设计优化

优化乘法器设计，去除冗余的注意力计算

双缓冲可以使访存操作和计算执行操作重叠执行

对于专用加速器设计，通常会把权重、KV cache、Q、K、V矩阵等原数据的值存储到临时存储中，此时该存储的设计带宽、延迟等指标影响了整体计算的性能。乒乓策略等

参考文献

AI云原生智能算力架构: AI智算及算力全产业链研究报告 2024 ↩︎
DOTA: Detect and Omit Weak Attentions for Scalable Transformer Acceleration ↩︎
秦育彬：高能效注意力模型计算架构设计 ↩︎ ↩︎