mlx-lm

Run LLMs with MLX

5,204 stars Python #llms#mlx
RAW Doc

File: README.md

MLX LM

MLX LM is a Python package for generating text and fine-tuning large language
models on Apple silicon with MLX.

Some key features include:

  • Integration with the Hugging Face Hub to easily use thousands of LLMs with a
    single command.
  • Support for quantizing and uploading models to the Hugging Face Hub.
  • Low-rank and full model
    fine-tuning

    with support for quantized models.
  • Distributed inference and fine-tuning with mx.distributed

The easiest way to get started is to install the mlx-lm package:

With pip:

sh
pip install mlx-lm

With conda:

sh
conda install -c conda-forge mlx-lm

Quick Start

To generate text with an LLM use:

bash
mlx_lm.generate --prompt "How tall is Mt Everest?"

To chat with an LLM use:

bash
mlx_lm.chat

This will give you a chat REPL that you can use to interact with the LLM. The
chat context is preserved during the lifetime of the REPL.

Commands in mlx-lm typically take command line options which let you specify
the model, sampling parameters, and more. Use -h to see a list of available
options for a command, e.g.:

bash
mlx_lm.generate -h

The default model for generation and chat is
mlx-community/Llama-3.2-3B-Instruct-4bit. You can specify any MLX-compatible
model with the --model flag. Thousands are available in the
MLX Community Hugging Face
organization.

Python API

You can use mlx-lm as a module:

python
from mlx_lm import load, generate

model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")

prompt = "Write a story about Einstein"

messages = [{"role": "user", "content": prompt}]
prompt = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True,
)

text = generate(model, tokenizer, prompt=prompt, verbose=True)

To see a description of all the arguments you can do:

text
>>> help(generate)

Check out the generation
example

to see how to use the API in more detail. Check out the batch generation
example

to see how to efficiently generate continuations for a batch of prompts.

The mlx-lm package also comes with functionality to quantize and optionally
upload models to the Hugging Face Hub.

You can convert models using the Python API:

python
from mlx_lm import convert

repo = "mistralai/Mistral-7B-Instruct-v0.3"
upload_repo = "mlx-community/My-Mistral-7B-Instruct-v0.3-4bit"

convert(repo, quantize=True, upload_repo=upload_repo)

This will generate a 4-bit quantized Mistral 7B and upload it to the repo
mlx-community/My-Mistral-7B-Instruct-v0.3-4bit. It will also save the
converted model in the path mlx_model by default.

To see a description of all the arguments you can do:

text
>>> help(convert)

Streaming

For streaming generation, use the stream_generate function. This yields
a generation response object.

For example,

python
from mlx_lm import load, stream_generate

repo = "mlx-community/Mistral-7B-Instruct-v0.3-4bit"
model, tokenizer = load(repo)

prompt = "Write a story about Einstein"

messages = [{"role": "user", "content": prompt}]
prompt = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True,
)

for response in stream_generate(model, tokenizer, prompt, max_tokens=512):
    print(response.text, end="", flush=True)
print()

Sampling

The generate and stream_generate functions accept sampler and
logits_processors keyword arguments. A sampler is any callable which accepts
a possibly batched logits array and returns an array of sampled tokens. The
logits_processors must be a list of callables which take the token history
and current logits as input and return the processed logits. The logits
processors are applied in order.

Some standard sampling functions and logits processors are provided in
mlx_lm.sample_utils.

Command Line

You can also use mlx-lm from the command line with:

text
mlx_lm.generate --model mistralai/Mistral-7B-Instruct-v0.3 --prompt "hello"

This will download a Mistral 7B model from the Hugging Face Hub and generate
text using the given prompt.

For a full list of options run:

text
mlx_lm.generate --help

To quantize a model from the command line run:

text
mlx_lm.convert --model mistralai/Mistral-7B-Instruct-v0.3 -q

For more options run:

text
mlx_lm.convert --help

You can upload new models to Hugging Face by specifying --upload-repo to
convert. For example, to upload a quantized Mistral-7B model to the
MLX Hugging Face community you can do:

text
mlx_lm.convert \
    --model mistralai/Mistral-7B-Instruct-v0.3 \
    -q \
    --upload-repo mlx-community/my-4bit-mistral

Models can also be converted and quantized directly in the
mlx-my-repo Hugging
Face Space.

Long Prompts and Generations

mlx-lm has some tools to scale efficiently to long prompts and generations:

  • A rotating fixed-size key-value cache.
  • Prompt caching

To use the rotating key-value cache pass the argument --max-kv-size n where
n can be any integer. Smaller values like 512 will use very little RAM but
result in worse quality. Larger values like 4096 or higher will use more RAM
but have better quality.

Caching prompts can substantially speedup reusing the same long context with
different queries. To cache a prompt use mlx_lm.cache_prompt. For example:

bash
cat prompt.txt | mlx_lm.cache_prompt \
  --model mistralai/Mistral-7B-Instruct-v0.3 \
  --prompt - \
  --prompt-cache-file mistral_prompt.safetensors

Then use the cached prompt with mlx_lm.generate:

text
mlx_lm.generate \
    --prompt-cache-file mistral_prompt.safetensors \
    --prompt "\nSummarize the above text."

The cached prompt is treated as a prefix to the supplied prompt. Also notice
when using a cached prompt, the model to use is read from the cache and need
not be supplied explicitly.

Prompt caching can also be used in the Python API in order to avoid
recomputing the prompt. This is useful in multi-turn dialogues or across
requests that use the same context. See the
example
for more usage details.

Supported Models

mlx-lm supports thousands of LLMs available on the Hugging Face Hub. If the
model you want to run is not supported, file an
issue or better yet, submit
a pull request. Many supported models are available in various quantization
formats in the MLX Community Hugging
Face organization.

For some models the tokenizer may require you to enable the trust_remote_code
option. You can do this by passing --trust-remote-code in the command line.
If you don't specify the flag explicitly, you will be prompted to trust remote
code in the terminal when running the model.

Tokenizer options can also be set in the Python API. For example:

python
model, tokenizer = load(
    "qwen/Qwen-7B",
    tokenizer_config={"eos_token": "<|endoftext|>", "trust_remote_code": True},
)

Large Models

[!NOTE]
This requires macOS 15.0 or higher to work.

Models which are large relative to the total RAM available on the machine can
be slow. mlx-lm will attempt to make them faster by wiring the memory
occupied by the model and cache. This requires macOS 15 or higher to
work.

If you see the following warning message:

[WARNING] Generating with a model that requires ...

then the model will likely be slow on the given machine. If the model fits in
RAM then it can often be sped up by increasing the system wired memory limit.
To increase the limit, set the following sysctl:

bash
sudo sysctl iogpu.wired_limit_mb=N

The value N should be larger than the size of the model in megabytes but
smaller than the memory size of the machine.


File: mlx_lm/README.md

Generate Text with MLX and :hugs: Hugging Face

This an example of large language model text generation that can pull models from
the Hugging Face Hub.

For more information on this example, see the README in the
parent directory.

This package also supports fine tuning with LoRA or QLoRA. For more information
see the LoRA documentation.