Enabling Energy-Efficient Deployment of Large Language Models on Memristor Crossbar: A Synergy of Large and Small

TL;DR

A novel architecture for the memristor crossbar is presented that enables the deployment of state-of-the-art LLM on a single chip or package, eliminating the energy and time inefficiencies associated with off-chip communication.

Abstract

Large language models (LLMs) have garnered substantial attention due to their promising applications in diverse domains. Nevertheless, the increasing size of LLMs comes with a significant surge in the computational requirements for training and deployment. Memristor crossbars have emerged as a promising solution, which demonstrated a small footprint and remarkably high energy efficiency in computer vision (CV) models. Memristors possess higher density compared to conventional memory technologies, making them highly suitable for effectively managing the extreme model size associated with LLMs. However, deploying LLMs on memristor crossbars faces three major challenges. Firstly, the size of LLMs increases rapidly, already surpassing the capabilities of state-of-the-art memristor chips. Secondly, LLMs often incorporate multi-head attention blocks, which involve non-weight stationary multiplications that traditional memristor crossbars cannot support. Third, while memristor crossbars excel at performing linear operations, they are not capable of executing complex nonlinear operations in LLM such as softmax and layer normalization. To address these challenges, we present a novel architecture for the memristor crossbar that enables the deployment of state-of-the-art LLM on a single chip or package, eliminating the energy and time inefficiencies associated with off-chip communication. Our testing on BERT_Large showed negligible accuracy loss. Compared to traditional memristor crossbars, our architecture achieves enhancements of up to 39X in area overhead and 18X in energy consumption. Compared to modern TPU/GPU systems, our architecture demonstrates at least a 68X reduction in the area-delay product and a significant 69% energy consumption reduction.

Figures (14)

• Figure 1: Two basic components that compose layers in LLM. For example, a layer in BERT consists of one component A and one component B.
• Figure 2: (a) A multi-bit memristor crossbar with $2\times 2$ cells ; (b) The area breakdown of a 128x128 memristor crossbar shafiee2016isaac, using a shared ADC.
• Figure 3: The sub-operation has two phases. The first phase involves the linear multiplication between X and Y, while the second phase involves the addition function $F$ applied to the multiplication result Z.
• Figure 4: We decompose the softmax into three parts, and two of these parts can be integrated with the preceding and subsequent matrix multiplication operations, resulting in three standardized sub-operations.
• Figure 5: Overview of our proposed architecture and the traditional architecture. We use unique colors to indicate the locations of weight matrices, denoted as $\textbf{W}_1$, $\textbf{W}_2$, and $\textbf{W}_3$. The dark arrows among the crossbars indicate the data flow direction.