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

A 7-billion-parameter model at 16 bits is about 14 GB of weights. Phones and many laptops cannot hold that in DRAM, but they do have far more flash storage. Reading the whole model from flash for each generated token is the naive fix, and it takes two seconds or more per token, because flash bandwidth is an order of magnitude below DRAM's and much worse still for small, scattered reads.

Apple's insight is that most of the model is not needed for any one token. In models with ReLU activations, the feed-forward layers' hidden neurons are overwhelmingly zero: about 97% for OPT-6.7B. If you knew in advance which neurons a token would activate, you would only need those neurons' weights. Feed-forward layers hold about two-thirds of a model's parameters, so that is most of the model.

The paper builds three techniques on top of that. A low-rank predictor per layer takes the attention output and guesses which feed-forward neurons will fire, with negligible effect on accuracy. A sliding window keeps in DRAM the neurons used by the last few tokens, since consecutive tokens activate overlapping sets, so each new token only loads neurons it has not needed recently. Row-column bundling stores each neuron's up-projection column next to its down-projection row in flash, so one read fetches both and reads are twice as large, which flash handles far more efficiently. Attention layers and embeddings stay in DRAM.

With DRAM for only half of OPT-6.7B, flash reads fell from 13.4 GB to about 0.2 GB per token, and latency fell from 3.2 seconds to 0.67 seconds on an M1 Max CPU and from 2.2 seconds to 84 milliseconds on a GPU.

Most deployment advice assumes the model fits in memory. This paper is about what to do when it does not, and it explains how on-device assistants can run models larger than RAM.

It needs sparse activations. The method works because ReLU-family models leave most neurons at zero. Most modern models use SwiGLU or GELU, which are not sparse in this way; the paper uses ReLU-converted versions of Falcon and Llama 2. Before promising this approach, check the customer's model architecture.

It is a single-request technique. Results are for one sequence at a time on one device; the authors list batched inference and prompt processing as future work. For server deployments, quantization and more GPUs are still the answer.

Storage layout is part of the design. Flash is fastest for large contiguous reads. Bundling the weights a token needs together roughly doubled read throughput. The broader lesson for edge deployments is that how weights are laid out on disk can matter as much as how many there are.