加速AI算法中矩阵乘法的FPGA实现

作者: 张海军 来源:傅里叶的猫

随着大模型越来越火,算力成了绝对的稀缺资源,因此现在很多的研究方向都是如何降低AI算法中的运算,就像我们前面那篇博客中讲到的,MatMul-free LM方法通过完全消除矩阵乘法,为大规模语言模型设计了一种新型的算法。这不仅在内存使用上带来了显著的节省,还在GPU和FPGA上实现了高效的计算。

加州大学颠覆性大模型架构,FPGA重回AI舞台?

这篇文章我们继续看一篇优化矩阵算法的论文,题目是《Fast, Scalable, Energy-Efficient Non-element-wise Matrix Multiplication on FPGA》

现代神经网络(NN)架构严重依赖大量的乘法累加算术运算,构成了主要的计算成本。因此,本文在FPGA上提出了一种高吞吐量、可扩展且高能效的非元素矩阵乘法单元,作为神经网络的基本组件。论文简化了基于LUT的近似矩阵乘法MADDNESS算法的层间和层内冗余,以设计一个快速、高效的可扩展近似矩阵乘法模块,称为“近似乘法单元(AMU)”。AMU 通过专用的内存管理和访问设计进一步优化了基于 LUT 的矩阵乘法,将计算开销与输入分辨率分离,并显著提高了基于 FPGA 的 NN 加速器效率。实验结果表明,与基于FPGA的量化神经网络(QNN)加速器的先进解决方案相比,使用我们的AMU可实现高达9×的吞吐量和112×的能效。

AMU架构概述

AMU 基于 MADDNESS 架构,由分配器、编码器和聚合器组成。然而,AMU根据提出的三种优化技术优化每个计算块。这三个处理单元分别提供分配块、查找原型、获取部分结果和打包乘法结果函数。

AMU.png

AMU基于MADDNESS架构,主要包含三个关键组件:分配器(Allocator)、编码器(Encoder)和聚合器(Aggregator)。这些组件协同工作,实现了对输入特征图的高效处理和矩阵乘法的近似计算。

分配器(Allocator)
  • 功能:接收并解包输入特征图的数据包,将数据分配给多个并行的编码器。
  • 工作流程:输入特征图通过一系列数据包(每个包含M块)到达AMU。分配器负责解包这些数据包,并根据输入特征图的组织方式,将数据块分配给N个并行的编码器进行处理。
编码器(Encoder)
  • 功能:每个编码器负责编码输入数据包中的块,以确定对应原型的地址。
  • 工作流程:编码器按顺序读取输入块,使用当前块和相应的分割值来确定下一轮的ID。这一过程涉及到查找数组中的分割值,并将其与当前块进行比较,以识别出原型地址。
聚合器(Aggregator)
  • 功能:聚合来自编码器的部分结果,以产生最终的输出特征图。
  • 工作流程:聚合器收集所有编码器的结果,并将它们合并,以形成完整的输出特征图。

参数压缩

AMU通过参数压缩进一步优化计算。在原始的MADDNESS算法中,需要加载所有查找表(LUTs)来计算输入特征图与权重矩阵的近似矩阵乘法。但是,通过I/O修剪,AMU只需要加载远小于原算法所需数量的LUTs,这大大减少了资源占用,提高了效率。

内存管理与访问设计

AMU优化了内存管理,将LUTs重塑并分配到一个二维数组中,以便于并行计算M块并按顺序发送O包。这种设计允许AMU以较低的延迟和较高的并行性进行矩阵乘法计算。

总体架构

AMU架构通过流水线输入和输出数据包,从前一个AMU接收输入,处理后传递给下一个AMU,形成一个连续的处理链。每个AMU处理输入特征图的一部分,然后将结果传递给下一个单元,直到生成最终的输出特征图。

性能优化

AMU的设计使得它可以处理非逐元素的矩阵乘法,通过减少计算开销和提高资源利用率,实现了高吞吐量和能源效率。与传统矩阵乘法器相比,AMU在处理大规模神经网络时展现出显著的优势,特别是在量化神经网络(QNN)加速器上。

AMU架构通过上述组件和策略,不仅优化了计算流程,减少了不必要的计算和存储,而且提高了矩阵乘法的执行效率,降低了功耗,非常适合用于高性能、低功耗的神经网络加速器设计。

Hardware Design Space Exploration

AMU是一种替代传统的逐元素矩阵乘法的算法,它通过LUT(查找表)-基存储和优化技术,实现了更高效的矩阵运算。为了充分利用FPGA资源,研究者们探索了几个关键点:

  1. 聚合器设计:聚合器在矩阵运算中扮演重要角色,它等待所有ID准备就绪后开始工作。每个CodeBook都有一个独特的原型ID,范围从0b0000到0b1111(总计2I个原型)。例如,如果ID 1101是第7个CodeBook的原型地址,那么位于(13, 6)位置的所有LUT上的单元会被挑选作为输入特征图和权重的点积结果。这些被选中的单元(灰色单元)被加总,并在需要时加上偏置,然后通过连续的阈值操作,相当于量化神经网络中的缩放、批量标准化和均匀量化激活,从而获得O×M的区块,再管道化地发送O个包到下一个AMU。聚合器包含从离线训练阶段获得的O×M个LUT,每个LUT有I×N个近似的部分乘积结果。通过调整O/I和M/N的比例,可以达到期望的资源效率-精度平衡,以适应基于AMU的神经网络推理加速。
  2. 内存分配设计:为了充分利用设备的存储资源,研究者分析了LUT的阵列结构。由于每个CodeBook只能被识别为特定的原型,因此仅需要一个部分乘积。然而,为了存储所有原型和权重矩阵的点积结果,LUT可能会大量占用主内存资源。因此,合理分配LUT内存和设计聚合器的内存访问方式对于提升AMU的资源利用和吞吐量至关重要。
  3. 硬件设计空间的探索:与基于元素的矩阵乘法不同,AMU通过非结构化的剪枝和不典型的内存访问模式实现矩阵乘法,这导致了不规则的内存访问行为,为设计高效的加速器带来了额外挑战。在编码和聚合阶段,主内存的大部分访问发生,因此优化LUT的内存分配和聚合器的内存访问设计是提高AMU资源利用率和吞吐量的关键。

通过这些探索,研究者们展示了AMU如何克服逐元素运算带来的瓶颈,实现与问题规模无关的计算开销,以及如何通过硬件设计的优化达到资源效率和准确性的平衡。这些努力使AMU能够提供比现有解决方案更高的吞吐量和能效,尤其是在量化神经网络加速器上。

AMU-based QNN accelerator

  1. AMU模块的优化:通过在MADDNESS算法的基础上消除冗余运算,AMU模块实现了快速、高效、可扩展的近似矩阵乘法。AMU通过三种优化策略——输入输出数据剪枝、特征图重组和参数压缩,解耦了乘法计算开销与问题规模的关系,从而显著提高了MADDNESS算法的性能。
  2. 硬件设计空间的探索:研究者对FPGA上的AMU模块进行了硬件设计空间的探索,以优化不规则剪枝带来的非典型内存分配和访问。这种优化对提升AMU模块的性能至关重要。
  3. AMU模块的性能提升:实验结果证实,采用AMU的QNN加速器相比现有技术,在相同的神经网络模型复杂度下,吞吐量提高了9.9倍,能效(GOPS/W)提高了112倍。
  4. AMU模块的神经网络推理加速器:AMU模块被用于构建神经网络推理加速器,该加速器处理二进制特征图,采用8位精度的LUT。AMUNet-2在资源成本相同的情况下,相对于LUTNet和MVAU分别提高了9.1倍和13.7倍的效率,同时精度损失不超过2.6%。
  5. 性能对比:为了验证AMU模块在神经网络推理加速中的性能,研究者在ZCU104和XCZU19EG平台上实现了基于AMU的神经网络加速器,用于SFC模型(四层多层感知器)和CNV模型(VGG启发的卷积神经网络)。通过将AMU模块替换到CNV模型的卷积层中,研究者展示了AMU加速器在遇到卷积操作时的性能。实验显示,AMU加速器在不同问题规模下均展现出卓越的峰值吞吐量和峰值效率,即使随着问题复杂度的增加,如从57.8 MOPs增加至109 MOPs,AMU加速器依然维持了高效率,同时资源使用率也在可接受范围内。
  6. 结论:论文的结论部分强调,AMU模块的提出为神经网络加速器设计开辟了新路径,尤其是在FPGA平台上,它能够实现高吞吐量和高能效,同时保持可接受的精度损失。通过优化MADDNESS算法和专门设计的内存管理,AMU加速器在处理复杂神经网络时表现出色,证明了其在深度学习领域的潜在应用价值。然而,AMU模块仍存在一定的精度损失,这源于其在近似矩阵乘法中采用的聚类策略,每个AMU仅关注自身输出,而忽略了与最终网络推理结果的联系。未来的研究方向可能是探索针对深层神经网络中常见操作的有效优化,以进一步提升性能。

AMU-2.png

AMU的创新点和改进主要体现在以下几个方面:

  1. 消除冗余运算:AMU优化了MADDNESS算法中层间的冗余运算,使得在FPGA上可以快速高效地实现近似矩阵乘法,同时保持计算的可扩展性。
  2. 资源效率提升:通过专门的内存管理及访问设计,AMU能够有效减少计算开销,不再受限于输入特征图的分辨率,从而显著提高了FPGA平台上的神经网络加速器的资源效率和能量效率。
  3. 性能提升:在评估中,AMU模块在Xilinx ZCU104平台上的实现,相比现有技术,在单个AMU模块的最佳吞吐量和能量效率权衡方面进行了探索。实验结果表明,AMU在量化神经网络加速器中,与同类解决方案相比,能够达到最高9倍的吞吐量和112倍的能量效率提升。
  4. 优化策略影响:论文还评估了AMU基神经网络加速器在不同优化策略下的性能,以及与先前工作的性能对比,验证了AMU方法的有效性和优越性。
  5. 硬件设置优化:AMU的设计还探索了FPGA上的硬件设计空间,优化了非结构化剪枝引起的不典型内存分配和访问问题,进一步提高了AMU的性能。

综上所述,AMU方法的创新性和改进主要体现在其对MADDNESS算法的优化,通过消除冗余运算、提高资源效率和优化硬件设计,实现了矩阵乘法的快速、可扩展和高能效计算,尤其是在FPGA平台上,为神经网络加速器的性能和能效提升做出了显著贡献。

AMU-3.JPG

最新文章

最新文章