MLX Tutorial: Apple's Machine Learning Framework for Apple Silicon

MLX is an open-source machine learning framework from Apple, designed specifically for Apple Silicon chips (M1, M2, M3, M4). The framework offers a familiar API for NumPy and PyTorch users, with performance optimized for Apple hardware. MLX has become the go-to choice for developers who want to run ML models locally on Mac devices without requiring NVIDIA GPUs.

In this tutorial, we will cover installation, basic usage, LLM fine-tuning, model inference, and best practices to maximize MLX performance on Apple Silicon devices.

Why MLX?

Before diving in, let's understand why MLX is worth learning:

Installation

Prerequisites

• macOS 13.5 or later
• Apple Silicon (M1/M2/M3/M4)
• Python 3.9 or later

Installing MLX Core

pip install mlx

Installing MLX-LM (for Language Models)

pip install mlx-lm

Installing MLX-VLM (for Vision-Language Models)

pip install mlx-vlm

Installing from Source (Optional)

git clone https://github.com/ml-explore/mlx.git
cd mlx
pip install -e .

Verify Installation

import mlx.core as mx
import mlx.nn as nn

print(f"MLX version: {mx.version}")
print(f"Default device: {mx.defaultdevice()}")


Simple test
a = mx.array([1, 2, 3, 4, 5])
print(f"Array: {a}")
print(f"Sum: {mx.sum(a)}")

MLX version: 0.x.x
Default device: Device(gpu, 0)
Array: array([1, 2, 3, 4, 5], dtype=int32)
Sum: array(15, dtype=int32)

import mlx.core as mx

Creating arrays
a = mx.array([1.0, 2.0, 3.0, 4.0])
b = mx.ones((3, 4))
c = mx.zeros((2, 3))
d = mx.random.normal((5, 5))

print(f"Shape of b: {b.shape}")
print(f"Dtype of a: {a.dtype}")

Mathematical operations
x = mx.array([[1, 2], [3, 4]], dtype=mx.float32)
y = mx.array([[5, 6], [7, 8]], dtype=mx.float32)

Element-wise operations
print(f"Addition: {x + y}")
print(f"Multiplication: {x  y}")


Matrix multiplication
print(f"MatMul: {x @ y}")

Reduction operations
print(f"Sum: {mx.sum(x)}")
print(f"Mean: {mx.mean(x)}")
print(f"Max: {mx.max(x)}")

import mlx.core as mx

a = mx.ones((1000, 1000))
b = mx.ones((1000, 1000))

These operations are not yet executed
c = a + b
d = c  2

Evaluation happens when we need the result
mx.eval(d)
print(d)

Or when converting to Python/NumPy
result = d.tolist()

import mlx.core as mx

Check default device
print(f"Default device: {mx.defaultdevice()}")

Run on CPU
mx.setdefaultdevice(mx.cpu)
a = mx.ones((100, 100))
print(f"Device: {mx.defaultdevice()}")


Switch back to GPU
mx.setdefaultdevice(mx.gpu)

b = mx.ones((100, 100))
print(f"Device: {mx.defaultdevice()}")

Building Neural Networks with MLX

Simple Model

MLX provides the mlx.nn module, similar to PyTorch:

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim

class SimpleNet(nn.Module):
    def init(self):
        super().init()
        self.layers = [
            nn.Linear(784, 256),
            nn.Linear(256, 128),
            nn.Linear(128, 10),
        ]

    def call(self, x):
        for i, layer in enumerate(self.layers[:-1]):
            x = nn.relu(layer(x))
        return self.layers-1

Initialize model
model = SimpleNet()

Check parameters
params = model.parameters()
print(f"Model initialized with multiple layers")

Training Loop

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np

class MLP(nn.Module):
    def init(self, inputdim, hiddendim, outputdim):

        super().init()
        self.linear1 = nn.Linear(inputdim, hiddendim)

        self.linear2 = nn.Linear(hiddendim, outputdim)


    def call(self, x):
        x = nn.relu(self.linear1(x))
        return self.linear2(x)

Synthetic data
np.random.seed(42)
X = np.random.randn(1000, 10).astype(np.float32)
y = (X[:, 0] + X[:, 1] > 0).astype(np.int32)

Xtrain = mx.array(X[:800])
ytrain = mx.array(y[:800])

Xtest = mx.array(X[800:])
ytest = mx.array(y[800:])


Model and optimizer
model = MLP(10, 64, 2)
optimizer = optim.Adam(learningrate=1e-3)

Loss function
def lossfn(model, x, y):

    logits = model(x)
    return mx.mean(nn.losses.crossentropy(logits, y))

Training
lossandgradfn = nn.valueandgrad(model, lossfn)

for epoch in range(100):
    loss, grads = lossandgradfn(model, Xtrain, ytrain)

    optimizer.update(model, grads)
    mx.eval(model.parameters(), optimizer.state)

    if (epoch + 1) % 20 == 0:
        testlogits = model(Xtest)

        testpreds = mx.argmax(testlogits, axis=1)

        accuracy = mx.mean(testpreds == ytest)

        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}, Test Acc: {accuracy.item():.4f}")

import mlx.core as mx
import mlx.nn as nn

class ConvNet(nn.Module):
    def init(self, numclasses=10):
        super().init()
        self.conv1 = nn.Conv2d(1, 32, kernelsize=3, padding=1)

        self.conv2 = nn.Conv2d(32, 64, kernelsize=3, padding=1)
        self.pool = nn.MaxPool2d(kernelsize=2, stride=2)

        self.fc1 = nn.Linear(64  7  7, 128)
        self.fc2 = nn.Linear(128, numclasses)
        self.dropout = nn.Dropout(p=0.5)

    def call(self, x):
        x = self.pool(nn.relu(self.conv1(x)))
        x = self.pool(nn.relu(self.conv2(x)))
        x = x.reshape(x.shape[0], -1)
        x = self.dropout(nn.relu(self.fc1(x)))
        return self.fc2(x)

model = ConvNet()
print("ConvNet initialized successfully")

Using MLX-LM for Language Models

MLX-LM is a library that simplifies working with Large Language Models on Apple Silicon.

Inference with Hugging Face Models

# Download and run model
mlxlm.generate --model mlx-community/Llama-3.2-3B-Instruct-4bit \

  --prompt "Explain machine learning in simple terms" \
  --max-tokens 200

from mlxlm import load, generate

Load model (automatically downloads from Hugging Face)
model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

Generate text
prompt = "Explain what machine learning is in simple terms:"

messages = [{"role": "user", "content": prompt}]
formattedprompt = tokenizer.applychattemplate(

    messages, tokenize=False, addgenerationprompt=True

)

response = generate(
    model,
    tokenizer,
    prompt=formattedprompt,
    maxtokens=500,

    temp=0.7,
)

print(response)

from mlxlm import load, streamgenerate

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

prompt = "Write a Python function to calculate fibonacci numbers:"

messages = [{"role": "user", "content": prompt}]
formattedprompt = tokenizer.applychattemplate(
    messages, tokenize=False, addgenerationprompt=True
)

Streaming output
for token in streamgenerate(

    model,
    tokenizer,
    prompt=formattedprompt,
    maxtokens=500,

):
    print(token, end="", flush=True)
print()

# Convert with 4-bit quantization
mlxlm.convert \
  --hf-path meta-llama/Llama-3.2-3B-Instruct \
  --mlx-path ./mlx-llama-3.2-3b-4bit \
  --quantize \
  --q-bits 4

Convert without quantization
mlxlm.convert \

  --hf-path microsoft/Phi-3-mini-4k-instruct \
  --mlx-path ./mlx-phi-3-mini

from mlxlm import convert

convert(
    hfpath="meta-llama/Llama-3.2-1B-Instruct",

    mlxpath="./mlx-llama-1b",
    quantize=True,
    qbits=4,

    qgroupsize=64,

)

# Prepare training data in JSONL format
train.jsonl format:
{"text": "instruction: ... response: ..."}
or chat format:
{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

Create a configuration file loraconfig.yaml:

model: "mlx-community/Llama-3.2-3B-Instruct-4bit"
train: true
data: "./trainingdata"

seed: 42
loralayers: 16
batchsize: 4

iters: 1000
valbatches: 25
learningrate: 1e-5

loraparameters:
  keys: ["selfattn.qproj", "selfattn.vproj"]
  rank: 8
  alpha: 16
  dropout: 0.0

Run fine-tuning:

mlxlm.lora \
  --model mlx-community/Llama-3.2-3B-Instruct-4bit \
  --train \
  --data ./trainingdata \
  --batch-size 4 \
  --lora-layers 16 \
  --iters 1000 \
  --learning-rate 1e-5

Preparing Training Data

import json

Format data for fine-tuning
trainingdata = [

    {
        "messages": [
            {"role": "user", "content": "What is Python?"},
            {"role": "assistant", "content": "Python is a high-level programming language designed for code readability. Python supports multiple programming paradigms including procedural, object-oriented, and functional programming."}
        ]
    },
    {
        "messages": [
            {"role": "user", "content": "How do you create a list in Python?"},
            {"role": "assistant", "content": "To create a list in Python, use square brackets []. For example: mylist = [1, 2, 3, 'hello']. Lists in Python can store various data types."}
        ]
    },
]

Save as JSONL
with open("trainingdata/train.jsonl", "w") as f:

    for item in trainingdata:
        f.write(json.dumps(item) + "\n")

Validation data
with open("trainingdata/valid.jsonl", "w") as f:

    for item in trainingdata[:10]:
        f.write(json.dumps(item) + "\n")

Using Fine-Tuned Adapters

from mlxlm import load, generate

Load model with adapter
model, tokenizer = load(
    "mlx-community/Llama-3.2-3B-Instruct-4bit",
    adapterpath="./adapters"
)

Generate with fine-tuned model
response = generate(
    model,
    tokenizer,
    prompt="What is Python?",
    maxtokens=200

)
print(response)

mlxlm.fuse \
  --model mlx-community/Llama-3.2-3B-Instruct-4bit \
  --adapter-path ./adapters \
  --save-path ./fused-model

MLX for Computer Vision

Image Classification

import mlx.core as mx
import mlx.nn as nn

class ResidualBlock(nn.Module):
    def init(self, channels):
        super().init()
        self.conv1 = nn.Conv2d(channels, channels, kernelsize=3, padding=1)

        self.bn1 = nn.BatchNorm(channels)
        self.conv2 = nn.Conv2d(channels, channels, kernelsize=3, padding=1)
        self.bn2 = nn.BatchNorm(channels)

    def call(self, x):
        residual = x
        x = nn.relu(self.bn1(self.conv1(x)))
        x = self.bn2(self.conv2(x))
        return nn.relu(x + residual)

class SimpleResNet(nn.Module):
    def init(self, numclasses=10):

        super().init()
        self.conv1 = nn.Conv2d(3, 64, kernelsize=7, stride=2, padding=3)
        self.bn1 = nn.BatchNorm(64)
        self.pool = nn.MaxPool2d(kernelsize=3, stride=2, padding=1)

        self.block1 = ResidualBlock(64)
        self.block2 = ResidualBlock(64)
        self.globalpool = nn.AvgPool2d(kernelsize=7)

        self.fc = nn.Linear(64, numclasses)

    def call(self, x):
        x = self.pool(nn.relu(self.bn1(self.conv1(x))))
        x = self.block1(x)
        x = self.block2(x)
        x = self.globalpool(x)

        x = x.reshape(x.shape[0], -1)
        return self.fc(x)

model = SimpleResNet(numclasses=10)
dummyinput = mx.random.normal((1, 3, 224, 224))

output = model(dummyinput)
print(f"Output shape: {output.shape}")

MLX Server: OpenAI-Compatible API

MLX-LM provides a server compatible with the OpenAI API:

# Start server
mlxlm.server --model mlx-community/Llama-3.2-3B-Instruct-4bit --port 8080

from openai import OpenAI

client = OpenAI(
    baseurl="http://localhost:8080/v1",
    apikey="not-needed"

)

response = client.chat.completions.create(
    model="mlx-community/Llama-3.2-3B-Instruct-4bit",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What are the benefits of machine learning?"}
    ],
    maxtokens=300,
    temperature=0.7
)

print(response.choices[0].message.content)

Or use curl:

curl http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "mlx-community/Llama-3.2-3B-Instruct-4bit",
    "messages": [{"role": "user", "content": "Hello!"}],
    "maxtokens": 100

  }'

import mlx.core as mx
import time

def benchmarkmatmul(size, iterations=100):
    a = mx.random.normal((size, size))
    b = mx.random.normal((size, size))
    mx.eval(a, b)

    start = time.time()
    for  in range(iterations):

        c = a @ b
        mx.eval(c)
    elapsed = time.time() - start

    flops = 2  size3  iterations
    gflops = flops / elapsed / 1e9
    print(f"Size {size}x{size}: {gflops:.1f} GFLOPS ({elapsed/iterations1000:.2f} ms/iter)")


for size in [512, 1024, 2048, 4096]:
    benchmarkmatmul(size)

from mlxlm import load, generate
import time

modelstotest = [

    "mlx-community/Llama-3.2-1B-Instruct-4bit",
    "mlx-community/Llama-3.2-3B-Instruct-4bit",
    "mlx-community/Mistral-7B-Instruct-v0.3-4bit",
]

prompt = "Explain the concept of neural networks in detail:"

for modelname in modelstotest:
    model, tokenizer = load(modelname)


    messages = [{"role": "user", "content": prompt}]
    formatted = tokenizer.applychattemplate(

        messages, tokenize=False, addgenerationprompt=True

    )

    start = time.time()
    response = generate(model, tokenizer, prompt=formatted, maxtokens=200)
    elapsed = time.time() - start

    tokens = len(tokenizer.encode(response))
    print(f"{modelname}: {tokens/elapsed:.1f} tokens/sec")

import mlx.core as mx

Automatic differentiation
def f(x):
    return mx.sum(x  2)


Gradient
gradf = mx.grad(f)
x = mx.array([1.0, 2.0, 3.0])
print(f"Gradient: {gradf(x)}")


Value and gradient
valandgradf = mx.valueandgrad(f)
value, gradient = valandgradf(x)

print(f"Value: {value}, Gradient: {gradient}")

Vectorized map
def singlefn(x):
    return x  2 + 2  x + 1

batch = mx.array([[1.0], [2.0], [3.0], [4.0]])
results = mx.vmap(singlefn)(batch)

print(f"Vectorized results: {results}")

import mlx.core as mx
import mlx.nn as nn

class MultiHeadSelfAttention(nn.Module):
    def init(self, dims, numheads):
        super().init()
        self.numheads = numheads
        self.headdim = dims // numheads

        self.query = nn.Linear(dims, dims)
        self.key = nn.Linear(dims, dims)
        self.value = nn.Linear(dims, dims)
        self.out = nn.Linear(dims, dims)

    def call(self, x):
        B, T, C = x.shape

        q = self.query(x).reshape(B, T, self.numheads, self.headdim).transpose(0, 2, 1, 3)
        k = self.key(x).reshape(B, T, self.numheads, self.headdim).transpose(0, 2, 1, 3)
        v = self.value(x).reshape(B, T, self.numheads, self.headdim).transpose(0, 2, 1, 3)

        scale = self.headdim * -0.5

        attn = (q @ k.transpose(0, 1, 3, 2))  scale
        attn = mx.softmax(attn, axis=-1)

        out = (attn @ v).transpose(0, 2, 1, 3).reshape(B, T, C)
        return self.out(out)

Test
attention = MultiHeadSelfAttention(256, 8)
x = mx.random.normal((2, 10, 256))
output = attention(x)
print(f"Attention output shape: {output.shape}")

import mlx.core as mx

Good: batch evaluations
a = mx.random.normal((1000, 1000))
b = mx.random.normal((1000, 1000))
c = a @ b
d = c + mx.oneslike(c)
e = mx.sum(d)
mx.eval(e)  # Evaluate once at the end

Less optimal: evaluating too frequently
a = mx.random.normal((1000, 1000))
mx.eval(a)  # Unnecessary evaluation
b = mx.random.normal((1000, 1000))
mx.eval(b)  # Unnecessary evaluation
c = a @ b
mx.eval(c)  # Only this is needed

2. Use Quantization for Large Models

from mlxlm import load

4-bit quantization saves ~75% memory
model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

8-bit for balance between quality and memory
model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-8bit")

import mlx.core as mx

Delete arrays that are no longer needed
largearray = mx.random.normal((10000, 10000))
result = mx.sum(largearray)

mx.eval(result)
del largearray

Use metal.clearcache() when needed

mx.metal.clearcache()

4. Batch Processing for Efficiency

from mlxlm import load, generate

model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

prompts = [
    "What is Python?",
    "Explain Docker briefly.",
    "What is Kubernetes?"
]

for prompt in prompts:
    messages = [{"role": "user", "content": prompt}]
    formatted = tokenizer.applychattemplate(

        messages, tokenize=False, addgenerationprompt=True

    )
    response = generate(model, tokenizer, prompt=formatted, maxtokens=100)
    print(f"Q: {prompt}")
    print(f"A: {response}\n")

5. Monitor GPU Usage

import mlx.core as mx

Check memory usage
peakmemory = mx.metal.getpeakmemory()

activememory = mx.metal.getactivememory()
cachememory = mx.metal.getcachememory()


print(f"Peak Memory: {peakmemory / 1e9:.2f} GB")
print(f"Active Memory: {activememory / 1e9:.2f} GB")

print(f"Cache Memory: {cachememory / 1e9:.2f} GB")

MLX Ecosystem

MLX has a growing ecosystem of libraries:

| Library | Purpose |

|---------|---------|

| mlx | Core framework (array ops, neural networks, optimizers) |

| mlx-lm | Language model inference and fine-tuning |

| mlx-vlm | Vision-language models |

| mlx-audio | Audio processing and speech models |

| mlx-data | Data loading and preprocessing |

| mlx-graphs | Graph neural networks |

| mlx-image | Computer vision utilities |

Available Models in MLX Community

The MLX community on Hugging Face provides hundreds of models pre-converted to MLX format:

• Llama 3.x (1B, 3B, 8B, 70B) in various quantization levels
• Mistral/Mixtral for efficient MoE models
• Phi-3/Phi-4 small models from Microsoft
• Qwen 2.5 from Alibaba
• Gemma 2 from Google
• DeepSeek models
• Whisper for speech-to-text
• Stable Diffusion for image generation

Browse models at: https://huggingface.co/mlx-community

Conclusion

MLX is an incredibly powerful framework for running machine learning workloads on Apple Silicon. With its familiar API, lazy evaluation support, and unified memory architecture, MLX enables developers to fully leverage the capabilities of Apple chips.

Key takeaways:

MLX is ideal for prototyping, local development, research, and even deployment on Apple devices. With an ever-growing number of models available in MLX format, this framework has become the primary choice for ML on Apple Silicon.