通过张量压缩优化在FPGA上实现超高效内存利用的Transformer训练

《IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems》：Ultra Memory-Efficient On-FPGA Training of Transformers via Tensor-Compressed Optimization

Transformer模型训练在FPGA加速器上的优化与实现，通过低秩张量压缩和双向收缩流算法降低计算与内存需求，片上内存存储减少通信延迟，自定义计算内核和流水线技术提升能效，在AMD FPGA上实现6MB内存预算的端到端训练，内存减少30-51倍，能耗降低4倍。

　　

摘要：

Transformer模型在广泛的机器学习任务中取得了先进的性能。由于隐私保护、领域适应以及设备上的科学机器学习等方面的考虑，人们在资源受限的边缘设备上训练Transformer的兴趣日益增加。然而，Transformer训练所需的巨大计算和内存需求往往超出了边缘设备的处理能力。本文利用低秩张量压缩技术，提出了首个用于Transformer训练的FPGA加速器。在算法层面，我们提出了一种双向收缩流方法，用于张量化的Transformer训练，与现有的张量操作相比，显著降低了计算FLOPS和层内内存成本。在硬件层面，我们将所有高度压缩的模型参数和梯度信息存储在芯片上，为训练的每个阶段创建了一个仅依赖片上内存的框架。这减少了芯片间通信，从而最小化了延迟和能耗。此外，我们还为每个训练阶段实现了定制的计算内核，并利用层内并行性和流水线技术进一步提高运行时效率和内存利用率。通过对ATIS数据集上大小在36.7 MB到93.5 MB之间的Transformer模型进行实验（使用FP-32数据格式），我们的FPGA加速器能够在AMD Alevo U50 FPGA上进行单批次端到端训练，其内存预算低于6 MB的BRAM和22.5 MB的URAM。与在NVIDIA RTX 3090 GPU上进行未压缩训练相比，我们的FPGA训练实现了高达30倍的内存节省。同时，与在NVIDIA RTX 3090 GPU上进行张量Transformer训练相比，我们的FPGA加速器每个训练阶段的能耗降低了4.0倍。作为初步成果，本文突显了在边缘设备上进行大规模张量训练的巨大潜力。