For years, running machine learning on a Mac came with a qualification: fine for development, but serious GPU work usually meant an NVIDIA machine. Apple's Metal backend for PyTorch improved the situation, but it still meant working within a framework designed for a much broader range of hardware.
MLX takes a different approach. Released by Apple's machine learning research team in December 2023 under the MIT license, it was designed around Apple silicon from the start. You can use it to experiment with neural networks, run quantized language models, and fine-tune smaller models without sending the work to a server.
The interesting part isn't simply that MLX runs on a Mac. It's how the framework handles memory, execution, and devices.
What is MLX?
MLX is an array framework with a NumPy-like API, automatic differentiation, GPU acceleration, and packages for neural networks and optimizers. If you've worked with NumPy, JAX, or PyTorch, much of it will look familiar.
The Python API is useful for experimentation, while C, C++, and Swift APIs extend the same underlying ideas to other applications. MLX Swift also makes it possible to build native features for macOS, iOS, iPadOS, and other Apple platforms.
MLX is deliberately small at the conceptual level. You work with arrays and functions, decide when to evaluate computations, and let the framework handle much of the scheduling.
Why MLX feels different
Unified memory changes the programming model
On Apple silicon, the CPU and GPU share a physical memory pool. MLX allocates arrays in unified memory, allowing supported operations to run on either device without explicitly transferring those arrays between separate memory pools.
This needs a little context: PyTorch's MPS backend also benefits from Apple silicon's unified memory. MLX's distinction is how directly it exposes that architecture. Rather than placing a tensor on a device and moving it when necessary, you can choose the device for an operation.
import mlx.core as mx
a = mx.array([1.0, 2.0, 3.0])
b = mx.array([4.0, 5.0, 6.0])
cpu_result = mx.add(a, b, stream=mx.cpu)
gpu_result = mx.multiply(a, b, stream=mx.gpu)
mx.eval(cpu_result, gpu_result)For comparison, a typical PyTorch MPS operation explicitly places its tensor on the MPS device:
import torch
x = torch.tensor([1.0, 2.0, 3.0], device="mps")
y = x * 2The difference isn't proof that one framework will always be faster. It is a difference in the way the developer expresses where computation happens.
Unified memory also gives Macs a useful advantage for local language models: the GPU isn't restricted to a small, separate VRAM allocation. But the full advertised RAM capacity isn't available to a model. macOS, other applications, model activations, caches, and runtime buffers all need room too.
Lazy computation
MLX normally builds a computation graph before executing it. Operations are evaluated when their values are needed or when you explicitly call mx.eval.
import mlx.core as mx
a = mx.array([1.0, 2.0, 3.0])
b = a * 2 + 1 # Builds a computation graph
mx.eval(b)
print(b) # array([3, 5, 7], dtype=float32)That separation gives MLX opportunities to optimize execution, and it avoids calculating values a program never uses. It also means you need to be deliberate about evaluation inside training loops; otherwise, the pending computation graph can grow unnecessarily.
Dynamic graphs and function transformations
MLX builds graphs dynamically, so ordinary operations can work with changing inputs without a separate tracing step. It also provides function transformations such as mx.grad, mx.value_and_grad, mx.vmap, and mx.compile.
def squared_sum(x):
return (x ** 2).sum()
x = mx.array([1.0, 2.0, 3.0])
print(mx.grad(squared_sum)(x))
# array([2, 4, 6], dtype=float32)There is an important distinction between dynamic execution and explicit compilation. With mx.compile, changing input shapes can trigger recompilation. shapeless=True can avoid that for suitable functions, but shape-dependent logic needs extra care.
Training a small model
The mlx.nn and mlx.optimizers packages will feel familiar to PyTorch users. Here's a complete training example that doesn't require downloading a dataset.
We'll generate two-dimensional points and train a small network to predict whether the sum of their coordinates is positive. It's a toy task, but it demonstrates the forward pass, gradients, optimizer updates, and evaluation in one place.
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
mx.random.seed(0)
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.l1 = nn.Linear(2, 16)
self.l2 = nn.Linear(16, 2)
def __call__(self, x):
return self.l2(nn.relu(self.l1(x)))
# Two features per sample; class 1 if their sum is positive.
x = mx.random.normal((256, 2))
y = (x[:, 0] + x[:, 1] > 0).astype(mx.int32)
def loss_fn(model, x, y):
return nn.losses.cross_entropy(model(x), y).mean()
model = MLP()
mx.eval(model.parameters())
optimizer = optim.Adam(learning_rate=1e-2)
loss_and_grad = nn.value_and_grad(model, loss_fn)
for epoch in range(51):
loss, grads = loss_and_grad(model, x, y)
optimizer.update(model, grads)
mx.eval(loss, model.parameters(), optimizer.state)
if epoch % 10 == 0:
print(f"Epoch {epoch}: loss={loss.item():.4f}")The loss should generally fall as training proceeds, although the exact numbers depend on initialization. The key line is mx.eval(loss, model.parameters(), optimizer.state): it evaluates the loss and updated training state at each iteration instead of letting pending work accumulate.
For an example using real handwritten-digit data, Apple's MLX MNIST implementation is a sensible next step.
Running and fine-tuning LLMs with mlx-lm
For many developers, the easiest entry point is mlx-lm, the companion package for language-model inference and fine-tuning.
Install it and generate text with a quantized, open-weight model:
python -m pip install mlx-lm
mlx_lm.generate \
--model mlx-community/Llama-3.2-3B-Instruct-4bit \
--prompt "Explain unified memory in one paragraph." \
--max-tokens 200The model downloads from Hugging Face on first use. You can also start an interactive session with mlx_lm.chat.
From Python, apply the instruction model's chat template rather than treating a user message as unformatted text:
from mlx_lm import load, generate
model, tokenizer = load(
"mlx-community/Llama-3.2-3B-Instruct-4bit"
)
messages = [
{"role": "user", "content": "Explain unified memory in one paragraph."}
]
prompt = tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
)
response = generate(
model,
tokenizer,
prompt=prompt,
max_tokens=200,
)
print(response)The MLX Community on Hugging Face maintains a large collection of converted and quantized models, so in many cases you can use an existing checkpoint rather than converting one yourself.
Fine-tuning with LoRA and QLoRA
MLX-LM also supports parameter-efficient fine-tuning. LoRA trains small adapter weights instead of updating every weight in the base model. When the base model is quantized, MLX-LM can use QLoRA, which further reduces the memory needed for the frozen weights.
Install the training dependencies:
python -m pip install "mlx-lm[train]"Prepare a directory named my_data with a train.jsonl file. Each line needs to follow one of the supported dataset formats. A chat-format record looks like this:
{"messages":[{"role":"user","content":"What is MLX?"},{"role":"assistant","content":"MLX is an array framework optimized for Apple silicon."}]}Use a meaningful collection of training examples, not just the single illustrative record above. An optional valid.jsonl file lets you monitor validation loss.
mlx_lm.lora \
--model mlx-community/Llama-3.2-3B-Instruct-4bit \
--train \
--data ./my_data \
--iters 100Memory requirements still depend on the base model, sequence length, batch size, and training settings. A small quantized model is the sensible first experiment on a laptop.
Quantization is central to making local inference practical. Four-bit weights theoretically require one-quarter the raw storage of 16-bit weights, but total memory savings are smaller once scales, metadata, activations, and the KV cache are included.
Where the ecosystem stands
MLX is no longer confined to short research examples.
Native Apple apps: MLX Swift supports on-device workloads across Apple's platforms. Apple dedicated a WWDC 2025 session to the framework.
Local AI tooling: Ollama introduced an MLX-powered Apple silicon inference preview in March 2026, another indication of MLX's growing role beyond research code.
More than macOS: The Python project now documents CUDA-enabled and CPU-only Linux installation options. Apple silicon remains its original design target; check platform requirements rather than assuming every backend has identical support.
Multi-machine work:
mx.distributedsupports distributed computation, including setups spanning multiple Apple silicon machines.
The project changes quickly. For current versions and compatibility, consult its release history instead of relying on an old installation guide.
When MLX is — and isn't — the right choice
MLX is a strong option when you want to explore model internals on a Mac, fine-tune a model within a local memory budget, prototype an ML idea, or add inference to a native Apple application. Keeping data local can also be useful where cloud processing is undesirable.
It isn't a drop-in replacement for the entire PyTorch ecosystem. Specialized architectures and research libraries may not have MLX implementations, and large-scale training still benefits from multi-GPU infrastructure. Performance comparisons also depend on the model, precision, batch size, and workload; a unified-memory architecture doesn't guarantee a universal speed advantage.
If all you want is to chat with a downloaded model, a higher-level application may be simpler. If you want to work with the arrays, training steps, and inference code yourself, MLX gives you that control.
Getting started
Use an Apple silicon Mac with macOS 14 or later and a native ARM Python 3.10 or newer.
Install MLX with
python -m pip install mlx; addmlx-lmif you want to work with language models.Try the array and training snippets above, then browse the official MLX examples.
Start LLM experiments with a small quantized model that fits your available memory.
Read the official explanations of unified memory and lazy evaluation. They explain most of the framework's less familiar behavior.
Final thoughts
MLX doesn't turn a MacBook into a GPU server, and it doesn't need to. It makes Apple silicon a more practical place to experiment with models, run local inference, and fine-tune without building the workflow around a remote machine. For developers who want control over the underlying code rather than just a chat interface, that's where MLX earns its place.
Comments(0)
Join the discussion by logging into your account.
No comments yet. Be the first to comment!