SpinTran: A Spintronic Transformer Accelerator Using Full-Precision Computing in Memory With Approximate Softmax Functions

Transformer models have made remarkable contributions to various artificial intelligence (AI) applications. However, high energy consumption and delay caused by matrix multiplication (MM) and softmax are bottleneck for hardware to implement transformer models. Computing in memory (CIM) is a promising architecture to solve these problems by reducing data movement and enabling highly parallel data processing. In this work, we propose a software-hardware co-design SpinTran accelerator using the full-precision spintronic CIM and approximate softmax functions for performing energy-efficient transformer models. First, a full-precision CIM scheme (FPCIM) based on toggle spin torque magnetic random-access memory (TST-MRAM) is proposed to realize fast MM operation. Second, we build a highly-reconfigurable array by utilizing the characteristic of TST-MRAM, which can write multiple bits into a column simultaneously to significantly reduce write-back cycles for dynamic matrix in transformer models. Third, we design an efficient Taylor expansion method that adopts approximation and range reduction techniques to achieve rapid convergence, thereby improving the computational efficiency of softmax function. Results show that the delay and energy efficiency of the proposed SpinTran accelerator are 3.5 ns and 56.7 TOPS/W in INT8 MAC operation, respectively. Meanwhile, the energy consumption is 4.71 μJ/token, 12.82 μJ/token, and 4.32 μJ/token when executing Bert-base, Bert-large, and ViT-small models, respectively, which achieves 2.81 to 4.30× energy reduction compared to other state-of-the-art works.