Running Mistral 7B with Core ML

WWDC’ 24 is the moment Apple officially unveiled Apple Intelligence and
reiterated their commitment to efficient, private, and on-device AI.
During the keynote and the sessions that followed, they demonstrated
Apple Intelligence, which powers a huge array of AI-enhanced features
that show practical uses for everyday tasks. These are not
*AI-for-the-sake-of-AI* shiny demos. These are time-saving,
appropriate (and fun!) helpers that are deeply integrated with apps and
the OS, that also offer developers a number of ways to include these
features within their own apps.

Apple Intelligence features can only work this well
because of the vertically integrated software stack that harnesses
Apple Silicon’s capabilities to the fullest. Apple also offers a platform for developers to run models on-device, known as
Core ML. This software stack allows you to run ML models across all 3
compute units (CPU, GPU & Neural Engine) available on Apple Silicon hardware.

In this blog post, we’ll be exploring some of the best new Core ML
features to replicate the Mistral 7B example Apple showcased in the
WWDC’24 Deploy machine learning and AI models on-device with Core
ML
session, where they use a fork of
swift-transformers
to run a state-of-the-art LLM on a Mac. This is a high-quality model
with more than 7 billion parameters that pushes the capabilities of
consumer hardware today. You can also check out WWDC’24 Bring your
machine learning and AI models to Apple
silicon
session, where part of the Mistral 7B conversion process is shown.

Let’s see what steps to take to run it as efficiently as possible, and
learn the new tools available in iOS 18 & macOS Sequoia.

This is what we’ll be building today:

TL;DR

By the end of this blog post, you will have learnt all the new goodies
accompanying the latest macOS release AND you will have successfully run
a 7B parameter model using less than 4GB of memory on your Mac.

Step 1: Clone the preview branch of the swift-transformers repo: git clone -b preview https://github.com/huggingface/swift-transformers
Step 2: Download the converted Core ML models from this Hugging Face repo
Step 3: Run inference using Swift: swift run transformers "Best recommendations for a place to visit in Paris in August 2024:" --max-length 200 Mistral7B-CoreML/StatefulMistralInstructInt4.mlpackage

Best new Core ML features from WWDC’ 24

Here are some of the most impactful Core ML features from WWDC’ 24 we
will use to run Mistral 7B on a Mac.

Swift Tensor

The first feature we want to highlight is an entirely new Swift type to
work with ML tensors. These are multi-dimensional data structures every
ML framework uses. Python developers working on ML are familiar with
numpy arrays or torch tensors, which provide convenient,
high-level interfaces to manipulate these large multi-dimensional
matrices easily. The new MLTensor type provides a high-level
abstraction that mimics the ones available in Python frameworks, greatly
simplifying working with tensor data in Swift.

Core ML already had multi-dimensional data types in the form of
MLMultiArray
and
MLShapedArray.
However, they were only meant for data storage and simple operations
like wrapping your data and sending it as input to a Core ML model, or
unwrapping results from a Core ML model. However, manipulating tensor
data with these APIs is difficult. Only a few primitive operations are
provided, and you may have to write your own by accessing the underlying
storage as an opaque pointer to number data. This is time-consuming and
error-prone.

The new Tensor type provides a high-level abstraction that mimics
the ones available in Python frameworks, greatly simplifying working
with tensor data in Swift. Consider a language model like the one we
want to port to Core ML. Language models take in an input sequence of
tokens, and they output an estimation of the probabilities of all the
tokens in the vocabulary, meaning that tokens with a high probability
have a high chance of being plausible continuations of the input. The
application’s job is to select the best next token to append to the
sequence based on those probabilities. Tensor type makes it easy to
handle these operations without custom code.

When we released swift-transformers,
we wrote a lot of code (later extended by the community, thanks! ❤️) to
help with input preparations (convert words to tokens) and output
post-processing. For example, check out our softmax operation
using Accelerate. All this can be removed when using MLTensor, as
softmax is provided out of the box!

Stateful Buffers

Before WWDC’ 24, a Core ML model was essentially a pure stateless
function where you provide inputs and return some outputs. However,
sometimes you need to keep a state that depends on previous
computations. The functional programming method for maintaining state is
to add an additional input/output pair. So, based on your inputs and
state, the model computes the output and the new state. There is nothing
wrong with this approach, and in fact, that’s the way high-performance
frameworks like JAX work.

However, there are practical limitations: the stateful data needs to be
sent to the model as an input and retrieved as an output every time you
call the model. If the stateful data is large, then all this going back
and forth increases overhead and slows things down. This is particularly
important for LLMs because you have to run many iterations to generate a
sequence. The performance bottleneck is usually your computer’s memory
bandwidth (i.e., how fast you can move things to your GPU and back).
Stateful models solve this problem by reserving a block of memory for
state data and keeping it on the GPU so you don’t have to send and
receive it every time you use the model.

Stateful buffers were introduced in this WWDC’ 24 session
using a toy example that is easy to understand but not representative of
practical uses with big models such as LLMs. An LLM performance trick
for transformers-based models is key-value caching (known as
kv-caching). As shown in the following illustration, it avoids costly
matrix multiplications in the crucial attention block by caching the
result of previous operations performed in previous steps. We won’t go
into details, but the takeaways are: kv-cache dramatically increases
performance, and it requires a large block of memory that is the perfect
candidate for using stateful buffers. Here is a coremltools user guide
update about stateful models.

New Quantization Techniques

In WWDC 23, we explored a very cool technique called palletization, and
we showed how it could help bring text-to-image models, such as Stable
Diffusion, to Macs and iPhones.

Whilst these techniques allow you to reduce the size considerably, if
pushed too far, the impact on quality is drastic. Bigger models suffer
more from this, as the weight data has an extensive dynamic range.
Trying to create a small lookup table (LUT) that captures all possible
values becomes increasingly difficult. The solution introduced in WWDC
24 is to focus on a smaller portion of the data at a time, and create
multiple lookup tables for different areas of the same tensor.

These methods (block-wise quantization) allow us to compress models to
as low as 4-bit precision. Instead of using 4 bytes (the size of a
float32 number) to represent each model parameter, we can get away
with half a byte (a nibble) for each. This is an 8-fold reduction in
model size (minus some overhead to account for the block-wise
quantization tables), or 4 times smaller when compared to float16
precision.

Multifunction Support

We won’t use this feature for this example but we wanted to mention it
here as it was introduced at WWDC 24, and we will be showcasing it in
some upcoming work. Multifunction support essentially allows you to
package LoRA adapters into generative models to use the same model (with
a small set of additional parameters, called adapters) for different
tasks. LoRA is the preferred community technique for large model
fine-tuning. In diffusion models, for example, you can use LoRA to
generate images with different styles, such as photorealistic or
cartoonish. We believe LoRA is part of the solution that powers Apple’s
Genmoji implementation. For language models, LoRA adapters can be used
to adapt a generic LLM to specific tasks or domains.

To read more about LoRA, you can check this post.

To read more about Multifunction, you can check out Apple coremltools
user guide here.

Converting Mistral 7B to Core ML

The single most important component for running a large language model
efficiently is the kv-cache. As mentioned above, this is a great
candidate for the new stateful model feature
released at WWDC’ 24. Models in the transformers library already use
efficient attention implementations that rely heavily on kv-caching.
However, the default implementations are optimized for Nvidia GPUs, and
this hardware has a different set of constraints than Apple Silicon
does. In the case of Core ML, we need to pre-allocate the full cache
buffer beforehand and ensure that each time we call the model, we update
the buffer in place. This avoids inefficient memory allocations and
tensor concatenations and is also a requirement for Core ML stateful
buffers.

To achieve this goal, we have to use a different attention
implementation that considers these factors. This requires modifying the
transformers modeling code for the Mistral architecture, and it’s done
in this fragment of code.

Note: If you want to follow along and replicate the conversion (or
convert another Mistral-based model, like a different fine-tune), you
can use this script
to run all the conversion steps.

Tracing & Conversion

The first step is to load the model. We’ll use the patched
implementation with the in-place cache method.

MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
torch_model = StatefulMistralForCausalLM(MODEL_ID)
torch_model.eval()

Before running Core ML conversion, we need to trace the model with
example inputs. This process records the tensor operations performed on
those inputs, and the traced graph will be translated to Core ML
operations during conversion. We use sample inputs to trace the model;
we don’t need real data.

input_ids = torch.zeros((1, 2), dtype=torch.int32)
causal_mask = torch.zeros((1, 1, 2, 5), dtype=torch.float32)

traced_model = torch.jit.trace(torch_model, [input_ids, causal_mask])

The input to a language model is a sequence of tokens of varying length.
We’ll allow the input to grow from a single token to a maximum context
length of 2048. We can use
coremltools range
dimensions to specify these bounds.

query_length = ct.RangeDim(lower_bound=1, upper_bound=2048, default=1)
end_step_dim = ct.RangeDim(lower_bound=1, upper_bound=2048, default=1)

inputs = [
    ct.TensorType(shape=(1, query_length), dtype=np.int32, name="inputIds"),
    ct.TensorType(shape=(1, 1, query_length, end_step_dim), dtype=np.float16, name="causalMask"),
]

outputs = [ct.TensorType(dtype=np.float16, name="logits")]

In addition to the sequence tokens (called inputIds in the example
above), there’s another input called causalMask, which specifies the
tokens the model needs to pay attention to. This is mostly used when
generating multiple sequences at the same time using batching. Check out
how these inputs are used in an example runner
here.

In this situation, all the input sequences inside a batch must have the
same length, so we use padding tokens and the causal mask to tell the
model that the padding tokens are not to be considered as inputs.

State Preparation

The PyTorch modeling code uses keyCache and valueCache as the
names of the cache buffers to hold the kv-cache. Those blocks are
allocated for the maximum context length (2048). We use coremltools‘
new
StateType
to specify that those blocks must be converted to a stateful Core ML
buffer during conversion.

states = [
    ct.StateType(
        wrapped_type=ct.TensorType(shape=torch_model.kv_cache_shape, dtype=np.float16),
        name="keyCache",
    ),
    ct.StateType(
        wrapped_type=ct.TensorType(shape=torch_model.kv_cache_shape, dtype=np.float16),
        name="valueCache",
    ),
]

Core ML Conversion

To convert the model to Core ML, we need to specify the input and output
types, as well as the states. The converted model will use float16
precision because that’s what we specified for the input data. We also
need to indicate the minimum deployment target as iOS18, as that’s where
these features are available. (We can also use macOS15, which refers
to the same conversion target.)

mlmodel_fp16 = ct.convert(
    traced_model,
    inputs=inputs,
    states=states,
    outputs=outputs,
    minimum_deployment_target=ct.target.iOS18,
    skip_model_load=True,
)

Model Compression

Using the new block-size quantization strategies described above, we use
4-bit linear quantization with block size 32. This will greatly reduce
model size and make the model run faster. Even though computation will
still be performed in float16, weights are transferred in 4-bit mode
and decompressed on the fly, which is more efficient than transferring a
large amount of 16-bit weights.

The quantization parameters are configured as follows:

op_config = ct.optimize.coreml.OpLinearQuantizerConfig(
    mode="linear_symmetric",
    dtype="int4",
    granularity="per_block",
    block_size=32,
)
config = ct.optimize.coreml.OptimizationConfig(global_config=op_config)

Let’s use that configuration to quantize the model. The following line
will take a few minutes to run:

mlmodel_int4 = ct.optimize.coreml.linear_quantize_weights(mlmodel_fp16, config=config)

mlmodel_int4.save("StatefulMistral7BInstructInt4.mlpackage")

There’s a final step after conversion and quantization are done. We need
to include a piece of additional metadata that indicates the model
identifier we used (mistralai/Mistral-7B-Instruct-v0.3). The Swift
code will use this to download the tokenizer files from the Hub.
Tokenization is converting text data to the numerical representations
used by models, and it’s different for every model.

mlmodel_int4._spec.description.metadata.userDefined.update({
    "co.huggingface.exporters.name": MODEL_ID
})

The generated model is a mlpackage of about 3.8G, compared with the
14G that a float16 conversion would produce. You can find it
here on the
Hub.

Running Mistral 7B with Swift

If you followed the steps above or downloaded the model from the Hub,
you can run it locally using the preview branch of
swift-transformers. Apple engineers contributed it to the project,
including the following important features:

• Full Tensor support, which greatly simplifies pre- and
  post-processing tasks, and allows us to delete many lines of
  low-level, confusing and fragile code.

• Support for the Swift counterpart of the Stateful API.

Since adopting these features is a breaking change and requires iOS 18
or macOS 15, we’ll keep them in a preview branch for now.

To run the model from the command line, please first clone the preview
branch from the GitHub repo:

    git clone -b preview https://github.com/huggingface/swift-transformers

And then run the CLI to test the model:

swift run transformers "Best recommendations for a place to visit in Paris in August 2024:" --max-length 128 Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage

For easier testing, you can also use swift-chat, a simple app we
wrote to show how to integrate the swift-transformers package
inside. You have to use the preview branch as well. An example of
swift-chat running the converted Mistral model was shown at the
beginning of this post.

Running Mistral 7B with Python

For those of you who are more familiar with Python, it’s just as easy!

python3 generate.py Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage --prompt "Best recommendations for a place to visit in Paris in August 2024:"

coremltools makes it just as easy to run Core ML models with Python.

What’s Next?

We are extremely excited about the progress in Core ML and
coremltools this year,
and we are looking forward to seeing lots of third-party apps leveraging
ML models to solve real tasks people need. On our side, we are committed
to making this as easy as possible so developers can concentrate on
creating cool apps. There are a few things on our drawing board:

• The model updates presented here are excellent for GPUs on Mac
  computers. Core ML can use the Neural Engine, which is particularly
  efficient on iPhones. Getting the most performance out of the Neural
  Engine requires some additional adaptations, which we plan to carry
  out on a few example models. This work will be based on the
  learnings discussed in this 2022 (and still very relevant) article by Apple.
  We won’t run Mistral 7B on iPhone, but there are several smaller
  models, like Apple’s OpenELM or DCLM that make for great
  candidates to explore!

• The code presented here is highly experimental. As summer goes on,
  we plan to adopt these methods and incorporate them into
  exporters, a Python tool designed to convert transformers models
  to Core ML. Hopefully, you’ll soon be able to convert many
  interesting model architectures very easily.

• We’ll keep working on the preview branch of
  swift-transformers to incorporate new features or API changes as
  they are released. If you are interested, keep an eye on it!

How can you help?

The tools released by Apple in WWDC help us on our long-term goal to
make AI easy and accessible to all, and we’d love to see where you can
take them. The example we showed is experimental, but you can use it to
convert any Mistral fine-tune to Core ML – please let us know if you do!
If you want to try other model architectures, please feel free to open
issues or PRs to the preview branch of swift-transformers –
we’ll try to help you get going!

There’s never been a better time than today to apply your creativity to
solve problems that interest you! Go try things, have fun, and tell us
how we can help.