Running Mistral 7B with Core ML

Hugging Face has detailed a workflow for running the Mistral 7B model on Mac hardware by leveraging new Core ML features introduced at WWDC 24. By combining stateful buffers, a new Swift Tensor type, and advanced quantization, a 7-billion parameter model can be run using less than 4GB of memory.

New Core ML Features from WWDC 24

Apple's latest updates to Core ML provide several critical optimizations for on-device large language models (LLMs).

Swift Tensor (MLTensor)

The new MLTensor type provides a high-level abstraction for multi-dimensional data structures in Swift, mimicking the functionality of numpy arrays or torch tensors. This replaces the need for low-level manipulations of MLMultiArray or MLShapedArray, which previously required accessing underlying storage as opaque pointers. MLTensor includes built-in operations like softmax, which simplifies the pre- and post-processing pipelines for language models.

Stateful Buffers

Traditionally, Core ML models operated as stateless functions. Stateful buffers allow a model to reserve a block of memory on the GPU to maintain state across iterations. For LLMs, this is used to implement key-value (KV) caching. By keeping the KV-cache on the GPU, the system avoids the overhead of sending and receiving large stateful data between the CPU and GPU during every token generation step, significantly increasing performance by reducing memory bandwidth bottlenecks.

Block-wise Quantization

To reduce model size without drastic quality loss, Core ML now supports block-wise quantization. This technique creates multiple lookup tables (LUTs) for different areas of the same tensor rather than a single table for the entire tensor. This allows models to be compressed to 4-bit precision, resulting in an 8-fold reduction in size compared to float32 and a 4-fold reduction compared to float16.

Multifunction Support

Multifunction support allows developers to package LoRA (Low-Rank Adaptation) adapters into generative models. This enables a single base model to be used for various tasks or styles by swapping small sets of additional parameters (adapters) without needing to reload the entire model.

Converting Mistral 7B to Core ML

Running Mistral 7B efficiently requires a custom attention implementation that pre-allocates the full KV-cache buffer and updates it in place to satisfy Core ML's stateful buffer requirements.

Tracing and Conversion Process

  1. Tracing: The model is loaded using a patched implementation of StatefulMistralForCausalLM and traced with example inputs using torch.jit.trace.
  2. Input Definition: Range dimensions are used via coremltools to allow input sequences to grow from a single token up to a maximum context length of 2048.
  3. State Preparation: The keyCache and valueCache buffers are defined using ct.StateType to ensure they are converted into stateful Core ML buffers.
  4. Conversion: The model is converted to Core ML with a minimum deployment target of iOS 18 or macOS 15 to access the new stateful features.

Model Compression

Using OpLinearQuantizerConfig, the model is compressed using 4-bit linear symmetric quantization with a block size of 32. This reduces the final mlpackage size from approximately 14GB (at float16) to about 3.8GB.

Execution and Implementation

Users can run the converted Mistral 7B model using the preview branch of the swift-transformers repository.

Running with Swift

The preview branch incorporates full MLTensor support and the Swift counterpart of the Stateful API. Inference can be executed via the command line:

git clone -b preview https://github.com/huggingface/swift-transformers
swift run transformers "Prompt text" --max-length 128 Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage

Running with Python

Inference is also supported via Python using coremltools:

python3 generate.py Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage --prompt "Prompt text"

Future Roadmap

Hugging Face plans to integrate these experimental methods into exporters, a Python tool for converting transformers models to Core ML. Additionally, they aim to optimize performance for the Apple Neural Engine (ANE) on iPhones by exploring smaller models like OpenELM or DCLM, as the current Mistral 7B implementation is primarily optimized for Mac GPUs.

Sources