介绍一种transformer稀疏神经网络的硬件加速器算法

全文4938字,阅读约需15分钟,帮我划重点

划重点

01本文介绍了一种在FPGA上的高效稀疏transformer加速器STA,旨在降低transformer的操作次数和内存大小,提高推理效率。

02STA采用多样矩阵乘法引擎DMME和可扩展的softmax模块,充分利用稀疏模式,实现高计算效率。

03实验结果显示,即使在稀疏度高达75%的情况下,稀疏变换器的性能与密集基准相比仅损失可忽略不计。

04与先前基于FPGA的解决方案相比,STA-4实现了与高端FPGA解决方案可比的延迟,而STA-8则更快1.27-1.49倍。

以上内容由腾讯混元大模型生成,仅供参考

欢迎关注下方公众号阿宝1990,本公众号专注于自动驾驶和智能座舱,每天给你一篇汽车干货,我们始于车,但不止于车。


基于transformer的网络在深度学习中是一股强大的力量。自transformer发明以来,在许多领域如神经机器翻译、语言理解和图像处理都产生了巨大影响。然而,transformer在感知识别处理的性能伴随着高昂的计算和内存成本,这成为基于transformer应用高效部署的重大障碍。

已有学者提出了一种在 FPGA 上的密集脉动阵列加速器以及一种分区方案,以实现transformer的低延迟推理。此外,利用基于块循环矩阵的权重表示来对transformer的加速运算。然而,很多先前由于transformer在高并行任务计算特性使得其无论是在计算资源需求还是运行功耗都很高,很难满足transformer高效推理的需求。

本文介绍一篇经典的论文,通过充分利用稀疏模式,在 FPGA 上提出了一种高效的稀疏transformer加速器,即 STA。它减少了transformer的操作次数和内存大小,减轻了计算和存储的负担。因此,STA 在保持高预测精度的同时实现了高效的transformer推理。


transformer关键架构及稀疏化原理

transformer的关键结构由多头注意力(MHA)残差块(ResBlock)和逐位置前馈网络(FFN)ResBlock 表征。其最消耗时间的计算主要体现在线性层级变换、矩阵乘法以及Softmax计算。   

图片
图1 transformer的关键结构示意图

FFN ResBlock 的输入和输出通过残差连接器连接。FFN ResBlock 内部还有两个线性变换模块以及一个 ReLU 激活函数。MHA ResBlock 的结构更复杂,MHA ResBlock 的输入被分成多个并行头,并对其应用相应的线性投影。它们的结果作为输入并行地送入注意力机制,最后,注意力头的结果连接在一起并传入一个线性层以获得输出线性投影。注意,注意力机制与线性层完全不同,它执行无参数的MatMuls。MHA ResBlock的残差连接器与FFN ResBlock 的组织方式相同。

下图中表示了N:M 稀疏模式下(a) MHA ResBlock 和(b)FFN ResBlock中的操作。两个 Resblock 都是 Transformer 的关键结构。(c) 线性层中 2:4 稀疏参数的图示。   

图片
图 2(a)和(b)分别展示了 MHA 和 FFN ResBlocks 的内部结构

使用稀疏模式可以压缩许多网络模型而不影响其输出精度。在稀疏模式下,参数矩阵被分成多组。每组有 M 个元素,最多包含 N 个非零元素。图 1 展示了一个具有 2:4 稀疏模式的参数矩阵示例。每组 4 个值中只需存储和计算 2 个非零值。在本文中,我们关注如何使用稀疏模式设计一个具有高计算效率的高效transformer硬件加速器。如图 2 所示,在具有稀疏模式的变换器关键结构中,既有稀疏 - 密集矩阵乘法(MatMuls),也有密集 - 密集型 MatMuls。因此,设计一个高效的计算引擎以支持不同类型的 MatMuls 至关重要。此外,softmax 函数嵌入在变换器的 MHA ResBlock 中。为避免 softmax 算子中间结果的频繁片外内存访问,在加速器上实现一个可扩展的 softmax 模块至关重要。


如何为 N:M 稀疏transformer神经网络设计高效的硬件加速器?

首先,需要基于N:M 稀疏 Transformer支持耗时的加速器单元设计。由此,需要考虑如下三种层级的设计思路:
(a) 线性层的稀疏密集 MatMuls
(b) 用于注意力机制的 Dense-dense MatMuls
(c) 用于注意力机制的Softmax

基于以上(a)(b)模块,需要采用多样化MatMul引擎(DMME),并设计可扩展的 Softmax 模块用于 (c)。   

1、STA系统架构

详细的STA 系统级视图如图 3 所示。整个系统是一个 ARM + FPGA 异构架构(如果是智驾领域,则是MPU+NPU的异构域控芯片架构),可分为三部分:ARM(相当于智驾域控的MCU或CPU) 上的软件程序、FPGA 板上的外部内存(相当于智驾域控内芯片外置的LDDR、EMMC或内部Flash内存单元)以及 FPGA 芯片(相当云智驾域控的SOC或AI运算单元)上的定制加速器。当然这种模式和在自动驾驶上的域控中央控制单元是几乎相同的芯片架构。

如下图3所示,整个加速器包括了三个重要的模块:计算模块、控制模块和存储模块。

1)计算模块
- DMME:密集-密集和稀疏-密集、高计算效率下的MatMuls;
- 可扩展的softmax;
- 矢量引擎:偏置添加、残差添加和激活;
- 重构网络:在写回内存之前重新排序数据;

2)控制模块
- 顶部控制器
- 地址生成器
- MatMul 控制器

3)Softmax 控制器
- 存储块
- 重量记忆
- 输入存储器    
- 中间记忆

图片
图 3. STA 的整体架构

软件部分通过 Avalon 总线与 FPGA 通信,传递控制信号和要计算的数据。在初始化阶段,ARM 上的软件程序将transformer结构的可配置参数发送到 FPGA。在 FPGA 加速器进行计算之前,输入和权重预先加载到 DRAM 中。当 FPGA 上的硬件加速器准备好时,软件在 DRAM 中接收相应结果。外部内存与 FPGA 芯片一起存储所有模型参数和输入样本。

基于图 1 所示,DMME 以高计算效率支持多样矩阵计算。可扩展 softmax 模块在变换器的 MHA 子层中执行 softmax 操作。向量单元负责计算密度低的操作,包括偏差加法、残差加法和 ReLU 激活。重排网络在每次 MatMul 计算过程之后以及数据连接之前使用,在将临时结果写回中间片内内存之前对其重新排序。至于片内存储,可分为三部分,包括权重内存、输入内存和中间内存。权重和输入内存分别从片外 DRAM 存储变换器的参数和输入数据。结果也写回输入内存,并传递到 DRAM。一个 ResBlock 中的所有临时结果都将存储在中间内存中,与外部内存无通信。

2、多样矩阵乘法引擎(DMME)

DMME 以高计算效率统一了N:M稀疏transformer中的密集-密集和稀疏-密集 MatMuls。当它执行稀疏 - 密集 MatMuls 时,只加载非零权重参数并选择相应的激活函数来计算,从而提高计算效率。  
DMME 的架构如下图所示。采用两级层次结构设计:
• (Hier-1)H 并行处理子阵列:用于注意力机制的并行密集-密集MatMuls、线性层的平铺稀疏-密集 MatMuls;
• (Hier-2)子阵列中的 TxT 脉动 PE:以时分方式执行稠密和稀疏稠密 MatMuls 复用方式。

图片

如上图它是一个两级层次设计,充分探索了 N:M 稀疏变换器中 MatMuls 内部的并行性。

其中,混合模式收缩期 PE以输出平稳方式,由MAC、非零组成元素选择器 (NZES),多个多路复用器 (MUX) 和寄存器 (REG)。MAC表示为16 位乘法和 32 位累加运算。NZES是根据需要选择合适的数据输入位掩码;MUX是选择正确的数据路径;REG是存储临时数据。

DMME 由 H 个并行的 T×T 脉动阵列组成,每个阵列都能以时分复用的方式高效地执行密集 - 密集和稀疏 - 密集 MatMuls。执行密集 - 密集和稀疏 - 密集 MatMuls 的能力来自内部的混合模式脉动 PE。它由一个 MAC、一个非零元素选择器、多个多路复用器和寄存器组成。MAC 接受两个 16 位输入进行乘法运算,然后将结果与本地 32 位输出部分和累加。非零元素选择器仅在稀疏 - 密集 MatMul 模式下激活,根据输入位掩码选择合适的激活。多路复用器和寄存器分别用于数据路径选择和临时数据存储。


图片
图 4 展示了 DMME 如何以计算高效的方式支持密集 - 密集和稀疏 - 密集 MatMuls

为简单起见,使用 1:2 稀疏模式展示。

1)稠密-稠密MatMuls:
对于稠密矩阵计算MatMuls而言,假设1:2 稀疏模式且点积长度为 4,DMME的每个收缩子阵列均以输出平稳方式执行总共5个周期,最周可以实现高计算效率。

2)稀疏-密集 MatMuls:
对比典型的5 周期运算与我们这里的3周期运算而言,可以明显跳过冗余的修剪操作,节省计算周期和能源,实现高计算效率。

3、可扩展 Softmax 模块

我们知道,softmax这类函数在自动驾驶感知识别分类中大大的缩小运算量。主要体现在,对于所有经过原始激活函数做类型判别后的属性目标经过如下Softmax函数处理后,这些原始分数被转换为了概率分布:
图片

其中,z 是原始分数向量,softmax(z)i是第 i个类别的概率值。对于特定的输入,因为类别的总数 K 是固定的,Softmax函数的分母在计算时是固定的。因此,通过指数化计算较大指数分子对固定的分子相除,对于原始分数较大的类别,其指数值会相对较大,从而导致对应的概率值更接近1。   

可扩展 Softmax 模块可以执行任意长度的 softmax 函数。它保持所有中间结果完全本地,避免片外数据通信。它由三个主要部分组成:一个面积高效的指数函数、一个部分和累加器和一个可扩展的除法器。指数函数使用查找表结合一阶泰勒展开来近似,如图 5 所示。

图片
图 5 展示了我们提出的可扩展 Softmax 架构的细节

一个指数算子可以仅用一个一乘法器和一个加法器来实现。可配置的部分和累加器可以适应各种长度的输入向量,提高了硬件的灵活性。为降低除法的延迟,我们通过级联多个带有流水线的除法器块设计了一个高度并行的除法器,其中一个除法器块仅由减法器和移位器组成。除法器的精度可以通过流水线长度调整。完全流水线化的 softmax 架构可以使前一个向量的除法和当前向量的指数运算重叠。如图 5 所示,该架构有两个可调整参数 P 和 Q,其中 P 表示架构的并行性,Q 表示流水线深度以及输出精度。P个并行输入被送入 softmax 模块,指数输出不仅使用加法树累加,还临时存储在数据缓冲区中。一旦累加过程完成,除法器模块接受累加结果和指数输出以执行 Q - 级流水线除法。

实验结果对比

1、关于模型稀疏性的准确性


首先使用各种 N:M 配置评估模型准确性,以确定最有效的硬件实现。在我们的实验中,从 TED Talks 和 IWSLT’15语料库中选择了一组不同的低资源翻译对作为数据集。我们使用中的 Transformer base 模型在这些数据集上进行训练。我们的密集基准 Transformer 使用中的方法进行训练。所有稀疏模型都使用基于密集基准的 SR - STE 方法微调 40 个 epoch,以进行公平比较。使用BLEU 这一常用的度量标准,用于评估我们的模型准确性。BLEU 分数越高,模型准确性越高。

如表 I 所示,即使稀疏度高达 75%,稀疏变换器的性能与密集基准相比也仅有可忽略不计的损失。还观察到,如果我们将 N 固定为 1,在相同稀疏度下我们可以实现与更高 N 类似的性能。因此,为避免昂贵的硬件开销,我们选择 1:4 作为 N:M 配置用于进一步的硬件性能评估。

图片

2、硬件性能

整个硬件系统使用 SystemVerilog 通过 Quartus Prime 18.1 进行综合。设计的平台是一个英特尔 Arria 10 SX660 SoC 设备。根据所选的 1:4 稀疏度,STA - 4 采用以下配置实现,这是由一个周期准确的模拟器为预先性能评估确定的。DMME 的 (H,T) 为 (4,16),因此 STA4 中有 1024 个 MAC。在 softmax 模块中,P 和 Q 分别为 8 和 16。表 II 展示了我们的 STA - 4 的资源利用情况。

图片

3、跨平台比较

对于跨平台比较,执行变换器推理任务的硬件设置如下。CPU 结果使用英特尔 (R) i9 - 9900X CPU @ 3.50GHz 测量,GPU 结果分别使用 Jeston Nano 边缘计算设备和 NVIDIA RTX2080Ti 卡测量。CPU 和 GPU 结果都在 PyTorch 框架下评估。先前基于 FPGA 的工作使用多达 4096 个 MAC 以及一个 64×64 脉动阵列,而 STA - 4 仅使用其四分之一的 MAC 即可实现可比的推理延迟和更高的计算效率。除了 STA - 4,还进一步构建了 STA - 8,其 N:M 配置为 1:8,用于那些对延迟和功耗比对准确性更关注的应用。

图 6 展示了不同硬件平台相对于 CPU 的理想化性能加速。

图片

这里我们考虑这些基于transformer的模型的所有 MHA 和 FFN ResBlocks 中的单批处理时间,所选模型针对不同应用。TinyBERT是一个用于多种语言任务的轻量级 BERT 模型 。Dino,一个小型视觉变换器,可以作为许多计算机视觉任务的骨干。Transformer - base 模型是神经机器翻译任务的经典模型。如图 6 所示,STA - 4 可以实现与先前高端 FPGA 解决方案可比的延迟。此外,与图 6 中的所有其他解决方案相比,STA - 8 可以实现最低的推理延迟,比先前的 FPGA 解决方案快 1.27 - 1.49 倍。

表 III 展示了在不同平台上无批处理的公平性能比较。

图片

总结

Transformer在深度学习中已不可或缺。然而,由于其大量的计算和内存需求,基于transformer的模型难以实现高效部署。为解决此问题,通过利用细粒度结构化稀疏性,在 FPGA 上提出了一种高效的稀疏transformer加速器 STA。文章设计不仅具有一个能以高计算效率执行稀疏和密集矩阵乘法的统一计算引擎,还具有一个可扩展的 softmax 模块,以避免频繁的片外内存访问,消除了中间片外数据通信的延迟。