基于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…