Apple Study Suggests Future iPhone and Mac Support for Large Language Models (LLM)

Apple is relatively late to the game when it comes to large language models (LLMs), but they have recently taken some interesting initiatives. The company's researchers have recently introduced a method that significantly reduces the memory required to run LLMs. This breakthrough relies on novel storage and memory management techniques that dynamically transfer the weights of the model between flash memory and DRAM while maintaining impressive low latency.

Their method allows large models to run at considerable throughput even when partially loaded onto device RAM. This could be important for Apple, a company that does not have a supercomputing service business, as it would enable significant benefits from AI on devices. This technology could allow future iPhones and Macs to leverage the capabilities of LLMs without overwhelming the system's memory.

Memory Requirements of LLMs

LLMs are notorious for their enormous memory requirements. For example, a 7 billion parameter model requires over 14GB of RAM in half precision. Such demands exceed the capacity of most edge devices.

One approach to address this issue is quantization, which compresses the model by converting parameters into 8-bit or 4-bit integers. However, quantization often requires modifications to the model and even complete retraining. In some cases, quantized LLMs are still too large to perform inference on devices.

Apple's research attempts to tackle the problem from a different angle: how can we deploy models on hardware with limited memory?

Loading LLMs from Flash Memory

One strategy to address the memory bottleneck is to store LLMs in flash memory and gradually load them into RAM for inference tasks. While flash memory on devices is more abundant than DRAM, it is at least an order of magnitude slower.

Naively performing inference using flash memory would require reloading the entire model for each forward pass. This process is slow, even for compressed models, often taking several seconds and consuming a significant amount of energy.

In their paper, Apple's researchers introduce a series of optimization techniques aimed at simplifying the process of loading models from flash memory to DRAM during inference while keeping memory consumption low.

Exploiting Model Sparsity

Language models rely on transformer blocks, which consist of two key elements: attention mechanisms and feed-forward (FFN) networks. Research has shown that the FFNs in LLMs exhibit high sparsity, with many values becoming zero or close to zero after activation, making them irrelevant for inference. Surprisingly, some FFN layers can have sparsity exceeding 95%.

Apple's researchers propose leveraging this sparsity to optimize model inference. By identifying and loading only the non-sparse elements during inference, memory costs can be significantly reduced. Their strategy involves fully loading smaller attention layers while selectively loading only the non-sparse parts of the FFNs into DRAM.

"Our approach involves iteratively transferring only the necessary non-sparse data from flash memory to DRAM for processing during inference," the researchers write.

To achieve this, they employ a "low-rank predictor" network to determine which parts of the FFNs will be non-sparse and reduce the amount of data that needs to be loaded. The active neurons identified by the predictor are then transferred to memory.

Sliding Window

Apple's researchers also design a "sliding window technique" to manage the loading and unloading of neurons during model inference. This approach keeps in memory only a small portion of the neurons identified by the recent input tokens and releases space for old tokens when new ones come in. Their research reveals that with each new token, only a small number of neurons need to be swapped, simplifying the data transfer from flash memory to RAM.

"This strategy allows for efficient memory utilization as it frees up memory resources previously allocated to old tokens that are no longer within the sliding window," the researchers write.

The size of the sliding window can be configured based on the memory limitations of the device running the model.

Optimizing Storage and Loading

Apple's team's research shows that neuron activations between different layers follow patterns that can be utilized for storage optimization. They suggest reorganizing the storage of model weights based on neurons that need to be loaded together. This arrangement allows the system to extract a continuous data block in a single read operation.

The researchers also investigate the benefits of storing co-activated neurons together. This approach would enable the system to extract data from a large flash memory segment. Although the results from this particular experiment did not meet expectations, these insights provide guidance for future research directions.

In addition, the team designs a novel data structure and dynamic memory allocation algorithm that reduces the frequency of memory allocation and allows the model to access approximately 25% of FFN neurons without significant performance degradation.

Accelerating LLMs on Low-Memory Devices

Apple's research team tested their techniques on two different models, OPT 6.7B and a sparse variant of Falcon 7B, and conducted experiments on two hardware setups. The first platform was an Apple M1 Max with 1TB SSD as flash memory. The second set of tests was performed on a Linux machine equipped with a 24GB NVIDIA GeForce RTX 4090 graphics card. In their experiments, they loaded only about 50% of the model weights at any given time.

On the M1 Max, loading the model from flash memory to RAM for each inference would result in a latency of 2.1 seconds per token. However, by implementing their novel techniques (sparsity prediction, windowing, and intelligent storage), the team reduced this latency to approximately 200 milliseconds. The improvements were even more significant on systems equipped with GPUs.

"We have demonstrated the ability to run LLMs that are twice the size of available DRAM, achieving 4-5x inference speedup on CPU and 20-25x on GPU compared to traditional loading methods," the researchers write.

The researchers believe that their findings could have important implications for future research and emphasize the importance of considering hardware characteristics in developing inference optimization algorithms.

"We believe that approaches like this work will be key to fully harnessing the potential of LLMs in a wide range of devices and applications as LLMs continue to grow in size and complexity," they write.