LLM in a flash: Efficient Large Language Model Inference with Limited Memory (2024)

Abstract: Large language models (LLMs) are central to modern natural language processing, delivering exceptional performance in various tasks. However, their substantial computational and memory requirements present challenges, especially for devices with limited DRAM capacity. This paper tackles the challenge of efficiently running LLMs that exceed the available DRAM capacity by storing the model parameters in flash memory, but bringing them on demand to DRAM. Our method involves constructing an inference cost model that takes into account the characteristics of flash memory, guiding us to optimize in two critical areas: reducing the volume of data transferred from flash and reading data in larger, more contiguous chunks. Within this hardware-informed framework, we introduce two principal techniques. First, “windowing” strategically reduces data transfer by reusing previously activated neurons, and second, “row-column bundling”, tailored to the sequential data access strengths of flash memory, increases the size of data chunks read from flash memory. These methods collectively enable running models up to twice the size of the available DRAM, with a 4-5x and 20-25x increase in inference speed compared to naive loading approaches in CPU and GPU, respectively. Our integration of sparsity awareness, context-adaptive loading, and a hardware-oriented design paves the way for effective inference of LLMs on devices with limited memory.

## introduction in recent years, large language models (llms) have demonstrated strong performance across a wide range of natural language tasks @xcite @xcite @xcite . however, the unprecedented capabilities of these models come with substantial computational and memory requirements for inference. llms can contain hundreds of billions or even trillions of parameters, which makes them challenging to load and run efficiently, especially on personal devices. currently, the standard approach is to load the entire model into dram (dynamic random access memory) for inference @xcite . however, this severely limits the maximum model size that can be run. for example, a 7 billion parameter model requires over 14gb of memory just to load the parameters in half-precision floating point format, exceeding the capabilities of most personal devices such as smartphones. while it is possible to employ techniques such as quantization to reduce the model size, still, this cannot address the main limitation of loading the entire model into dram. to address this limitation, we propose to store the model parameters in flash memory, which is at least an order of magnitude larger than dram. then, during inference, we directly load the required subset of parameters from the flash memory, avoiding the need to fit the entire model in dram. to this end, our work makes several contributions: ## flash memory & llm inference in this section, we explore the characteristics of memory storage systems (e.g., flash, dram), and their implications for large language model (llm) inference. we aim to understand the challenges and hardware-specific considerations essential for algorithm design, particularly in optimizing inference when working with flash memory. ## load from flash this section addresses the challenge of conducting inference on devices where the available dram is substantially smaller than the size of the model. this necessitates storing the full model weights in flash memory. our primary metric for evaluating various flash loading strategies is latency, dissected into three distinct components: the i/o cost of loading from flash, the overhead of managing memory with newly loaded data, and the compute cost for inference operations. our proposed solutions for reducing latency under memory constraints are categorized into areas: 1. reducing data load: aiming to decrease latency associated with flash i/o operations by loading less data foot_0 . streamlining the management of data once it is loaded into memory to minimize overhead. it is important to note that our focus is not on optimizing the compute, as it is orthogonal to the core concerns of our work. instead, we concentrate on optimizing flash memory interactions and memory management to achieve efficient inference on memory-constrained devices. we will elaborate on the implementation details of these strategies in the experimental setup section. ## experiments and results we start this section by briefly discussing our experimental setup and implementation details. next, we show that the techniques introduced in section 3 can improve the inference latency significantly across different models and runtime platforms. we postpone the some details to the appendix sections as follows: performance of our trained low-rank predictor (appendix b). ## related works as llms grow in size, reducing their computational and memory requirements for inference has become an active area of research. approaches broadly fall into two categories: model compression techniques such as pruning and quantization @xcite @xcite @xcite @xcite @xcite @xcite , and selective execution such as sparse activation @xcite , or conditional computation @xcite . our work is orthogonal to these directions, focusing mainly on minimizing data transfer from flash memory during inference. perhaps most related to our work is the literature on selective weight loading. sparsegpu @xcite exploits activation sparsity to load a subset of weights for each layer. however, it still requires loading from ram. flexgen @xcite offloads the weights and kv-cache from gpu memory to dram and dram to flash memory. in contrast, we consider only the cases where the full model can't reside in the whole dram and gpu memory on the edge devices. notably, flexgen is still theoretically bound by the slow throughput of flash to dram in such scenarios. an expanded discusion of related works is deferred to appendix e. overall, the primary assumption in the literature is that the model can fully reside in the gpu memory or system dram. however, considering the limited resources available on personal devices, we do not share this assumption in this work. instead, we concentrate on exploring how to store and load parameters on flash memory more efficiently, aiming to enhance inference efficiency. ## discussion in this study, we have tackled the significant challenge of running large language models (llms) on devices with constrained memory capacities. our approach, deeply rooted in the understanding of flash memory and dram characteristics, represents a novel convergence of hardware-aware strategies and machine learning. by developing an inference cost model that aligns with these hardware constraints, we have introduced two new techniques: 'windowing' and 'row-column bundling'. the practical outcomes of our research are noteworthy. we have demonstrated the ability to run llms up to twice the size of available dram. for example, on opt model, we achieve an acceleration in inference speed of 4-5x compared to traditional loading methods in cpu, and 20-25x in gpu. this is particularly crucial for deploying llms in resource-limited environments, thereby expanding their applicability and accessibility. while in this work we have studied the previously unexplored problem of serving llms from flash, we note that this work is only a first step in this direction, and has several limitations that we discuss in the next section, and we believe there are several interesting problems left to be explored in future works. for instance, from the algorithmic perspective, more optimized techniques of weight bundling and data structures can be crafted, while from the engineering perspective, the characteristics of specific hardware platforms can inform works on building more efficient inference stacks. ## limitations our study represents an initial endeavor in the pursuit of democratizing large language model (llm) inference, making it accessible to a wider array of individuals and devices. we recognize that this early effort has its limitations, which, in turn, open up compelling avenues for future research. a critical aspect for future exploration is the systematic analysis of power consumption and thermal limitations inherent in the methods we propose, particularly for on-device deployment. currently, our study is limited to single-batch inference. we provide some preliminary results on combining our proposed idea with speculative decoding, however, expanding this to include more complex scenarios like prompt processing and multi-batch inference are valuable areas for further investigation. in our initial proof of concept, we operated under the assumption of memory availability being half the size of the model. exploring the dynamics of working with varying memory sizes-both larger and smaller-introduces a fascinating balance between latency and accuracy, and is a compelling area for future exploration. in conclusion, our methodology is constructed on the foundation of sparsified networks. nonetheless, the underlying concept holds potential for broader applications. it can be adapted to selectively load weights in non-sparse networks or to dynamically retrieve model weights from flash storage. this adaptation would be contingent on the specific requirements of the input prompt or the contextual parameters provided. such an approach suggests a versatile strategy for managing model weights, and optimizing performance based on the nature of the input, thereby enhancing the efficiency, usefulness, and applicability of the proposed scheme in various scenarios dealing with large language models (llms). quantization-aware training for large language models. in advances in neural information processing systems, volume 35.