Ultra Memory-Efficient On-FPGA Training of Transformers via Tensor-Compressed Optimization (2501.06663v1)

Abstract: Transformer models have achieved state-of-the-art performance across a wide range of machine learning tasks. There is growing interest in training transformers on resource-constrained edge devices due to considerations such as privacy, domain adaptation, and on-device scientific machine learning. However, the significant computational and memory demands required for transformer training often exceed the capabilities of an edge device. Leveraging low-rank tensor compression, this paper presents the first on-FPGA accelerator for end-to-end transformer training. On the algorithm side, we present a bi-directional contraction flow for tensorized transformer training, significantly reducing the computational FLOPS and intra-layer memory costs compared to existing tensor operations. On the hardware side, we store all highly compressed model parameters and gradient information on chip, creating an on-chip-memory-only framework for each stage in training. This reduces off-chip communication and minimizes latency and energy costs. Additionally, we implement custom computing kernels for each training stage and employ intra-layer parallelism and pipe-lining to further enhance run-time and memory efficiency. Through experiments on transformer models within $36.7$ to $93.5$ MB using FP-32 data formats on the ATIS dataset, our tensorized FPGA accelerator could conduct single-batch end-to-end training on the AMD Alevo U50 FPGA, with a memory budget of less than $6$-MB BRAM and $22.5$-MB URAM. Compared to uncompressed training on the NVIDIA RTX 3090 GPU, our on-FPGA training achieves a memory reduction of $30\times$ to $51\times$. Our FPGA accelerator also achieves up to $3.6\times$ less energy cost per epoch compared with tensor Transformer training on an NVIDIA RTX 3090 GPU.

• The paper presents a method for on-FPGA training of transformer models using tensor-compressed optimization, significantly reducing memory and power requirements for resource-constrained hardware.
• Key techniques include a novel bi-directional contraction flow, an on-chip memory-only framework, and custom computing kernels optimized for FP-32 data formats.
• Experimental results demonstrate memory footprint reduction factors of 30x to 51x and energy cost reductions up to 3.6x compared to uncompressed GPU training on specific datasets.

Ultra Memory-Efficient On-FPGA Training of Transformers via Tensor-Compressed Optimization

The paper presents an efficient approach for training transformer models on FPGAs, addressing the substantial computational and memory resources typically required. The researchers exploit tensor-compressed optimization techniques to enable end-to-end transformer training on an FPGA, demonstrating significant memory savings while maintaining model performance, particularly in scenarios constrained by hardware resources. By leveraging tensor decomposition formats, specifically tensor-train (TT) and tensor-train-matrix (TTM), this paper paves the way for more accessible deployment of transformer models in edge environments without compromising accuracy.

Transformer models have driven advances across various domains due to their high performance, yet this success is accompanied by substantial resource demands, especially when training these models. This work acknowledges the role of edge devices in areas such as privacy-sensitive applications and real-time processing environments, where the model must adapt rapidly and operate efficiently. The paper introduces the first on-FPGA tensor-compressed transformer training system, presenting a comprehensive hardware and algorithm design to mitigate the typical memory and bandwidth limitations of FPGAs.

Key Contributions

• Bi-Directional Contraction Flow: The researchers propose a novel bi-directional contraction flow for tensorized transformer training. This significantly reduces the FLOPs and intra-layer memory costs seen with traditional tensor operations, making it much more efficient than existing methods.
• On-Chip Memory-Only Framework: The introduced framework ensures all highly compressed model parameters and gradient information remain on-chip during each stage of training. This design choice reduces the need for off-chip communications, thereby minimizing latency and energy consumption.
• Custom Computing Kernels and Pipeline Optimization: The hardware design incorporates custom computing kernels and employs intra-layer parallelism alongside fine-grained pipelining strategies. These techniques are crafted to improve run-time efficiency and resource utilization within the FPGA.

Experimental Results

The implementation of this tensorized FPGA accelerator on the AMD Alveo U50 FPGA was studied using the ATIS dataset, handling transformers with overall model sizes ranging from 36.7 MB to 93.5 MB using FP-32 data formats. Experimental results on models with two to six encoder blocks, indicated that this approach could vastly compact the memory footprint - down to less than 6 MB of BRAM and 22.5 MB of URAM. The findings exhibit a compelling memory reduction by factors of 30 to 51, compared to uncompressed training on NVIDIA RTX 3090 GPUs. Additionally, the energy costs per training epoch saw reductions up to 3.6 times over conventional GPU-based transformer training.

Theoretical and Practical Implications

The work significantly extends the frontier of transformer model training, specifically within edge computing environments. From a theoretical perspective, it advances the understanding of high-rank tensor compression and its effective integration within neural network training pipelines. Practically, this development holds promise for enabling robust machine learning capabilities on edge devices, facilitating privacy-preserving model deployment in real-world applications. The reduced memory and power footprint demonstrated here further highlights the utility of integrating tensorized training methods with specialized hardware, broadening the deployment scenarios for neural networks.

Future Directions

While the FPGA accelerator exhibits notable improvements, there remain several aspects for further exploration. Optimizing GPU implementations of tensor-compressed training, potentially by developing specialized CUDA kernels, could yield additional insights. Moreover, integrating low-precision data formats and exploring hybrid computation architectures could further enhance the scalability and applicability of this approach. As edge devices continue to evolve, aligning tensor-compressed optimization with these technological advances presents a fertile area for future research.

In summary, this paper marks a significant stride towards making high-performance transformer models viable on constrained hardware, wherein memory efficiency and energy conservation are paramount. It offers a solid foundation for future efforts aimed at enhancing the training and deployment of complex machine learning models beyond traditional, resource-intensive environments.