PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, Mehrdad Farajtabar

arXiv:2312.11514v3cs.CLcs.AIcs.LG

TL;DR

LLM inference can exceed available DRAM, limiting deployment on memory-constrained devices. This paper stores parameters in flash, loads needed weights on demand using hardware-aware transfer optimizations, and reports faster inference for models larger than DRAM. The study remains limited in deployment scope, including single-batch inference and unexamined power and thermal constraints.

  • Problem

    LLMs’ substantial computational and memory requirements make inference difficult on devices with limited DRAM, where full-model loading restricts runnable model size.

  • Method

    The paper stores model parameters in flash and selectively loads them into DRAM using windowing, row-column bundling, and hardware-aware transfer optimization.

  • Results

    Models up to twice the available DRAM size run with 4-5x faster CPU and 20-25x faster GPU inference than traditional loading methods.

  • Takeaways & Limitations

    The approach supports LLM deployment in resource-limited environments and expands their applicability and accessibility.

  • Takeaways & Limitations

    The study is limited to single-batch inference, with prompt processing and multi-batch inference left for future investigation.

Abstract

from arXiv · show

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.

1 Introduction

LLMs’ capabilities come with memory demands that make inference difficult on personal devices. The paper proposes storing parameters in flash and selectively loading them to run models beyond DRAM capacity.

  • Motivation: LLMs can contain hundreds of billions or trillions of parameters, creating substantial computational and memory requirements for inference.These requirements make efficient loading and execution challenging, especially on personal devices.
  • Motivation: A standard DRAM-resident approach limits runnable model size; a 7-billion-parameter half-precision model requires over 14GB just for parameters.Quantization reduces model size but does not remove the need to load the entire model into DRAM.
  • Approach: The paper stores model parameters in flash memory and loads only the required subset during inference.Flash memory is described as at least an order of magnitude larger than DRAM.
  • Approach: The method studies flash and DRAM constraints, then optimizes data-transfer volume, transfer throughput, and DRAM parameter management.The hardware analysis motivates techniques for serving LLMs from flash efficiently.
  • Results: Models up to 2x the device’s DRAM capacity run with inference speedups up to 4x in CPU, 7x in Metal, and 20x in NVIDIA GPU backends versus naive implementation.These results are reported as part of the paper’s contribution summary.

2 Flash Memory & LLM Inference

Flash provides substantially more capacity than DRAM but has lower bandwidth and poor efficiency for small random reads. The paper therefore targets selective transfers and larger, more contiguous reads to improve LLM inference.

  • Hardware constraints: Flash memory offers high capacity but lower latency and throughput than DRAM, while naively reloading an entire model can take seconds.Such transfers also consume more energy than moving data from DRAM to CPU or GPU memory.
  • Inference challenge: Selective reading of model weights uses activation sparsity to reduce response latency and avoid full-model loading.The motivation includes both initial loading penalties and repeated partial-weight loading challenges.
  • Read behavior: Flash throughput improves with larger sequential chunks because small random reads incur latency before data transfer begins.The read path includes multiple phases involving the operating system, drivers, interrupts, and flash controller.
  • Design response: The paper advocates jointly increasing read-chunk size and reducing transferred data to improve inference speed.Coalescing FFN row and column reads is presented as a way to pay startup latency less often.

3 Load From Flash

The section reduces flash-inference latency by loading less data, reading larger chunks, and managing loaded parameters efficiently in DRAM. It leverages activation sparsity and context-dependent neuron reuse while bundling corresponding FFN weights for higher-throughput reads.

  • 3.1 Reducing Data Transfer: 97% FFN sparsity in OPT 6.7B motivates selectively loading non-sparse FFN segments while retaining attention weights in DRAM.Falcon 7B and fine-tuned Llama 2 variants also exhibit high FFN sparsity, at 95% and 90%, respectively.
  • 3.1 Reducing Data Transfer: A low-rank predictor identifies elements zeroed by ReLU so the system loads only predictor-indicated weights without adversely affecting zero-shot performance.The predictor uses current-layer attention outputs and is trained with balanced positive and negative samples.
  • 3.1 Reducing Data Transfer: The sliding window retains predicted active neuron rows for recent tokens and incrementally loads only differences between the current token and predecessors.For a window of k tokens, the incremental load is s_agg(k + 1) − s_agg(k); decreasing aggregated usage makes larger windows reduce per-token loading.
  • 3.2 Increasing Transfer Throughput: Bundling each up-projection column with the corresponding down-projection row doubles chunks from d_model×num_bytes to 2d_model×num_bytes and increases throughput.The pairing is valid because both components correspond to the same intermediate neuron activation.
  • 3.3 Optimized Data Management in DRAM: Preallocating DRAM and replacing deleted elements with final elements before appending new weights reduces data movement when loaded FFN matrices change.This addresses the overhead of rewriting existing neuron data during reallocations.

4 Experiments and Results

Experiments evaluate flash-backed inference across models, datasets, and hardware setups, showing that predictor, windowing, and bundling improve loading efficiency and latency under constrained memory.

  • Experimental setup: The evaluation uses OPT 6.7B and sparsified Falcon 7B as primary models, with additional results for Phi-2, Persimmon 8B, and sparsified Llama 2.Latency measurements use 128-token C4 prompts followed by 256 generated tokens.
  • Experimental setup: Experiments span Apple M1 Max and M2 Ultra systems plus a Linux machine with an NVIDIA RTX 4090, using CPU and GPU execution.Approximately half of total DRAM and GPU memory is allocated for model computations.
  • Baselines: The primary baseline loads the required model subset from flash during each forward pass, with theoretical I/O latency used for fair comparison.A secondary hybrid baseline keeps half the model in memory and loads the remainder for every generated token without sparsity.
  • I/O efficiency: 1.25 GiB/s sparse versus 6.1 GiB/s dense reads shows the throughput cost of scattered access, while bundling mitigates it by combining projection weights.Windowing and low-rank prediction reduce transferred data, so the required sparse subset still loads faster and uses less DRAM overall.
  • End-to-end results: The efficient implementation improves loading efficiency over naive and hybrid approaches across all evaluated models, with further GPU improvement from speculative decoding.Table 3 evaluates end-to-end inference latency across different setups.

5 Ablation analysis

Ablations examine generation length, sampling, window size, speculative decoding, and energy use, showing robustness across longer generations but a total-energy trade-off for sparse inference.

  • Window size: More parameters retained in DRAM reduce GPU latency, establishing a memory-use versus latency trade-off.Figure 7 reports this relationship for OPT-6.7B.
  • Generation length: Generating 1000 tokens for OPT 6.7B on GPU does not cause SSD thermal throttling or lower performance.Average flash latency does not increase later in generation; the first few tokens have higher latency because DRAM must be filled.
  • Sampling: Nucleus sampling does not lower performance during long generations on either CPU or GPU.The authors use this result to assess whether more diverse activations disadvantage the method.
  • Speculative decoding: 1.4x decoding speedup is achieved with speculative decoding on OPT 6.7B using λ = 4.The reported speedup is close to the original 1.58x speculative-decoding speedup.
  • Power consumption: Sparse inference uses less power per unit time than a similarly sized dense model but consumes more total energy because token generation takes longer.A systematic quantitative evaluation of the exact power-use pattern remains future work.

6 Related Works

Related work includes model compression, selective execution, and weight offloading, whereas this paper targets cases where the full model cannot fit in device DRAM or GPU memory.

  • Research categories: Existing approaches broadly include model compression techniques such as pruning and quantization, and selective execution methods.These categories frame prior efforts to reduce inference requirements as LLMs grow.
  • Selective weight loading: DejaVu exploits activation sparsity to load subsets of weights for each layer but still requires loading from GPU memory.This differs from the flash-centered setting studied here.
  • Offloading: FlexGen offloads weights and KV-cache across GPU memory, DRAM, and flash, but remains theoretically bounded by flash-to-DRAM throughput when memory is insufficient.The paper considers cases where the full model cannot reside in aggregate device memory.
  • Problem setting: Unlike the literature’s usual assumption that the model fits fully in GPU memory or system DRAM, this work studies efficient parameter storage and loading from flash.The target setting is resource-limited personal devices.

7 Discussion

The paper presents hardware-aware flash-backed inference with windowing and row-column bundling, enabling larger-than-DRAM models and faster loading while identifying substantial future work.

  • Contribution: The method combines an inference cost model aligned with flash and DRAM characteristics with windowing and row-column bundling.These techniques target inference on devices with constrained memory capacities.
  • Results: LLMs up to twice the available DRAM size can run, with OPT inference accelerated 4-5x on CPU and 20-25x on GPU versus traditional loading.These are the paper’s reported practical outcomes for resource-limited deployment.
  • Limitations and future work: The work is described as a first step, with further opportunities in optimized weight bundling, data structures, and hardware-specific inference stacks.The authors identify both algorithmic and engineering directions for future work.

8 Limitations

The study identifies scope limits and open questions involving deployment conditions, memory availability, power and thermal behavior, and broader model architectures.

  • On-device power consumption and thermal limitations remain important areas for future analysis.
  • The evaluation is currently limited to single-batch inference, leaving prompt processing and multi-batch scenarios for future investigation.
  • The initial proof of concept assumes DRAM availability equal to half the model size; other memory budgets remain unexplored.
  • The approach is constructed on sparsified networks, although the authors suggest adapting selective loading to non-sparse networks and contextual retrieval.
  • Appendix D reports a negative result for bundling neurons by co-activation as a strategy for increasing chunk size.

B Low-Rank Activation Predictor: Additional Results

Additional results examine predictor sparsity, accuracy, efficiency, implementation conditions, and caching considerations across several models and hardware settings.

  • Sparsity patterns of predictors: False negatives constitute a small fraction of predictor outputs, but reducing them requires over-predicting redundant neurons.
  • Accuracy: Zero-shot accuracy does not drop with predictors, while larger predictors in selected later layers improve some Persimmon and Falcon metrics.
  • Efficiency: Predictors add limited overhead: OPT-6.7B averages less than 2.4% of nonembedding weights and FLOPs, with 2.75% CPU and 4.8% GPU inference-time shares.
  • Experimental setup: The experiments process one sequence at a time and allocate DRAM for the KV cache while focusing primarily on model size.
  • Caching considerations: Benchmarks disable caching to provide a conservative lower bound, although caching may improve throughput in practical systems.

C.1 Results for OPT 6.7B Model

Results across OPT, Falcon, Persimmon, and Phi-2 evaluate memory allocation, predictor behavior, latency, and accuracy under constrained DRAM.

  • OPT 6.7B: 52.1% of OPT-6.7B model size is retained in DRAM, including embeddings, attention weights, predictors, and the loaded FFN portion.
  • OPT 6.7B: Less than 162ms per token of memory-related latency compares with approximately 2196ms for the baseline loading 13.4GB per token.
  • Falcon 7B: 52.93% of Falcon 7B model size is allocated in DRAM with a four-token window and active FFN portion.
  • Falcon 7B: 250ms per token on an M1 Max is approximately 9 to 10 times faster than the Falcon baseline latency of around 2196 milliseconds.
  • Phi-2 and comparative results: Phi-2 achieves 2.35x speedup over the naive baseline, while LLM in Flash achieves 3x speedup in another reported setup.
  • Accuracy: Using predictors on sparse models did not hurt MMLU results; the reported scores are 41.8 for Llama 2, 38.96 for a sparsified model, and 38.63 after predictor training.

E Extended Related Works

The paper situates its flash-based LLM inference method among selective loading, speculative execution, hardware optimization, and sparsity-oriented approaches. It combines cost modeling, sparsity prediction, and hardware awareness to reduce flash weight loading and achieve substantial speedups.

  • Selective Weight Loading: Selective weight loading exploits activation sparsity, but prior approaches differ in whether weights remain in GPU memory or are offloaded across memory tiers.Dejavu loads subsets of weights from GPU memory, while Flexgen offloads weights and KV-cache across GPU memory, DRAM, and flash.
  • Sparsity Prediction: The closest friend of each neuron in OPT 6.7B almost always coactivates, while the third closest friend coactivates 86% of the time on average.
  • Speculative Execution: The method uses lightweight speculation tailored to adaptive weight loading to hide flash latency.
  • Hardware Optimizations: Hardware optimization research spans memory architectures, dataflows, evaluation frameworks, sparse kernels, and flash, while this paper focuses on algorithmic improvements.The paper notes that these hardware techniques could provide additional speedups.
  • Speculative Execution: Speculative decoding is orthogonal to the method and can further improve performance by updating its window with multiple tokens rather than one.
  • Mixture of Experts: Mixture-of-Experts feed-forward sparsity can leverage the method to enable larger models on-device.
  • Summary: 4-5x and 20-25x speedups are demonstrated on CPU and GPU, respectively, by combining cost modeling, sparsity prediction, and hardware awareness.
  • Sparsity Prediction: Quantized OPT 6.7B is evaluated against the original model through active-neuron percentages across different layers and 100 sequences.

G Qualitative Evaluations

The qualitative evaluation compares outputs from the original model with outputs generated using predictors. Examples use both a fixed completion prompt and random samples from the C4 dataset.

  • Evaluation Purpose: The evaluation qualitatively examines whether predictor-assisted model outputs remain comparatively reasonable.
  • Fixed Prompt: Tables 7 and 8 compare original-model outputs with predictor-assisted outputs for completions beginning with “once upon a time there was a”.
  • C4 Prompt: Table 9 compares original-model and predictor-assisted outputs after prompting with a random sample from the C4 dataset.
Loading 2312.11514v3…