通过张量压缩优化在FPGA上实现超内存高效的Transformer训练
机器学习
2025-08-07 v2 硬件体系结构
计算与语言
摘要
Transformer模型在广泛的机器学习任务中实现了最先进的性能。由于隐私、领域自适应和设备端科学机器学习等考量，在资源受限的边缘设备上训练Transformer的兴趣日益增长。然而，Transformer训练所需的巨大计算和内存需求往往超出了边缘设备的能力。利用低秩张量压缩，本文提出了首个用于端到端Transformer训练的片上FPGA加速器。在算法方面，我们提出了一种用于张量化Transformer训练的双向收缩流，与现有的张量操作相比，显著降低了计算FLOPS和层内内存成本。在硬件方面，我们将所有高度压缩的模型参数和梯度信息存储在芯片上，为训练的每个阶段创建了一个仅依赖片上内存的框架。这减少了片外通信，并最小化了延迟和能耗成本。此外，我们为每个训练阶段实现了定制计算内核，并采用层内并行和流水线技术以进一步提高运行时和内存效率。通过在ATIS数据集上使用FP-32数据格式对大小在至 MB范围内的Transformer模型进行实验，我们的张量化FPGA加速器能够在AMD Alevo U50 FPGA上进行单批次端到端训练，其内存预算不足-MB BRAM和-MB URAM。与在NVIDIA RTX 3090 GPU上进行未压缩训练相比，我们的片上FPGA训练实现了至的内存缩减。与在NVIDIA RTX 3090 GPU上进行张量化Transformer训练相比，我们的FPGA加速器在每个epoch的能耗成本上也降低了高达。
引用
@article{arxiv.2501.06663,
title = {Ultra Memory-Efficient On-FPGA Training of Transformers via Tensor-Compressed Optimization},
author = {Jiayi Tian and Jinming Lu and Hai Li and Xiangwei Wang and Cong Hao and Ian Young and Zheng Zhang},
journal= {arXiv preprint arXiv:2501.06663},
year = {2025}
}