Introduction to Apple Silicon MLX using Python
1.1 What is MLX?
MLX is a machine learning framework by Apple designed specifically for Apple Silicon (M# chips) created by Apple’s machine learning research team. It is useful for Data Science, Machine Learning, LLMs, and Deep Learning being similar to Pytorch which is based on Numpy. What makes MLX easily adoptable is that it is possible to code with C++ and Python.
If you’ve ever used NumPy for numerical computing or PyTorch for deep learning, MLX will feel familiar. It provides the same kind of tensor (multi-dimensional array) operations but is specifically designed to take advantage of the unified memory architecture in Apple’s M-series chips (M1, M2, M3, M4, and beyond).
In regards to software engineering, you can think of MLX as the “Node.js of machine learning frameworks”. It’s designed to be approachable, efficient, and flexible, removing many of the traditional pain points of ML development.
1.2 Why MLX Matters for Artificial Intelligence
Apple Silicon has provided “unified” memory in their computers which means the GPU and CPU have access to the full amount of system RAM. This is an advantage over traditional discrete GPUs by AMD/NVidia on a windows computer. Since the GPU is responsible for performing matrix operations for anything Ai, it becomes limited by its own onboard memory where Ai models are pre-trained, loaded and inferenced.
Alternatively, matrix operations can be performed on the CPU with severely reduced performance. This said, Nvidia still has a stronghold in the computing industry due to their hardware strength but more so for corporations instead of consumers due to their equipment’s inflated pricing.
Fortunately, for the average consumer like you and me, more Ai models are being open-sourced from distillations, fine-tuning and GGUF-MLX conversions giving Mac’s more efficiency to run them over discrete GPUs.
As an engineer, you might wonder why you should care about yet another ML framework. Need more reasons why MLX is particularly relevant?
Familiar API Design
MLX’s Python API closely follows NumPy, which means the learning curve is gentle if you’ve done any scientific Python work. The higher-level neural network API (mlx.nn) and optimizer API (mlx.optimizers) follow PyTorch conventions.
Unified Memory
Unlike x86 Architecture computers, Apple silicon has a unified memory architecture, so you don’t have to move data between CPU and GPU memory. Unified memory means sharing the same memory pool. MLX is designed to optimize this, meaning you don’t need to think about device placement the way you do in PyTorch or TensorFlow.
Lazy Evaluation
Operations in MLX are lazy; they build a computation graph but don’t execute until you need the results. This enables automatic optimizations and makes it easier to reason about complex computations.
Composable Transformations
MLX provides powerful function transformations for automatic differentiation (computing gradients), automatic vectorization (batching), and computation graph optimization. These can be freely composed, so grad(vmap(grad(fn))) is perfectly valid.
1.3 MLX Key Features
Welcome to this introductory MLX and basic reference source.
| Feature | Description |
|---|---|
| NumPy-like API | Familiar array operations with mlx.core |
| PyTorch-like NN | Neural network layers via mlx.nn |
| Lazy Computation | Operations deferred until results needed |
| Unified Memory | No data copies between CPU and GPU |
| Auto-differentiation | grad(), value_and_grad() for training |
| Auto-vectorization | vmap() for automatic batching |
| Compilation | compile() for graph optimization |
| GPU Acceleration | Metal backend for Apple silicon GPU |
| Multi-device | CPU and GPU support with stream control |
1.4 MLX vs Other Frameworks
| Aspect | MLX | PyTorch | JAX | TensorFlow |
|---|---|---|---|---|
| Memory Model | Unified | Separate CPU/GPU | Separate CPU/GPU | Separate CPU/GPU |
| Evaluation | Lazy | Eager (default) | Lazy | Eager (default) |
| Graph | Dynamic | Dynamic | Static (via jit) | Static (via tf.function) |
| Apple Silicon | Native | Via MPS | Limited | Limited |
| API Style | NumPy-like | Unique | NumPy-like | Unique |
| Transformations | Composable | Limited | Composable | Limited |
1.5 The MLX Ecosystem
MLX isn’t just a single library. It has a growing ecosystem:
- mlx – The core array framework (what this article covers)
- mlx-lm – A package for LLM text generation and fine-tuning
- mlx-data – Efficient data loading and preprocessing
- mlx-vlm – Vision Language Model support
- mlx-whisper – Speech recognition
- mlx-examples – A rich collection of example implementations
- mlx-community (HuggingFace) – Pre-converted model weights
The official resources you’ll want to bookmark:
- Documentation: https://ml-explore.github.io/mlx/build/html/index.html
- GitHub: https://github.com/ml-explore/mlx
- Examples: https://github.com/ml-explore/mlx-examples
- HuggingFace Community: https://huggingface.co/mlx-community
2: Setting Up Your Development Environment
2.1 System requirements
Before installing MLX, verify your system meets these requirements:
- Hardware: Mac with Apple silicon (M1 or later)
- OS: macOS 14.0 (Sonoma) or later
- Python: 3.10 or later (must be a native ARM64 Python, not x86 via Rosetta)
To check your Mac OS:
sw_vers
# ProductVersion should be 14.0 or higher2.2 Installing MLX
Use pip install in the command line:
pip install mlxAs with any Python project, it’s better to use a virtual environment:
# Using venv (built into Python)
python -m venv mlx_env
source mlx_env/bin/activate
# Install MLX
pip install mlxUsing the newer and faster package installer “uv”:
# Install uv (if not already installed)
curl -LsSf https://astral.sh/uv/install.sh | sh
# Create a virtual environment with uv
uv python venv mlx_env
# Activate the environment
source mlx_env/bin/activate
# Install MLX
uv pip install mlx2.3 Verify the installation (optional)
Create a file called verify_install.py and run it with python verify_install.py:
import mlx.core as mx
# Check MLX version
print(f"MLX version: {mx.__version__}")
# Check Metal GPU availability
print(f"Metal available: {mx.metal.is_available()}")
# Get device info
if mx.metal.is_available():
info = mx.metal.device_info()
print(f"GPU: {info.get('device_name', 'Unknown')}")
# Quick computation test
a = mx.array([1.0, 2.0, 3.0])
b = mx.array([4.0, 5.0, 6.0])
c = a + b
mx.eval(c) # Force evaluation
print(f"Computation test: {a} + {b} = {c}")
# Test GPU computation
d = mx.add(a, b, stream=mx.gpu)
mx.eval(d)
print(f"GPU computation test passed: {d}")
print("\nMLX is installed and working correctly!")
#Expected outout
#MLX version: 0.x.x
#Metal available: True
#GPU: Apple M1 (or M2, M3, M4, etc.)
#Computation test: array([1, 2, 3], dtype=float32) + array([4, 5, 6], #dtype=float32) = array([5, 7, 9], dtype=float32)
#GPU computation test passed: array([5, 7, 9], dtype=float32)
#MLX is installed and working correctly!2.5 Installing Supporting Packages
For this training guide, you’ll also want to install supporting packages:
pip install numpy # For data interop
pip install matplotlib # For visualization
pip install datasets # HuggingFace datasets
pip install transformers # Tokenizers and model configs
pip install sentencepiece # Tokenization
pip install tqdm # Progress bars
pip install mlx-lm # LLM package (later chapters)2.6 Write your first program – Hello World
Instead of boring documentation, let’s develop a working program so you can have immediate tangible results. This is to write a “Hello World” version of Machine Learning called linear regression. This will give you a feel for the MLX workflow:
import mlx.core as mx
# Generate some synthetic data: y = 3x + 2 + noise
mx.random.seed(42)
X = mx.random.normal((100, 1))
noise = mx.random.normal((100, 1)) * 0.5
y = 3.0 * X + 2.0 + noise
# Initialize parameters as scalars
weight = mx.array(0.0)
bias = mx.array(0.0)
# Define the model
def predict(x, w, b):
return x * w + b
# Define the loss function (mean squared error)
def loss_fn(w, b, x, y):
pred = predict(x, w, b)
return mx.mean((pred - y) ** 2)
# Get the gradient function -- differentiate w.r.t. w AND b
loss_and_grad = mx.value_and_grad(loss_fn, argnums=[0, 1])
# Training loop
for step in range(200):
loss, (grad_w, grad_b) = loss_and_grad(weight, bias, X, y)
weight = weight - 0.1 * grad_w
bias = bias - 0.1 * grad_b
mx.eval(loss, weight, bias)
if step % 50 == 0:
print(f"Step {step}: loss = {loss.item():.4f}, "
f"weight = {weight.item():.4f}, bias = {bias.item():.4f}")
print(f"\nFinal: weight = {weight.item():.4f} (target: 3.0), "
f"bias = {bias.item():.4f} (target: 2.0)")When you run this, you should see the weight converge toward 3.0 and the bias toward 2.0. Don’t worry about understanding every line yet, we’ll cover each concept in detail in the coming chapters.
2.7 Development Tools
For the hello world example, you only need Terminal, no web browser, no IDEs, no simulator environments. However, you can use helpful tools beyond the command line. Here are a few tools that will enhance your MLX development experience:
Jupyter Notebooks
Great for interactive experimentation that can be executed by line with immediate output rendered below the code.
VS Code
Excellent Python support with the Python extension. The Jupyter extension also works well for notebook-style development.
Terminal
Since MLX scripts are just Python, any terminal works.
pip install jupyter
jupyter notebook2.8 Common Installation Issues
“No matching distribution found”
This usually means you’re running x86 Python via Rosetta.
“MLX requires macOS 14.0 or later”
Update your macOS to Sonoma (14.0) or later.
“Import errors after installation“
Make sure you’re in the correct virtual environment:
which python
# Should point to your virtual environment
pip list | grep mlx
# Should show mlx with version number2.9 Additional packages
In order to to complete the following tutorials, create the requirements.txt file and install the necessary packages:
mlx
numpy
matplotlib
datasets
transformers
sentencepiece
tqdm
mlx-lm
safetensors3: Arrays – The Foundation of MLX
3.1 What is an Array?
In MLX, the array is the fundamental data structure similar to how JavaScript has Array and Python has list, but designed for numerical computing. An MLX array is a multi-dimensional container of numerical data that lives in unified memory and can be operated on by both the CPU and GPU without copying. If you think of a spreadsheet as a 2D grid of numbers, an MLX array generalizes this concept to any number of dimensions.
import mlx.core as mx
# A scalar (0-dimensional)
scalar = mx.array(42)
print(scalar) # array(42, dtype=int32)
print(scalar.ndim) # 0
print(scalar.shape) # ()
# A vector (1-dimensional)
vector = mx.array([1, 2, 3, 4])
print(vector) # array([1, 2, 3, 4], dtype=int32)
print(vector.ndim) # 1
print(vector.shape) # (4,)
# A matrix (2-dimensional)
matrix = mx.array([[1, 2], [3, 4], [5, 6]])
print(matrix) # array([[1, 2], [3, 4], [5, 6]], dtype=int32)
print(matrix.ndim) # 2
print(matrix.shape) # (3, 2)
# A 3D tensor (think: batch of matrices)
tensor3d = mx.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
print(tensor3d.ndim) # 3
print(tensor3d.shape) # (2, 2, 2)3.2 Data Types
MLX arrays are typed. The type determines how much memory each element uses and what operations are valid:
# Integer types
a_int32 = mx.array([1, 2, 3]) # int32 (default for integers)
a_int8 = mx.array([1, 2, 3], dtype=mx.int8)
a_int16 = mx.array([1, 2, 3], dtype=mx.int16)
a_int64 = mx.array([1, 2, 3], dtype=mx.int64)
a_uint8 = mx.array([1, 2, 3], dtype=mx.uint8)
a_uint16 = mx.array([1, 2, 3], dtype=mx.uint16)
a_uint32 = mx.array([1, 2, 3], dtype=mx.uint32)
# Floating point types
b_float32 = mx.array([1.0, 2.0]) # float32 (default for floats)
b_float16 = mx.array([1.0, 2.0], dtype=mx.float16)
b_bfloat16 = mx.array([1.0, 2.0], dtype=mx.bfloat16)
# Boolean
c_bool = mx.array([True, False, True]) # bool
# Complex
d_complex = mx.array([1+2j, 3+4j]) # complex64
# Check the dtype
print(a_int32.dtype) # int32
print(b_float32.dtype) # float32Software engineer analogy: Think of dtypes like choosing between int, float, and boolean in JavaScript, but with more granular control over precision and memory usage.
Why dtypes matter: Using lower precision (like float16 instead of float32) halves memory usage and can speed up computation. For training language models, float16 or bfloat16 are common choices. For inference, quantized types (4-bit, 8-bit) can further reduce memory.
3.3 Creating Arrays
MLX provides many ways to create arrays:
import mlx.core as mx
# From Python lists
a = mx.array([1, 2, 3])
# From nested lists (multi-dimensional)
b = mx.array([[1.0, 2.0], [3.0, 4.0]])
# Zeros and ones
zeros = mx.zeros((3, 4)) # 3x4 matrix of zeros
ones = mx.ones((2, 3, 4)) # 2x3x4 tensor of ones
# Full (constant value)
full = mx.full((3, 3), 7.0) # 3x3 matrix filled with 7.0
# Identity matrix
eye = mx.eye(4) # 4x4 identity matrix
# Sequences
arange = mx.arange(0, 10, 2) # [0, 2, 4, 6, 8]
linspace = mx.linspace(0, 1, 5) # [0.0, 0.25, 0.5, 0.75, 1.0]
# Random arrays
mx.random.seed(42)
rand_normal = mx.random.normal((3, 3)) # Normal distribution
rand_uniform = mx.random.uniform(shape=(3, 3)) # Uniform [0, 1)
rand_bernoulli = mx.random.bernoulli(shape=(3, 3)) # Bernoulli (0 or 1)
rand_int = mx.random.randint(0, 10, (3, 3)) # Random integers
# Like existing arrays (same shape/dtype)
x = mx.array([1.0, 2.0, 3.0])
zeros_like = mx.zeros_like(x) # [0.0, 0.0, 0.0]
ones_like = mx.ones_like(x) # [1.0, 1.0, 1.0]3.4 Array Properties
Every array has useful properties for inspecting its shape and metadata:
a = mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
print(a.shape) # (2, 3) dimensions
print(a.ndim) # 2 number of dimensions
print(a.size) # 6 total number of elements
print(a.dtype) # float32 data type
print(a.nbytes) # 24 bytes used (6 elements * 4 bytes/float32)
print(a.itemsize) # 4 bytes per element
print(a.T) # Transposed: shape (3, 2)3.5 Indexing and Slicing
MLX indexing follows NumPy conventions, which should feel familiar:
a = mx.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# Single element
print(a[0, 0]) # 1
print(a[1, 2]) # 6
# Row and column slicing
print(a[0]) # First row: [1, 2, 3]
print(a[:, 0]) # First column: [1, 4, 7]
print(a[0:2]) # First two rows
print(a[0:2, 1:3]) # Sub-matrix
# Negative indexing
print(a[-1]) # Last row: [7, 8, 9]
print(a[-1, -1]) # Last element: 9
# Step slicing
b = mx.arange(10)
print(b[::2]) # [0, 2, 4, 6, 8] every other element
print(b[1::2]) # [1, 3, 5, 7, 9] odd indices
print(b[::-1]) # [9, 8, 7, 6, 5, 4, 3, 2, 1, 0] reversed
# Boolean indexing
mask = a > 5
print(mask) # [[False, False, False], [False, False, True], [True, True, True]]
# Integer array indexing
indices = mx.array([0, 2])
print(a[indices]) # First and third rows3.6 Reshaping and Broadcasting
Reshaping changes the view of an array without changing its data. Broadcasting allows operations between arrays of different shapes.
Web developer analogy: Broadcasting is like a CSS flexbox. It automatically handles size mismatches in predictable ways. A (1, 3) array plus a (3, 1) array “stretches” to produce a (3, 3) result.
# Reshaping
a = mx.arange(12) # shape: (12,)
b = a.reshape(3, 4) # shape: (3, 4)
c = a.reshape(2, 2, 3) # shape: (2, 2, 3)
# Use -1 to auto-infer a dimension
d = a.reshape(3, -1) # shape: (3, 4) MLX infers the 4
# Flatten
e = c.reshape(-1) # Back to shape (12,)
f = c.flatten() # Same as reshape(-1)
# Squeeze and expand dims
g = mx.array([[1, 2, 3]]) # shape: (1, 3)
h = g.squeeze() # shape: (3,)
i = mx.expand_dims(h, 0) # shape: (1, 3)
j = mx.expand_dims(h, 1) # shape: (3, 1)
# Broadcasting -- operations on different shapes
x = mx.array([[1, 2, 3]]) # shape: (1, 3)
y = mx.array([[1], [2], [3]]) # shape: (3, 1)
z = x + y # shape: (3, 3) via broadcasting
# Broadcast explicitly
k = mx.broadcast_to(mx.array([1, 2, 3]), (3, 3)) # shape: (3, 3)3.7 Mathematical Operations
MLX provides a comprehensive set of mathematical operations:
import mlx.core as mx
a = mx.array([1.0, 2.0, 3.0])
b = mx.array([4.0, 5.0, 6.0])
# Element-wise arithmetic
print(a + b) # [5.0, 7.0, 9.0]
print(a - b) # [-3.0, -3.0, -3.0]
print(a * b) # [4.0, 10.0, 18.0]
print(a / b) # [0.25, 0.4, 0.5]
print(a ** 2) # [1.0, 4.0, 9.0]
# Using explicit functions
print(mx.add(a, b))
print(mx.multiply(a, b))
print(mx.divide(a, b))
print(mx.power(a, 2))
# Unary operations
print(mx.sqrt(a)) # [1.0, 1.414, 1.732]
print(mx.exp(a)) # Exponential
print(mx.log(a)) # Natural logarithm
print(mx.abs(mx.array([-1, -2]))) # [1, 2]
print(mx.sign(mx.array([-3, 0, 5]))) # [-1, 0, 1]
# Trigonometric
print(mx.sin(a))
print(mx.cos(a))
print(mx.tan(a))
# Rounding
print(mx.round(mx.array([1.4, 2.5, 3.6]))) # [1, 2, 4]
print(mx.ceil(mx.array([1.1, 2.9]))) # [2, 3]
print(mx.floor(mx.array([1.1, 2.9]))) # [1, 2]
# Reductions (aggregations)
c = mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
print(mx.sum(c)) # 21.0 sum of all elements
print(mx.sum(c, axis=0)) # [5.0, 7.0, 9.0] sum along rows
print(mx.sum(c, axis=1)) # [6.0, 15.0] sum along columns
print(mx.mean(c)) # 3.5
print(mx.max(c)) # 6.0
print(mx.min(c)) # 1.0
print(mx.prod(c)) # 720.0
print(mx.var(c)) # Variance
print(mx.std(c)) # Standard deviation
# Sorting
d = mx.array([3, 1, 4, 1, 5, 9, 2, 6])
print(mx.sort(d)) # [1, 1, 2, 3, 4, 5, 6, 9]
print(mx.argsort(d)) # Indices that would sort the array
# Comparison
print(mx.equal(a, b)) # Element-wise equality
print(mx.greater(a, b)) # Element-wise >
print(mx.less(a, b)) # Element-wise <
print(mx.maximum(a, b)) # Element-wise max
print(mx.minimum(a, b)) # Element-wise min3.8 Matrix Operations
Matrix operations are central to neural networks. MLX provides optimized implementations:
import mlx.core as mx
# Matrix multiplication
A = mx.random.normal((3, 4))
B = mx.random.normal((4, 5))
C = mx.matmul(A, B) # shape: (3, 5)
# Or using the @ operator
C = A @ B # Same as matmul
# Dot product
v1 = mx.array([1.0, 2.0, 3.0])
v2 = mx.array([4.0, 5.0, 6.0])
print(mx.inner(v1, v2)) # 32.0
# Outer product
print(mx.outer(v1, v2)) # shape: (3, 3)
# Transpose
M = mx.random.normal((3, 4))
print(M.T) # shape: (4, 3)
print(mx.transpose(M)) # Same
# Matrix with specific dimensions
print(mx.transpose(M, (1, 0))) # Explicit axis order
# Einsum (Einstein summation) powerful for complex tensor operations
a = mx.random.normal((2, 3))
b = mx.random.normal((3, 4))
c = mx.einsum("ij,jk->ik", a, b) # Matrix multiplication via einsum
# Trace, diagonal
square = mx.array([[1, 2], [3, 4]])
print(mx.trace(square)) # 5.0
print(mx.diag(mx.array([1, 2, 3]))) # Diagonal matrix3.9 Concatenation, Stacking, and Splitting
import mlx.core as mx
a = mx.array([[1, 2], [3, 4]])
b = mx.array([[5, 6], [7, 8]])
# Concatenate along existing axis
print(mx.concatenate([a, b], axis=0)) # Vertical: shape (4, 2)
print(mx.concatenate([a, b], axis=1)) # Horizontal: shape (2, 4)
# Stack along new axis
print(mx.stack([a, b], axis=0)) # shape: (2, 2, 2)
# Split
c = mx.arange(12)
print(mx.split(c, 3)) # Split into 3 equal parts
print(mx.split(c, [3, 7])) # Split at indices 3 and 7
# Tile (repeat)
d = mx.array([[1, 2], [3, 4]])
print(mx.tile(d, (2, 3))) # Repeat 2x vertically, 3x horizontally3.10 Saving and Loading Arrays
import mlx.core as mx
a = mx.random.normal((100, 50))
b = mx.random.normal((100,))
# Save a single array
mx.save("weights_a.npy", a)
# Save multiple arrays
mx.savez("weights.npz", {"a": a, "b": b})
# Save compressed
mx.savez_compressed("weights_compressed.npz", {"a": a, "b": b})
# Save in safetensors format (recommended for ML models)
mx.save_safetensors("model.safetensors", {"weight": a, "bias": b})
# Load arrays
loaded_a = mx.load("weights_a.npy") # Single array
loaded_dict = mx.load("weights.npz") # Dict of arrays
# Interop with NumPy
import numpy as np
np_array = np.array(a) # MLX -> NumPy (triggers eval)
mlx_array = mx.array(np_array) # NumPy -> MLX4: Lazy Evaluation and Computation Graphs
4.1 Understanding Lazy Evaluation
This is one of the most important concepts in MLX, and it differs significantly from how PyTorch works (by default). In MLX, operations are lazy when you write c = a + b, no actual addition happens. Instead, MLX records the operation in a computation graph and defers the actual computation until you explicitly request the result.
Think of it like a recipe: writing c = a + b adds a step to the recipe, but doesn’t cook anything. You only start cooking when you call mx.eval().
import mlx.core as mx
a = mx.array([1.0, 2.0, 3.0])
b = mx.array([4.0, 5.0, 6.0])
# No computation happens here!
c = a + b
# The computation only happens when we evaluate
mx.eval(c)
print(c) # Now the result exists: array([5, 7, 9], dtype=float32)4.2 When Does Evaluation Happen Automatically?
While mx.eval() is the explicit way to trigger computation, there are implicit triggers you should be aware of:
import mlx.core as mx
a = mx.array([1.0, 2.0])
b = mx.array([3.0, 4.0])
c = a + b
# These all trigger implicit evaluation:
# 1. Printing
print(c) # Evaluates, then prints
# 2. Converting to NumPy
import numpy as np
np_array = np.array(c) # Evaluates, then converts
# 3. Using .item() on a scalar
d = mx.sum(c)
print(d.item()) # Evaluates, then returns Python float
# 4. Saving to disk
mx.save("output.npy", c) # Evaluates, then saves
# 5. Using in control flow (be careful!)
if d > 5: # Evaluates d to check the condition
print("Greater than 5")4.3 Why Lazy Evaluation Matters
Lazy evaluation provides several benefits such as graph optimization, memory efficiency, and composable transformations. Let’s explain them with examples.
Graph Optimization
Since MLX sees the full computation graph before executing, it can optimize it (e.g., fusing operations, eliminating dead code).
def process(x):
a = expensive_op(x) # This might be eliminated if unused
b = cheap_op(x)
return b # Only b is returned, a is never computed
result = process(input_data)
mx.eval(result) # expensive_op is never executed!Memory Efficiency
No intermediate results are stored until needed. You can create large models without immediately using memory:
model = BigModel() # Graph is built, but no weights allocated yet
model.load_weights("weights_fp16.safetensors") # Only fp16 memory usedComposable Transformations
Lazy evaluation makes it possible to transform the computation graph before executing it. This is how grad() (automatic differentiation) and vmap() (vectorization) work, rewriting the graph.
4.4 Best Practices for Evaluation
The right granularity for mx.eval() matters. Evaluate at each iteration of your outer training loop. This gives MLX enough graph to optimize while keeping overhead manageable.
# BAD: Too many small evaluations (high overhead)
for i in range(1000):
result = mx.add(a, b)
mx.eval(result) # 1000 separate graph evaluations!
# BAD: Too few evaluations (huge graph, growing overhead)
results = []
for i in range(1000):
results.append(mx.add(a, b))
mx.eval(results) # One huge graph
# GOOD: Evaluate at natural boundaries (e.g., per training step)
for batch in dataset:
loss, grads = loss_and_grad_fn(model, batch)
optimizer.update(model, grads)
mx.eval(loss, model.parameters()) # One eval per step4.5 Computation Graphs Explained
A computation graph is a directed acyclic graph (DAG) where:
- Nodes are operations (add, multiply, etc.) or array values
- Edges represent data flow
When you call mx.eval(f), MLX traverses this graph from f backward, computing only what’s needed to produce f.
Input: a, b
|
+-- [add] --> c = a + b
| |
+-- [multiply] --> d = a * b
|
[multiply] --> e = c * d
|
[sum] --> f = sum(e)4.6 Gradient Computation and Graphs
The lazy evaluation model is what makes MLX’s grad() transformation work. When you apply grad() to a function, MLX does 3 things (trace, reverse, returns):
- traces the forward computation graph
- reverses it to build a gradient computation graph
- returns a new function that builds both graphs
import mlx.core as mx
def f(x):
return mx.sum(x ** 2)
# Get gradient function
grad_fn = mx.grad(f)
x = mx.array([1.0, 2.0, 3.0])
gradient = grad_fn(x)
mx.eval(gradient)
print(gradient) # array([2, 4, 6], dtype=float32) which is 2*x4.7 Checkpointing for Memory
In machine learning, checkpointing for memory is a technique that trades increased computation for reduced memory usage. Instead of storing all intermediate activation tensors from the forward pass to compute gradients during back-propagation, the system stores only a selected subset of these values, called checkpoint.
import mlx.core as mx
def expensive_block(x):
# Many operations, large intermediate tensors
x = mx.matmul(x, w1)
x = mx.relu(x)
x = mx.matmul(x, w2)
x = mx.relu(x)
return x
# Without checkpointing: all intermediates stored
def forward(x):
return expensive_block(x)
# With checkpointing: intermediates discarded, recomputed during backward
def forward_checkpointed(x):
return mx.checkpoint(expensive_block)(x)5: Function Transformations
5.1 Automatic Differentiation with grad()
Automatic differentiation (autodiff) is the backbone of neural network training. It computes the gradient (derivative) of a function automatically, without you having to derive calculus formulas by hand.
import mlx.core as mx
# Simple example: derivative of sin(x) is cos(x)
x = mx.array(0.0)
# Forward: sin(0) = 0
print(mx.sin(x)) # array(0, dtype=float32)
# Gradient: cos(0) = 1
print(mx.grad(mx.sin)(x)) # array(1, dtype=float32)
# Second derivative: -sin(0) = 0
print(mx.grad(mx.grad(mx.sin))(x)) # array(-0, dtype=float32)5.2 grad() with Multiple Arguments
When your function has multiple arguments, you can control which ones get differentiated:
import mlx.core as mx
def loss_fn(weights, bias, inputs, targets):
predictions = inputs @ weights + bias
return mx.mean((predictions - targets) ** 2)
# By default, grad differentiates the FIRST argument
grad_fn = mx.grad(loss_fn)
# This gives d(loss)/d(weights)
# To differentiate specific arguments, use argnums
grad_fn = mx.grad(loss_fn, argnums=[0, 1])
# This returns a tuple: (d(loss)/d(weights), d(loss)/d(bias))5.3 value_and_grad() Get Both
Often you want both the function value and the gradient (e.g., to track loss during training):
import mlx.core as mx
def loss_fn(weights, inputs, targets):
predictions = inputs @ weights
return mx.mean((predictions - targets) ** 2)
# Get both value and gradient simultaneously
loss_and_grad = mx.value_and_grad(loss_fn)
weights = mx.random.normal((3, 1))
inputs = mx.random.normal((10, 3))
targets = mx.random.normal((10, 1))
loss, grads = loss_and_grad(weights, inputs, targets)
mx.eval(loss, grads)
print(f"Loss: {loss.item():.4f}")
print(f"Gradients shape: {grads.shape}")5.4 Higher-Order Gradients
MLX supports nested grad() calls for higher-order derivatives:
import mlx.core as mx
# f(x) = x^3
def f(x):
return x ** 3
# f'(x) = 3x^2
df = mx.grad(f)
# f''(x) = 6x
d2f = mx.grad(df)
# f'''(x) = 6
d3f = mx.grad(d2f)
x = mx.array(2.0)
print(f"f(2) = {f(x).item()}") # 8
print(f"f'(2) = {df(x).item()}") # 12
print(f"f''(2) = {d2f(x).item()}") # 24
print(f"f'''(2) = {d3f(x).item()}") # 65.5 Vectorization with vmap()
vmap() (vectorized map) automatically transforms a function that operates on a single example into one that operates on a batch. The in_axes parameter specifies which axis to vectorize over for each argument:
Nonemeans don’t vectorize (use the same value for all)0means vectorize along the first dimension
import mlx.core as mx
# A function that works on a single vector
def predict(x, weight, bias):
return mx.matmul(x, weight) + bias
x = mx.array([1.0, 2.0, 3.0]) # Single input
w = mx.array([0.5, 0.3, 0.2])
b = mx.array(0.1)
# Apply to one example
result = predict(x, w, b)
# Now batch it with vmap -- apply to a batch of inputs
batch_x = mx.random.normal((32, 3)) # 32 inputs, each of size 3
# vmap over the first (0th) argument
batched_predict = mx.vmap(predict, in_axes=(0, None, None))
batch_results = batched_predict(batch_x, w, b)
print(batch_results.shape) # (32,)5.6 Jacobian-Vector Products (jvp) and Vector-Jacobian Products (vjp)
For advanced use cases, MLX provides forward-mode (jvp) and reverse-mode (vjp) autodiff primitives:
import mlx.core as mx
def f(x):
return mx.stack([mx.sin(x[0]), mx.cos(x[1])])
x = mx.array([1.0, 2.0])
v = mx.array([1.0, 0.0]) # Direction to compute Jacobian-vector product
# Forward-mode: J * v (Jacobian times vector)
# Useful when output dimension > input dimension
result = mx.jvp(f, (x,), (v,))
print(result) # (f(x), J(x) @ v)
# Reverse-mode: v^T * J (vector times Jacobian)
# This is what grad() uses internally
# Useful when input dimension > output dimension
result = mx.vjp(f, (x,))
print(result) # (f(x), vjp_fn) where vjp_fn(cotangent) gives gradients5.7 Compilation with compile()
MLX’s compile() function optimizes the computation graph by fusing operations and reducing memory overhead. Important notes about compile():
- Works well combined with training loops where shapes are consistent
- The first call compiles the graph (has overhead)
- Subsequent calls with the same shapes/dtypes reuse the compiled graph
- If shapes change, a recompilation is triggered
import mlx.core as mx
@mx.compile
def compiled_function(x, y):
return mx.sum(mx.sin(x) * mx.cos(y))
x = mx.random.normal((1000, 1000))
y = mx.random.normal((1000, 1000))
# First call: compilation (slightly slower)
result = compiled_function(x, y)
mx.eval(result)
# Subsequent calls: uses compiled graph (faster)
result = compiled_function(x, y)
mx.eval(result)5.8 Combining Transformations
The real power of MLX is that transformations compose freely:
import mlx.core as mx
def model_fn(params, x):
w, b = params
return mx.sum(mx.relu(x @ w + b))
# Compose: vectorize, then differentiate
batched_grad = mx.vmap(mx.grad(model_fn), in_axes=(None, 0))
# Compose: differentiate, then differentiate again
hessian_diag = mx.grad(mx.grad(model_fn))
# Compose: compile the gradient function
compiled_grad = mx.compile(mx.grad(model_fn))5.9 nn.value_and_grad() for Neural Networks
For the common pattern of training neural networks, MLX provides a specialized value_and_grad in the nn module that works directly with nn.Module objects.
This pattern nn.value_and_grad + optimizer.update + mx.eval is the core training loop you’ll use throughout this guide. Think of it like a React component’s lifecycle: loss_and_grad computes (render), optimizer.update applies changes (commit), and mx.eval triggers the actual computation (DOM update).
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
# Define a model
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(10, 20)
self.linear2 = nn.Linear(20, 1)
def __call__(self, x):
x = mx.relu(self.linear1(x))
return self.linear2(x)
model = MyModel()
# Loss function takes model and batch
def loss_fn(model, X, y):
pred = model(X)
return mx.mean((pred - y) ** 2)
# Get value and grad with respect to model parameters
loss_and_grad = nn.value_and_grad(model, loss_fn)
# Training loop
optimizer = optim.Adam(learning_rate=1e-3)
for batch_x, batch_y in dataset:
loss, grads = loss_and_grad(model, batch_x, batch_y)
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state, loss)5.10 Custom Functions
For operations that need custom gradient computation
import mlx.core as mx
@mx.custom_function
def my_custom_op(x):
# Forward pass
return mx.exp(x)
@my_custom_op.vjp
def my_custom_op_vjp(primals, cotangent):
x, = primals
# The gradient of exp(x) is exp(x) * cotangent
return mx.exp(x) * cotangent,6: NLP Tokenization and Embeddings
6.1 Why Text Needs Processing
Neural networks operate on numbers, not text. Before we can train a model to understand language, we need to convert text into numerical representations. This process involves three key steps:
- Tokenization: Breaking text into pieces (tokens)
- Numericalization: Mapping tokens to integer IDs
- Embedding: Converting integer IDs to dense vector representations
As a web developer, you can think of this as similar to how a browser processes HTML: raw text is parsed into a DOM tree (tokenization), elements get assigned identifiers (numericalization), and then the rendering engine uses those to compute layouts and paint pixels (embedding/forward pass).
You can think of this as similar to how a browser processes HTML: raw text is parsed into a DOM tree (tokenization), elements get assigned identifiers (numericalization), and then the rendering engine uses those to compute layouts and paint pixels (embedding/forward pass).
6.2 Tokenization Strategies
The beginnings as character-level tokenization: each character is a token.
Pros: Tiny vocabulary, handles any text
Cons: No semantic meaning, long sequences, hard to learn patterns
text = "Hello, world!"
tokens = list(text)
print(tokens) # ['H', 'e', 'l', 'l', 'o', ',', ' ', 'w', 'o', 'r', 'l', 'd', '!']
# Build vocabulary
vocab = sorted(set(text))
char_to_id = {ch: i for i, ch in enumerate(vocab)}
id_to_char = {i: ch for ch, i in char_to_id.items()}
vocab_size = len(vocab)
print(f"Vocabulary: {vocab}")
print(f"Vocabulary size: {vocab_size}")
print(f"'H' -> {char_to_id['H']}")
# Encode text
encoded = [char_to_id[ch] for ch in text]
print(f"Encoded: {encoded}")
# Decode back
decoded = ''.join(id_to_char[i] for i in encoded)
print(f"Decoded: {decoded}")Then there’s word-level tokenization split on whitespace and punctuation.
Pros: More semantic meaning per token, shorter sequences
Cons: Huge vocabulary, can’t handle unknown words, different forms (“run”/”running”) are separate tokens
import re
text = "The cat sat on the mat."
# Simple split
words = text.lower().split()
print(words) # ['the', 'cat', 'sat', 'on', 'the', 'mat.']
# Better: regex split keeping punctuation
words = re.findall(r'\w+|\S', text.lower())
print(words) # ['the', 'cat', 'sat', 'on', 'the', 'mat', '.']
# Build vocabulary
vocab = sorted(set(words))
word_to_id = {w: i for i, w in enumerate(vocab)}
vocab_size = len(vocab)The more recent standard is subword tokenization; a balance between character and word level. Common algorithms include BPE (Byte Pair Encoding), WordPiece, and SentencePiece.
Let’s implement a simple BPE tokenizer:
"""Simple BPE (Byte Pair Encoding) Tokenizer implementation."""
class SimpleBPETokenizer:
"""A minimal BPE tokenizer for educational purposes."""
def __init__(self):
self.merges = {} # pair -> merged token
self.vocab = {} # token -> id
self.id_to_token = {} # id -> token
def train(self, text, vocab_size=256, num_merges=100):
"""Train BPE on text to learn merge rules."""
# Start with character-level tokens
tokens = list(text)
vocab = sorted(set(tokens))
self.vocab = {t: i for i, t in enumerate(vocab)}
# Count pairs
for i in range(num_merges):
pairs = {}
for j in range(len(tokens) - 1):
pair = (tokens[j], tokens[j + 1])
pairs[pair] = pairs.get(pair, 0) + 1
if not pairs:
break
# Find most common pair
best_pair = max(pairs, key=pairs.get)
new_token = best_pair[0] + best_pair[1]
# Record merge
self.merges[best_pair] = new_token
if new_token not in self.vocab:
self.vocab[new_token] = len(self.vocab)
# Apply merge to token list
new_tokens = []
j = 0
while j < len(tokens):
if (j < len(tokens) - 1 and
tokens[j] == best_pair[0] and
tokens[j + 1] == best_pair[1]):
new_tokens.append(new_token)
j += 2
else:
new_tokens.append(tokens[j])
j += 1
tokens = new_tokens
if len(self.vocab) >= vocab_size:
break
self.id_to_token = {i: t for t, i in self.vocab.items()}
def encode(self, text):
"""Encode text to token IDs."""
tokens = list(text)
# Apply merges in order
for pair, merged in self.merges.items():
new_tokens = []
i = 0
while i < len(tokens):
if (i < len(tokens) - 1 and
tokens[i] == pair[0] and
tokens[i + 1] == pair[1]):
new_tokens.append(merged)
i += 2
else:
new_tokens.append(tokens[i])
i += 1
tokens = new_tokens
return [self.vocab.get(t, 0) for t in tokens]
def decode(self, ids):
"""Decode token IDs back to text."""
return ''.join(self.id_to_token.get(i, '?') for i in ids)
@property
def vocab_size(self):
return len(self.vocab)
# Usage
tokenizer = SimpleBPETokenizer()
training_text = "the cat sat on the mat. the dog ran on the mat. the cat and the dog."
tokenizer.train(training_text, vocab_size=64, num_merges=30)
print(f"Vocabulary size: {tokenizer.vocab_size}")
print(f"Merges learned: {len(tokenizer.merges)}")
test = "the cat"
encoded = tokenizer.encode(test)
print(f"Encoded '{test}': {encoded}")
print(f"Decoded: {tokenizer.decode(encoded)}")6.3 Using HuggingFace Tokenizers in Practice
If you want to apply use production-ready tokenizers from HuggingFace like gpt2 or Llama-3. However, by the time you read, there’s probably a more efficieent tokenizer.
from transformers import AutoTokenizer
# Load a pre-trained tokenizer (e.g., GPT-2 style)
tokenizer = AutoTokenizer.from_pretrained("gpt2")
text = "Hello, world! Machine learning is amazing."
# Tokenize
tokens = tokenizer.encode(text)
print(f"Token IDs: {tokens}")
print(f"Number of tokens: {len(tokens)}")
# See the actual tokens
token_strings = tokenizer.convert_ids_to_tokens(tokens)
print(f"Tokens: {token_strings}")
# Decode back
decoded = tokenizer.decode(tokens)
print(f"Decoded: {decoded}")
# Vocabulary size
print(f"Vocabulary size: {tokenizer.vocab_size}")6.4 From Tokens to Arrays in MLX
Here’s where the MLX part is used. Once we have token IDs, we convert them to MLX arrays:
import mlx.core as mx
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
def tokenize(text, max_length=128):
"""Convert text to MLX array of token IDs."""
ids = tokenizer.encode(
text,
max_length=max_length,
truncation=True,
padding="max_length",
)
return mx.array(ids)
# Single example
text = "The quick brown fox jumps over the lazy dog."
input_ids = tokenize(text)
print(f"Shape: {input_ids.shape}") # (128,)
print(f"Dtype: {input_ids.dtype}") # int32
# Batch of examples
texts = [
"The quick brown fox.",
"Machine learning with MLX.",
"Apple silicon is fast.",
]
batch_ids = mx.stack([tokenize(t) for t in texts])
print(f"Batch shape: {batch_ids.shape}") # (3, 128)6.5 Embeddings — Giving Tokens Meaning
An embedding converts sparse token IDs into dense vectors that capture semantic relationships. Similar words get similar vectors.
import mlx.core as mx
import mlx.nn as nn
# An embedding layer maps token IDs to vectors
vocab_size = 10000
embed_dim = 256
embedding = nn.Embedding(vocab_size, embed_dim)
# Token IDs (batch of 4 sequences, each 32 tokens long)
token_ids = mx.random.randint(0, vocab_size, (4, 32))
# Get embeddings
embedded = embedding(token_ids)
print(f"Token IDs shape: {token_ids.shape}") # (4, 32)
print(f"Embedded shape: {embedded.shape}") # (4, 32, 256)
# Each token ID is now a 256-dimensional vector
# Similar tokens will have similar vectors after training6.6 Understanding Embedding Geometry
The magic of embeddings is that they capture relationships in their vector geometry:
import mlx.core as mx
import mlx.nn as mx_nn
# After training, embeddings capture relationships like:
# vec("king") - vec("man") + vec("woman") ≈ vec("queen")
# vec("paris") - vec("france") + vec("italy") ≈ vec("rome")
# Let's simulate this concept:
vocab_size = 10000
embed_dim = 128
embedding = mx_nn.Embedding(vocab_size, embed_dim)
# Compute cosine similarity between two embeddings
def cosine_similarity(a, b):
"""Compute cosine similarity between two vectors."""
a_norm = a / mx.sqrt(mx.sum(a ** 2))
b_norm = b / mx.sqrt(mx.sum(b ** 2))
return mx.sum(a_norm * b_norm)
# Before training, embeddings are random
king_emb = embedding(mx.array(100)) # Token 100 ("king")
queen_emb = embedding(mx.array(101)) # Token 101 ("queen")
apple_emb = embedding(mx.array(200)) # Token 200 ("apple")
# Similarity (before training -- random)
sim_king_queen = cosine_similarity(king_emb, queen_emb)
sim_king_apple = cosine_similarity(king_emb, apple_emb)
mx.eval(sim_king_queen, sim_king_apple)
print(f"Similarity(king, queen): {sim_king_queen.item():.4f}")
print(f"Similarity(king, apple): {sim_king_apple.item():.4f}")
# Before training, these are random -- after training, king/queen should be higher6.7 Positional Encoding
Transformers don’t inherently understand token order. We need to add position information:
import mlx.core as mx
import math
def sinusoidal_position_encoding(max_length, embed_dim):
"""
Create sinusoidal positional encodings as described in
'Attention Is All You Need' (Vaswani et al., 2017).
Args:
max_length: Maximum sequence length
embed_dim: Embedding dimension
Returns:
Array of shape (max_length, embed_dim)
"""
position = mx.arange(max_length).reshape(max_length, 1)
div_term = mx.exp(
mx.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim)
)
pe = mx.zeros((max_length, embed_dim))
pe[:, 0::2] = mx.sin(position * div_term) # Even dimensions
pe[:, 1::2] = mx.cos(position * div_term) # Odd dimensions
return pe
# Create positional encodings
max_len = 512
d_model = 256
pos_enc = sinusoidal_position_encoding(max_len, d_model)
print(f"Position encoding shape: {pos_enc.shape}") # (512, 256)
# Add to embeddings
seq_len = 32
batch_size = 4
token_embeddings = mx.random.normal((batch_size, seq_len, d_model))
position_embeddings = pos_enc[:seq_len] # (32, 256)
# Broadcasting: position_embeddings is added to each sequence in the batch
combined = token_embeddings + position_embeddings
print(f"Combined shape: {combined.shape}") # (4, 32, 256)MLX also provides built-in positional encoding layers:
import mlx.nn as nn
# Sinusoidal positional encoding (built-in)
sin_pe = nn.SinusoidalPositionalEncoding(d_model)
# Rotary Position Embedding (RoPE) -- used in modern models like LLaMA
rope = nn.RoPE(dims=d_model)
# ALiBi (Attention with Linear Biases) -- another positional encoding approach
alibi = nn.ALiBi()6.8 Building a Complete Text Preprocessing Pipeline
Let’s put it all together into a reusable preprocessing pipeline:
"""Text preprocessing pipeline for NLP tasks in MLX."""
import mlx.core as mx
import mlx.nn as nn
from transformers import AutoTokenizer
import math
class TextPreprocessor:
"""Complete text preprocessing for transformer models."""
def __init__(self, model_name="gpt2", max_length=128, embed_dim=256):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.max_length = max_length
self.embed_dim = embed_dim
self.vocab_size = self.tokenizer.vocab_size
# Padding token
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
# Embedding layer
self.token_embedding = nn.Embedding(self.vocab_size, embed_dim)
# Positional encoding
self.pos_encoding = self._create_pos_encoding(max_length, embed_dim)
def _create_pos_encoding(self, max_length, embed_dim):
"""Sinusoidal positional encoding."""
position = mx.arange(max_length).reshape(max_length, 1)
div_term = mx.exp(
mx.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim)
)
pe = mx.zeros((max_length, embed_dim))
pe[:, 0::2] = mx.sin(position * div_term)
pe[:, 1::2] = mx.cos(position * div_term)
return pe
def process(self, texts):
"""
Process a list of texts into embedded tensors.
Args:
texts: List of strings
Returns:
tuple of (input_ids, embeddings, attention_mask)
"""
# Tokenize
encoded = self.tokenizer(
texts,
max_length=self.max_length,
truncation=True,
padding="max_length",
return_tensors=None,
)
# Convert to MLX arrays
input_ids = mx.array(encoded["input_ids"])
attention_mask = mx.array(encoded["attention_mask"])
# Embed tokens
token_emb = self.token_embedding(input_ids)
# Add positional encoding
positions = self.pos_encoding[:self.max_length]
embeddings = token_emb + positions
return input_ids, embeddings, attention_mask
def decode(self, ids):
"""Decode token IDs back to text."""
if isinstance(ids, mx.array):
ids = ids.tolist()
return self.tokenizer.decode(ids, skip_special_tokens=True)7: NLP Tasks and Data Preparation
7.1 Common NLP Tasks
As a machine learning engineer, here are some NLP tasks and patterns you’ll encounter. This section includes task-specific python classes and not self-contained executable scripts. They are considered foundational in machine learning and therefore adapted here to MLX.
| Task | Description | Example |
|---|---|---|
| Text Classification | Assign a category to text | Spam detection, sentiment analysis |
| Token Classification | Assign a label to each token | Named Entity Recognition (NER) |
| Text Generation | Generate new text | Chatbots, code completion |
| Sequence-to-Sequence | Transform input to output | Translation, summarization |
| Masked Language Modeling | Predict masked tokens | BERT pre-training |
| Causal Language Modeling | Predict next token | GPT pre-training |
7.2 The Dataset Object
MLX doesn’t have a built-in dataset class, but we can use the HuggingFace datasets library and wrap it for MLX. The data examples and scripts are for demonstrative purposes. For practice, feel free to use your own data. Free datasets from Kaggle and HuggingFace are good options as well.
"""Dataset utilities for MLX-based NLP training."""
import mlx.core as mx
import numpy as np
from torch.utils.data import Dataset as TorchDataset
class TextDataset:
"""A simple text dataset for language modeling."""
def __init__(self, texts, tokenizer, max_length=128):
self.tokenizer = tokenizer
self.max_length = max_length
self.examples = []
for text in texts:
encoding = tokenizer(
text,
max_length=max_length,
truncation=True,
padding="max_length",
)
self.examples.append({
"input_ids": encoding["input_ids"],
"attention_mask": encoding["attention_mask"],
})
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
example = self.examples[idx]
return {
"input_ids": mx.array(example["input_ids"]),
"attention_mask": mx.array(example["attention_mask"]),
}
def get_batch(self, indices):
"""Get a batch of examples."""
batch = [self[i] for i in indices]
return {
"input_ids": mx.stack([b["input_ids"] for b in batch]),
"attention_mask": mx.stack([b["attention_mask"] for b in batch]),
}
class CausalLMDataset(TextDataset):
"""Dataset for causal (autoregressive) language modeling.
For causal LM, the labels are the input_ids shifted by one position.
The model learns to predict the next token.
"""
def __getitem__(self, idx):
example = self.examples[idx]
input_ids = mx.array(example["input_ids"])
# For causal LM: input is [t1, t2, t3, ...], target is [t2, t3, t4, ...]
return {
"input_ids": input_ids[:-1], # All tokens except last
"targets": input_ids[1:], # All tokens except first
"attention_mask": mx.array(example["attention_mask"][:-1]),
}7.3 Loading Real Datasets
This example loads IMDB’s movie review data from Hugging Face https://huggingface.co/datasets/stanfordnlp/imdb
from datasets import load_dataset
from transformers import AutoTokenizer
# Load a dataset (e.g., IMDB for sentiment analysis)
dataset = load_dataset("imdb")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# Explore the data
print(f"Train size: {len(dataset['train'])}")
print(f"Test size: {len(dataset['test'])}")
print(f"Example: {dataset['train'][0]['text'][:200]}...")
print(f"Label: {dataset['train'][0]['label']}") # 0 = negative, 1 = positive
# Tokenize the dataset
def tokenize_fn(example):
return tokenizer(
example["text"],
truncation=True,
max_length=256,
padding="max_length",
)
tokenized = dataset.map(tokenize_fn, batched=True)7.4 Batching for Training
Efficient batching is crucial for training speed.
import mlx.core as mx
import numpy as np
class DataLoader:
"""Simple DataLoader for MLX."""
def __init__(self, dataset, batch_size=32, shuffle=True):
self.dataset = dataset
self.batch_size = batch_size
self.shuffle = shuffle
self.indices = np.arange(len(dataset))
def __iter__(self):
if self.shuffle:
np.random.shuffle(self.indices)
for start in range(0, len(self.indices), self.batch_size):
batch_indices = self.indices[start:start + self.batch_size]
yield self.dataset.get_batch(batch_indices)
def __len__(self):
return (len(self.dataset) + self.batch_size - 1) // self.batch_size7.5 The Training Loop Pattern
Here’s the standard training loop you’ll use throughout this guide. Think of it as the “request-response cycle” of ML training:
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from tqdm import tqdm
def train(model, dataset, num_epochs=10, batch_size=32, lr=1e-3):
"""Standard training loop for an MLX model."""
# Define loss function
def loss_fn(model, batch):
input_ids = batch["input_ids"]
targets = batch["targets"]
logits = model(input_ids)
# Cross-entropy loss
return nn.losses.cross_entropy(logits, targets, reduction="mean")
# Create gradient function
loss_and_grad = nn.value_and_grad(model, loss_fn)
# Set up optimizer
optimizer = optim.AdamW(learning_rate=lr)
# Set model to training mode
model.train()
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
for epoch in range(num_epochs):
total_loss = 0.0
num_batches = 0
for batch in tqdm(dataloader, desc=f"Epoch {epoch+1}"):
# Forward pass + gradient computation
loss, grads = loss_and_grad(model, batch)
# Update parameters
optimizer.update(model, grads)
# Evaluate (actually run the computation)
mx.eval(model.parameters(), optimizer.state, loss)
total_loss += loss.item()
num_batches += 1
avg_loss = total_loss / num_batches
print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}")
return model7.6 Evaluating Models
Evaluations are crucial in determining model quality. Data is batched processed measuring the total loss value then determining average loss and perplexity.
def evaluate(model, dataset, batch_size=32):
"""Evaluate model on a dataset."""
model.eval() # Set to evaluation mode (disables dropout, etc.)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
total_loss = 0.0
num_batches = 0
for batch in dataloader:
input_ids = batch["input_ids"]
targets = batch["targets"]
logits = model(input_ids)
loss = nn.losses.cross_entropy(logits, targets, reduction="mean")
mx.eval(loss)
total_loss += loss.item()
num_batches += 1
avg_loss = total_loss / num_batches
perplexity = float(mx.exp(mx.array(avg_loss)).item())
return avg_loss, perplexity7.7 Compute and Measure Perplexity of Language Models
Perplexity is the standard metric for language models. It measures how “surprised” the model is by the test data:
- Perplexity of vocab_size = random guessing
- Lower perplexity = better model (less surprised)
- Perplexity of 1 = perfect predictions
import mlx.core as mx
def compute_perplexity(model, dataset, batch_size=32):
"""Compute perplexity of a language model."""
avg_loss, perplexity = evaluate(model, dataset, batch_size)
print(f"Average Loss: {avg_loss:.4f}")
print(f"Perplexity: {perplexity:.2f}")
return perplexity
# Interpretation guide:
# Perplexity > 100: Model is basically random for this data
# Perplexity ~ 30-50: Decent small model
# Perplexity ~ 10-20: Good model
# Perplexity < 10: Strong model
# Perplexity ~ 1: Overfitting (memorized training data)7.8 Data Augmentation for NLP
Including randomness such as swapping words and inserting stop words are augmentation techniques used in NLP.
import random
import re
class TextAugmenter:
"""Simple text augmentation techniques for NLP."""
@staticmethod
def random_delete(text, p=0.1):
"""Randomly delete words with probability p."""
words = text.split()
return ' '.join(w for w in words if random.random() > p)
@staticmethod
def random_swap(text, n=1):
"""Randomly swap two words n times."""
words = text.split()
for _ in range(n):
if len(words) >= 2:
i, j = random.sample(range(len(words)), 2)
words[i], words[j] = words[j], words[i]
return ' '.join(words)
@staticmethod
def random_stopword_insert(text, stopwords=None):
"""Insert a random stopword at a random position."""
if stopwords is None:
stopwords = ["the", "a", "an", "is", "was", "are", "were",
"in", "on", "at", "to", "for", "of", "with"]
words = text.split()
if words:
pos = random.randint(0, len(words))
words.insert(pos, random.choice(stopwords))
return ' '.join(words)7.9 Handling Large Datasets with Streaming
For datasets too large to fit in memory, the preferred processing method is through streaming.
from datasets import load_dataset
# Stream a dataset (doesn't download all at once)
dataset = load_dataset("wikitext", "wikitext-103-raw-v1", streaming=True)
# Process in chunks
def stream_batches(dataset, tokenizer, batch_size=32, max_length=128):
"""Yield batches from a streaming dataset."""
buffer = []
for example in dataset:
text = example["text"]
if not text.strip():
continue
tokens = tokenizer.encode(
text,
max_length=max_length,
truncation=True,
)
if len(tokens) > 1:
buffer.append({
"input_ids": tokens[:-1],
"targets": tokens[1:],
})
if len(buffer) >= batch_size:
batch = buffer[:batch_size]
buffer = buffer[batch_size:]
yield {
"input_ids": mx.stack([mx.array(b["input_ids"]) for b in batch]),
"targets": mx.stack([mx.array(b["targets"]) for b in batch]),
}7.10 Training, Validating, and Testing Splits
Split ratios depend on the total size of the dataset. For small datasets (under ~10k rows), a 60/20/20 split is typical to ensure stable estimates. For medium datasets (~10k–1M rows), 80/10/10 are standard defaults. For large datasets (>1M rows), ratios like 98/1/1 are used because even 1% provides statistically significant validation and test sets.
- Small datasets 60/20/20
- Medium datasets 70/15/15
- Large datasets 80/10/10
from datasets import load_dataset
# Load with splits
dataset = load_dataset("imdb")
train_data = dataset["train"]
test_data = dataset["test"]
# Create validation split from training data
split = train_data.train_test_split(test_size=0.1, seed=42)
train_data = split["train"]
val_data = split["test"]
print(f"Train: {len(train_data)} examples")
print(f"Validation: {len(val_data)} examples")
print(f"Test: {len(test_data)} examples")8: Transformer Architecture in Depth
8.1 The Transformer Revolution
Before Transformers, NLP relied on Recurrent Neural Networks (RNNs) and LSTMs, which processed text one token at a time like reading a book word by word. The Transformer, introduced in the 2017 paper “Attention Is All You Need” by Vaswani, changed everything by processing all tokens in parallel and using “attention” to model relationships between any two positions in a sequence.
An analogy of RNNs is like synchronous sequential code, Transformers are like event-driven asynchronous processing where everything happens at once, with a coordination mechanism of attention to manage dependencies.
8.2 Self-Attention
Self-attention allows each position in a sequence to “attend to” all other positions. The mechanism uses three learned projections:
- Value (V) – What information do I provide?
- Query (Q) What am I looking for?
- Key (K) – What do I contain?
The attention score between two positions is the dot product of their Query and Key vectors, scaled and normalized. Here’s how it looks implemented step-by-step.
import mlx.core as mx
import math
def scaled_dot_product_attention(Q, K, V, mask=None):
"""
Compute scaled dot-product attention.
Args:
Q: Queries, shape (batch, seq_len, d_k)
K: Keys, shape (batch, seq_len, d_k)
V: Values, shape (batch, seq_len, d_v)
mask: Optional mask, shape (batch, seq_len, seq_len)
Returns:
Output: shape (batch, seq_len, d_v)
"""
d_k = Q.shape[-1]
# Compute attention scores: Q @ K^T
# Shape: (batch, seq_len, d_k) @ (batch, d_k, seq_len) -> (batch, seq_len, seq_len)
scores = mx.matmul(Q, mx.transpose(K, (0, 2, 1))) / math.sqrt(d_k)
# Apply mask (for causal/decoder attention)
if mask is not None:
scores = scores + (mask * -1e9)
# Softmax over the key dimension
attention_weights = mx.softmax(scores, axis=-1)
# Weighted sum of values
# Shape: (batch, seq_len, seq_len) @ (batch, seq_len, d_v) -> (batch, seq_len, d_v)
output = mx.matmul(attention_weights, V)
return output, attention_weights
# Example: Attention in action
batch_size = 2
seq_len = 4
d_model = 8
Q = mx.random.normal((batch_size, seq_len, d_model))
K = mx.random.normal((batch_size, seq_len, d_model))
V = mx.random.normal((batch_size, seq_len, d_model))
output, weights = scaled_dot_product_attention(Q, K, V)
mx.eval(output, weights)
print(f"Output shape: {output.shape}") # (2, 4, 8)
print(f"Weights shape: {weights.shape}") # (2, 4, 4)
print(f"Weights sum per position: {mx.sum(weights[0, 0])}") # Should be ~1.0 (softmax)8.3 Understanding Attention Weights
The attention weights form a matrix where entry (i, j) represents how much position i attends to position j.
# Visualize attention pattern (text-based)
def print_attention_pattern(weights, tokens):
"""Print attention weights as a readable pattern."""
seq_len = len(tokens)
print(" " + " ".join(f"{t:>5}" for t in tokens))
for i, token in enumerate(tokens):
row = weights[0, i] # First batch item
mx.eval(row)
scores = [f"{row[j].item():.2f}" for j in range(seq_len)]
print(f"{token:>5}: " + " ".join(f"{s:>5}" for s in scores))
tokens = ["The", "cat", "sat", "down"]
# After a trained model, you'd see patterns like:
# "sat" attending strongly to "cat" (subject-verb relationship)8.4 Multi-Head Attention
Instead of one attention computation, Transformers run multiple “heads” in parallel, each learning different attention patterns.
import mlx.core as mx
import mlx.nn as nn
import math
class MultiHeadAttention(nn.Module):
"""Multi-head self-attention mechanism."""
def __init__(self, d_model, num_heads):
super().__init__()
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads
# Linear projections for Q, K, V, and output
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def __call__(self, x, mask=None):
"""
Forward pass.
Args:
x: Input tensor, shape (batch, seq_len, d_model)
mask: Optional attention mask
Returns:
Output tensor, same shape as input
"""
batch_size, seq_len, _ = x.shape
# Project to Q, K, V
Q = self.W_q(x) # (batch, seq_len, d_model)
K = self.W_k(x)
V = self.W_v(x)
# Reshape for multi-head: (batch, seq_len, d_model) -> (batch, num_heads, seq_len, d_k)
Q = Q.reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
K = K.reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
V = V.reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
# Scaled dot-product attention
d_k = self.d_k
scores = mx.matmul(Q, K.transpose(0, 1, 3, 2)) / math.sqrt(d_k)
if mask is not None:
scores = scores + (mask * -1e9)
attention_weights = mx.softmax(scores, axis=-1)
context = mx.matmul(attention_weights, V)
# Reshape back: (batch, num_heads, seq_len, d_k) -> (batch, seq_len, d_model)
context = context.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model)
# Output projection
output = self.W_o(context)
return output
# Example usage
d_model = 256
num_heads = 8
mha = MultiHeadAttention(d_model, num_heads)
x = mx.random.normal((4, 32, d_model)) # batch=4, seq_len=32, d_model=256
output = mha(x)
mx.eval(output)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}") # Same as input: (4, 32, 256)8.5 Using MLX’s Built-in MultiHeadAttention
MLX provides a built-in MultiHeadAttention layer that you can use directly.
import mlx.core as mx
import mlx.nn as nn
# Built-in multi-head attention
mha = nn.MultiHeadAttention(dims=256, num_heads=8)
x = mx.random.normal((4, 32, 256))
output = mha(x, x, x) # Self-attention: queries, keys, values all from x
mx.eval(output)
print(f"Output shape: {output.shape}") # (4, 32, 256)MLX also provides a fast fused implementation via mx.fast.scaled_dot_product_attention
import mlx.core as mx
# Fast SDPA (fused kernel, more efficient)
# This is used internally by nn.MultiHeadAttention
Q = mx.random.normal((4, 8, 32, 32)) # (batch, heads, seq, d_k)
K = mx.random.normal((4, 8, 32, 32))
V = mx.random.normal((4, 8, 32, 32))
output = mx.fast.scaled_dot_product_attention(Q, K, V, scale=1.0 / 32.0)
mx.eval(output)8.6 Causal Masking For Autoregressive Models
For language generation (GPT-style), we need to prevent positions from attending to future positions.
import mlx.core as mx
def create_causal_mask(seq_len):
"""
Create a causal (lower-triangular) mask.
Position i can only attend to positions 0..i (not i+1..seq_len).
This is essential for autoregressive language models.
"""
# Upper triangular matrix of ones (above diagonal)
mask = mx.triu(mx.ones((seq_len, seq_len)), k=1)
return mask
# Example: 4x4 causal mask
mask = create_causal_mask(4)
mx.eval(mask)
print("Causal mask:")
print(mask)
# [[0, 1, 1, 1],
# [0, 0, 1, 1],
# [0, 0, 0, 1],
# [0, 0, 0, 0]]
# 1 means "blocked" -- position can't see future tokens
# Use in attention:
# scores = scores + (mask * -1e9) # Adds -infinity to masked positions8.7 The Feed-Forward Network
Each Transformer layer has a position-wise feed-forward network (FFN).
import mlx.nn as nn
class FeedForward(nn.Module):
"""Position-wise feed-forward network."""
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.w1 = nn.Linear(d_model, d_ff) # Expand
self.w2 = nn.Linear(d_ff, d_model) # Compress
self.dropout = nn.Dropout(dropout)
self.activation = nn.GELU() # GELU activation (standard in modern models)
def __call__(self, x):
return self.dropout(self.w2(self.activation(self.w1(x))))
# Typical dimensions: d_ff = 4 * d_model
ffn = FeedForward(d_model=256, d_ff=1024)
x = mx.random.normal((4, 32, 256))
output = ffn(x)
mx.eval(output)
print(f"FFN output shape: {output.shape}") # (4, 32, 256)8.8 Layer Normalization
Layer normalization stabilizes training by normalizing across the feature dimension.
import mlx.nn as nn
# Standard LayerNorm
layer_norm = nn.LayerNorm(dims=256)
x = mx.random.normal((4, 32, 256))
normalized = layer_norm(x)
mx.eval(normalized)
print(f"Mean: {mx.mean(normalized).item():.4f}") # Should be ~0
print(f"Std: {mx.std(normalized).item():.4f}") # Should be ~1
# RMSNorm (used in LLaMA and many modern models -- faster than LayerNorm)
rms_norm = nn.RMSNorm(dims=256)
normalized_rms = rms_norm(x)
mx.eval(normalized_rms)8.9 Residual Connections
Residual (skip) connections allow gradients to flow through deep networks:
# Pre-norm Transformer block (used in GPT-2, LLaMA)
# x -> LayerNorm -> Attention -> + x -> LayerNorm -> FFN -> + x
# Post-norm Transformer block (used in original Transformer)
# x -> Attention -> + x -> LayerNorm -> FFN -> + x -> LayerNorm
# In code:
def transformer_block_pre_norm(x, attention, ffn, norm1, norm2):
"""Pre-norm Transformer block with residual connections."""
# Attention sub-layer
residual = x
x = norm1(x) # Normalize
x = attention(x) # Self-attention
x = residual + x # Residual connection
# FFN sub-layer
residual = x
x = norm2(x) # Normalize
x = ffn(x) # Feed-forward
x = residual + x # Residual connection
return x9: Building a Complete Transformer
9.1 Encoder-Only Architecture (BERT-style)
Encoder-only models process the entire input simultaneously and are used for understanding tasks.
import mlx.core as mx
import mlx.nn as nn
import math
class TransformerEncoderBlock(nn.Module):
"""A single Transformer encoder block."""
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.attention = nn.MultiHeadAttention(d_model, num_heads)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(d_ff, d_model),
nn.Dropout(dropout),
)
self.dropout = nn.Dropout(dropout)
def __call__(self, x, mask=None):
# Self-attention with residual
residual = x
x = self.norm1(x)
x = self.attention(x, mask)
x = self.dropout(x)
x = residual + x
# FFN with residual
residual = x
x = self.norm2(x)
x = self.ffn(x)
x = residual + x
return x
class TransformerEncoder(nn.Module):
"""Full Transformer encoder stack."""
def __init__(self, vocab_size, d_model, num_heads, d_ff,
num_layers, max_seq_len, dropout=0.1):
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.pos_embedding = nn.Embedding(max_seq_len, d_model)
self.layers = [
TransformerEncoderBlock(d_model, num_heads, d_ff, dropout)
for _ in range(num_layers)
]
self.norm = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def __call__(self, input_ids, mask=None):
seq_len = input_ids.shape[1]
positions = mx.arange(seq_len)
# Embed tokens + positions
x = self.token_embedding(input_ids) + self.pos_embedding(positions)
x = self.dropout(x)
# Pass through encoder layers
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)
# Build an encoder for classification
class BertClassifier(nn.Module):
"""BERT-style classifier for text classification."""
def __init__(self, vocab_size, num_classes, d_model=256,
num_heads=8, num_layers=4, max_seq_len=128):
super().__init__()
self.encoder = TransformerEncoder(
vocab_size=vocab_size,
d_model=d_model,
num_heads=num_heads,
d_ff=d_model * 4,
num_layers=num_layers,
max_seq_len=max_seq_len,
)
self.classifier = nn.Sequential(
nn.Linear(d_model, d_model),
nn.GELU(),
nn.Linear(d_model, num_classes),
)
def __call__(self, input_ids):
# Encode
hidden = self.encoder(input_ids) # (batch, seq_len, d_model)
# Pool: use mean of all token representations
pooled = mx.mean(hidden, axis=1) # (batch, d_model)
# Classify
logits = self.classifier(pooled) # (batch, num_classes)
return logits9.2 Decoder-Only Architecture (GPT-style)
Decoder-only models generate text auto-regressively predicting one token at a time:
import mlx.core as mx
import mlx.nn as nn
import math
class TransformerDecoderBlock(nn.Module):
"""A single Transformer decoder block with causal attention."""
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.attention = nn.MultiHeadAttention(d_model, num_heads)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(d_ff, d_model),
nn.Dropout(dropout),
)
self.dropout = nn.Dropout(dropout)
def __call__(self, x, mask=None):
# Causal self-attention with residual
residual = x
x = self.norm1(x)
x = self.attention(x, mask)
x = self.dropout(x)
x = residual + x
# FFN with residual
residual = x
x = self.norm2(x)
x = self.ffn(x)
x = residual + x
return x
class GPTModel(nn.Module):
"""GPT-style causal language model."""
def __init__(self, vocab_size, d_model=256, num_heads=8,
num_layers=4, d_ff=1024, max_seq_len=256, dropout=0.1):
super().__init__()
self.d_model = d_model
self.max_seq_len = max_seq_len
# Token and position embeddings
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.pos_embedding = nn.Embedding(max_seq_len, d_model)
# Transformer blocks
self.layers = [
TransformerDecoderBlock(d_model, num_heads, d_ff, dropout)
for _ in range(num_layers)
]
self.dropout = nn.Dropout(dropout)
self.norm = nn.LayerNorm(d_model)
# Output head (tied with token embedding weights)
self.output_proj = nn.Linear(d_model, vocab_size)
def __call__(self, input_ids):
"""
Forward pass.
Args:
input_ids: Token IDs, shape (batch, seq_len)
Returns:
Logits, shape (batch, seq_len, vocab_size)
"""
batch_size, seq_len = input_ids.shape
# Create positions
positions = mx.arange(seq_len)
# Create causal mask
causal_mask = mx.triu(
mx.full((seq_len, seq_len), -1e9), k=1
)
# Embed tokens and positions
x = self.token_embedding(input_ids) + self.pos_embedding(positions)
x = self.dropout(x)
# Pass through transformer blocks
for layer in self.layers:
x = layer(x, mask=causal_mask)
# Final layer norm
x = self.norm(x)
# Project to vocabulary
logits = self.output_proj(x)
return logits
def generate(self, input_ids, max_new_tokens=50, temperature=1.0, top_k=None):
"""
Generate text autoregressively.
Args:
input_ids: Starting token IDs, shape (batch, seq_len)
max_new_tokens: Number of tokens to generate
temperature: Sampling temperature (lower = more deterministic)
top_k: If set, only sample from top k tokens
Returns:
Generated token IDs, shape (batch, seq_len + max_new_tokens)
"""
for _ in range(max_new_tokens):
# Truncate to max sequence length
context = input_ids[:, -self.max_seq_len:]
# Get logits for the last position
logits = self(context)
next_logits = logits[:, -1, :] # (batch, vocab_size)
# Apply temperature
if temperature > 0:
next_logits = next_logits / temperature
# Apply top-k filtering
if top_k is not None:
top_values, top_indices = mx.topk(next_logits, top_k)
probs = mx.softmax(top_values, axis=-1)
# Sample from top-k
idx = mx.random.categorical(mx.log(probs + 1e-10)[:, None, :])[:, None]
next_token = mx.take_along_axis(top_indices, idx, axis=1)[:, 0]
else:
# Sample from full distribution
probs = mx.softmax(next_logits, axis=-1)
next_token = mx.random.categorical(
mx.log(probs + 1e-10)[:, None, :]
)[:, 0]
# Append to sequence
input_ids = mx.concatenate(
[input_ids, next_token[:, None]], axis=1
)
return input_ids9.3 Using MLX’s Built-in Transformer
MLX provides a built-in nn.Transformer for convenience.
import mlx.nn as nn
# Built-in Transformer
transformer = nn.Transformer(
dims=256, # Model dimension
num_heads=8, # Number of attention heads
num_encoder_layers=4,
num_decoder_layers=4,
d_ff=1024, # Feed-forward dimension
dropout=0.1,
)9.4 RMSNorm and Modern Architecture Choices
Modern models like LLaMA use RMSNorm instead of LayerNorm, and SwiGLU instead of GELU in the FFN.
import mlx.core as mx
import mlx.nn as nn
class LLaMABlock(nn.Module):
"""LLaMA-style transformer block with RMSNorm and SwiGLU."""
def __init__(self, d_model, num_heads, d_ff, rope=None):
super().__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads
# RMSNorm instead of LayerNorm
self.norm1 = nn.RMSNorm(d_model)
self.norm2 = nn.RMSNorm(d_model)
# Attention projections (fused QKV)
self.w_q = nn.Linear(d_model, d_model, bias=False)
self.w_k = nn.Linear(d_model, d_model, bias=False)
self.w_v = nn.Linear(d_model, d_model, bias=False)
self.w_o = nn.Linear(d_model, d_model, bias=False)
# SwiGLU FFN
self.w_gate = nn.Linear(d_model, d_ff, bias=False) # Gate
self.w_up = nn.Linear(d_model, d_ff, bias=False) # Up projection
self.w_down = nn.Linear(d_ff, d_model, bias=False) # Down projection
# Rotary position embedding
self.rope = rope or nn.RoPE(dims=self.d_k)
def __call__(self, x, mask=None):
batch_size, seq_len, _ = x.shape
# Pre-norm attention
residual = x
x = self.norm1(x)
# Compute Q, K, V
Q = self.w_q(x).reshape(batch_size, seq_len, self.num_heads, self.d_k)
K = self.w_k(x).reshape(batch_size, seq_len, self.num_heads, self.d_k)
V = self.w_v(x).reshape(batch_size, seq_len, self.num_heads, self.d_k)
# Transpose to (batch, heads, seq, d_k)
Q = Q.transpose(0, 2, 1, 3)
K = K.transpose(0, 2, 1, 3)
V = V.transpose(0, 2, 1, 3)
# Apply RoPE to Q and K
Q = self.rope(Q)
K = self.rope(K)
# Scaled dot-product attention
scores = mx.matmul(Q, K.transpose(0, 1, 3, 2)) / mx.sqrt(
mx.array(self.d_k)
)
if mask is not None:
scores = scores + mask
attention_weights = mx.softmax(scores, axis=-1)
context = mx.matmul(attention_weights, V)
# Reshape and project output
context = context.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, -1)
x = self.w_o(context)
x = x + residual
# Pre-norm FFN with SwiGLU
residual = x
x = self.norm2(x)
gate = nn.silu(self.w_gate(x)) # SiLU (swish) activation on gate
up = self.w_up(x) # Up projection
x = self.w_down(gate * up) # Gated + down projection
x = x + residual
return x9.5 Weight Tying
In GPT-style models, the output projection weights are often shared with the token embedding.
class GPTWithWeightTying(nn.Module):
"""GPT model with weight tying between embedding and output projection."""
def __init__(self, vocab_size, d_model, num_heads, num_layers, max_seq_len):
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, d_model)
# ... other layers ...
def __call__(self, input_ids):
# ... transformer processing ...
x = self.norm(x)
# Use transposed embedding weights for output projection
# This saves memory and can improve training
logits = x @ self.token_embedding.weight.T
return logits9.6 Model Parameter Counting
Understanding how many parameters your model has is important.
from mlx.utils import tree_flatten
def count_parameters(model):
"""Count total and trainable parameters in a model."""
total = sum(v.size for _, v in tree_flatten(model.parameters()))
trainable = sum(v.size for _, v in tree_flatten(model.trainable_parameters()))
return total, trainable
def format_params(n):
"""Format parameter count in human-readable form."""
if n >= 1e9:
return f"{n/1e9:.1f}B"
elif n >= 1e6:
return f"{n/1e6:.1f}M"
elif n >= 1e3:
return f"{n/1e3:.1f}K"
return str(n)
# Example
model = GPTModel(vocab_size=50000, d_model=256, num_heads=8, num_layers=4)
total, trainable = count_parameters(model)
print(f"Total parameters: {format_params(total)}")
print(f"Trainable parameters: {format_params(trainable)}")9.7 Memory Estimation
def estimate_memory(model, batch_size=1, seq_len=128, dtype_bytes=4):
"""Estimate memory usage of a model."""
# Parameter memory
from mlx.utils import tree_flatten
total_params = sum(v.size for _, v in tree_flatten(model.parameters()))
param_memory = total_params * dtype_bytes
# Activation memory (rough estimate)
# Each transformer layer stores Q, K, V, attention weights, FFN intermediates
# Rough: ~5 * batch_size * seq_len * d_model * num_layers
d_model = model.d_model
num_layers = len(model.layers)
activation_memory = (
5 * batch_size * seq_len * d_model * num_layers * dtype_bytes
)
# Gradient memory (same as parameters for full fine-tuning)
gradient_memory = param_memory
# Optimizer state (Adam: 2x parameters for momentum and variance)
optimizer_memory = 2 * param_memory
total_memory = param_memory + activation_memory + gradient_memory + optimizer_memory
print(f"Parameters: {param_memory / 1e9:.2f} GB")
print(f"Activations: {activation_memory / 1e9:.2f} GB")
print(f"Gradients: {gradient_memory / 1e9:.2f} GB")
print(f"Optimizer: {optimizer_memory / 1e9:.2f} GB")
print(f"Total: {total_memory / 1e9:.2f} GB")
return total_memory9.8 Model Configuration Patterns
Pre-defining configurations for different model sizes and purposes.
# Small model (for learning/testing)
small_config = {
"vocab_size": 50000,
"d_model": 256,
"num_heads": 8,
"num_layers": 4,
"d_ff": 1024,
"max_seq_len": 256,
"dropout": 0.1,
}
# Medium model (for real tasks)
medium_config = {
"vocab_size": 50000,
"d_model": 512,
"num_heads": 8,
"num_layers": 8,
"d_ff": 2048,
"max_seq_len": 512,
"dropout": 0.1,
}
# Large model (needs significant memory)
large_config = {
"vocab_size": 50000,
"d_model": 1024,
"num_heads": 16,
"num_layers": 16,
"d_ff": 4096,
"max_seq_len": 1024,
"dropout": 0.1,
}10: Training a Language Model
10.1 Project Overview
We’ll build a character-level language model from scratch. This small model will teach you the complete training pipeline that scales to much larger models. The objective is to train a model to generate text character by character, learning from a training corpus.
10.2 The Complete Training Script
This section’s code lives in projects/text_generator/train.py. Here we’ll walk through each component:
"""
Character-level language model training with MLX.
This script trains a small transformer to generate text one character
at a time, learning patterns from a training corpus.
"""
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np
import math
from dataclasses import dataclass
@dataclass
class Config:
"""Model configuration."""
vocab_size: int = 65 # Unique characters in dataset
d_model: int = 128 # Embedding dimension
num_heads: int = 4 # Attention heads
num_layers: int = 4 # Transformer layers
d_ff: int = 512 # Feed-forward dimension (4x d_model)
max_seq_len: int = 256 # Maximum sequence length
dropout: float = 0.1 # Dropout rate
learning_rate: float = 3e-4 # Learning rate
batch_size: int = 64 # Batch size
num_epochs: int = 20 # Number of epochs
eval_interval: int = 100 # Evaluate every N steps10.3 The Model Architecture
class CharTransformer(nn.Module):
"""A small character-level transformer language model."""
def __init__(self, config: Config):
super().__init__()
self.config = config
# Embeddings
self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
self.pos_embedding = nn.Embedding(config.max_seq_len, config.d_model)
# Transformer layers
self.layers = [
nn.TransformerEncoderLayer(
dims=config.d_model,
num_heads=config.num_heads,
mlp_dims=config.d_ff,
dropout=config.dropout,
)
for _ in range(config.num_layers)
]
self.norm = nn.LayerNorm(config.d_model)
self.output_proj = nn.Linear(config.d_model, config.vocab_size)
self.dropout = nn.Dropout(config.dropout)
def __call__(self, x):
B, T = x.shape
# Token + position embeddings
tok_emb = self.token_embedding(x)
pos_emb = self.pos_embedding(mx.arange(T))
x = self.dropout(tok_emb + pos_emb)
# Causal mask
mask = mx.triu(mx.full((T, T), -1e9), k=1)
# Transformer layers
for layer in self.layers:
x = layer(x, mask)
x = self.norm(x)
logits = self.output_proj(x)
return logits10.4 Data Preparation
class CharDataset:
"""Character-level dataset for language modeling."""
def __init__(self, text, seq_len=256):
self.text = text
self.seq_len = seq_len
# Build vocabulary
self.chars = sorted(set(text))
self.vocab_size = len(self.chars)
self.char_to_id = {ch: i for i, ch in enumerate(self.chars)}
self.id_to_char = {i: ch for ch, i in self.char_to_id.items()}
# Encode entire text
self.encoded = np.array([self.char_to_id[ch] for ch in text], dtype=np.int32)
def __len__(self):
return max(0, (len(self.encoded) - self.seq_len) // self.seq_len)
def get_batch(self, batch_size):
"""Get a random batch of (input, target) pairs."""
# Random starting positions
indices = np.random.randint(
0, len(self.encoded) - self.seq_len - 1, size=batch_size
)
x = mx.array(np.stack([
self.encoded[i:i + self.seq_len] for i in indices
]))
y = mx.array(np.stack([
self.encoded[i + 1:i + self.seq_len + 1] for i in indices
]))
return x, y
def encode(self, text):
"""Encode text to token IDs."""
return [self.char_to_id[ch] for ch in text]
def decode(self, ids):
"""Decode token IDs to text."""
return ''.join(self.id_to_char[i] for i in ids)10.5 The Training Loop
def train_char_model():
"""Train a character-level language model."""
# Load data (Shakespeare is a common small dataset for this)
# You can replace this with any text file
text = open("shakespeare.txt", "r").read()
print(f"Dataset size: {len(text)} characters")
# Create dataset
config = Config()
dataset = CharDataset(text, seq_len=config.max_seq_len)
config.vocab_size = dataset.vocab_size
print(f"Vocabulary: {dataset.vocab_size} characters")
print(f"Vocabulary: {''.join(dataset.chars)}")
# Create model
model = CharTransformer(config)
# Count parameters
from mlx.utils import tree_flatten
total_params = sum(v.size for _, v in tree_flatten(model.parameters()))
print(f"Total parameters: {total_params:,}")
# Loss function
def loss_fn(model, x, y):
logits = model(x)
return nn.losses.cross_entropy(
logits.reshape(-1, config.vocab_size),
y.reshape(-1),
reduction="mean"
)
# Gradient function
loss_and_grad = nn.value_and_grad(model, loss_fn)
# Optimizer with cosine decay
lr_schedule = optim.cosine_decay(config.learning_rate, 1000, 1e-5)
optimizer = optim.AdamW(learning_rate=lr_schedule)
# Training loop
model.train()
steps_per_epoch = len(dataset) // config.batch_size
for epoch in range(config.num_epochs):
total_loss = 0.0
num_steps = 0
for step in range(steps_per_epoch):
x, y = dataset.get_batch(config.batch_size)
# Forward + backward
loss, grads = loss_and_grad(model, x, y)
# Update
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state, loss)
total_loss += loss.item()
num_steps += 1
if step % config.eval_interval == 0:
avg_loss = total_loss / max(num_steps, 1)
print(f"Epoch {epoch+1}, Step {step}: loss = {avg_loss:.4f}")
avg_loss = total_loss / num_steps
perplexity = math.exp(avg_loss)
print(f"Epoch {epoch+1} complete: loss = {avg_loss:.4f}, "
f"perplexity = {perplexity:.2f}")
# Generate sample
sample = generate(model, dataset, "The ", max_tokens=200)
print(f"Sample: {sample}\n")
return model, dataset
def generate(model, dataset, prompt, max_tokens=200, temperature=0.8):
"""Generate text from the model."""
model.eval()
tokens = dataset.encode(prompt)
input_ids = mx.array([tokens])
for _ in range(max_tokens):
# Get logits for the last position
logits = model(input_ids)
next_logits = logits[:, -1, :] / temperature
# Sample
probs = mx.softmax(next_logits, axis=-1)
next_token = mx.random.categorical(mx.log(probs + 1e-10)[:, None, :])[:, 0]
input_ids = mx.concatenate([input_ids, next_token[:, None]], axis=1)
mx.eval(input_ids)
return dataset.decode(input_ids[0].tolist())10.6 Understanding the Training Dynamics
The model captures sentence structure and style. Observe the following when training a character-level model:
Early training (loss ~4.0, perplexity ~55):
The the the the the the the the
The model learns common characters first.
Mid training (loss ~2.0, perplexity ~7):
The caled the shall the pake the soneThe model learns word-like patterns.
Late training (loss ~1.5, perplexity ~4.5):The king shall hear the common voice
And speak the truth of noble deed
11: Training a BERT-Style Classifier
11.1 Transfer Learning for Classification
Instead of training from scratch, we can fine-tune a pre-trained model for specific tasks. This is much more efficient and effective.
11.2 Building a Sentiment Classifier
The complete code for this project is in projects/sentiment_classifier/. Here’s the walkthrough.
"""
Sentiment classification using a Transformer encoder trained with MLX.
This trains a BERT-style model on the IMDB dataset for binary
sentiment classification (positive/negative).
"""
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from transformers import AutoTokenizer
from datasets import load_dataset
from tqdm import tqdm
import numpy as np
# Configuration
VOCAB_SIZE = 30522 # GPT-2 tokenizer vocab size
D_MODEL = 256
NUM_HEADS = 8
NUM_LAYERS = 4
D_FF = 1024
MAX_SEQ_LEN = 256
DROPOUT = 0.1
BATCH_SIZE = 32
LEARNING_RATE = 5e-5
NUM_EPOCHS = 311.3 The Classification Model
class SentimentClassifier(nn.Module):
"""Transformer-based text classifier."""
def __init__(self, vocab_size, num_classes=2):
super().__init__()
self.embedding = nn.Embedding(vocab_size, D_MODEL)
self.pos_embedding = nn.Embedding(MAX_SEQ_LEN, D_MODEL)
self.embed_dropout = nn.Dropout(DROPOUT)
self.layers = [
nn.TransformerEncoderLayer(
dims=D_MODEL,
num_heads=NUM_HEADS,
mlp_dims=D_FF,
dropout=DROPOUT,
)
for _ in range(NUM_LAYERS)
]
self.norm = nn.LayerNorm(D_MODEL)
self.classifier = nn.Sequential(
nn.Linear(D_MODEL, D_MODEL),
nn.GELU(),
nn.Dropout(DROPOUT),
nn.Linear(D_MODEL, num_classes),
)
def __call__(self, input_ids, mask=None):
B, T = input_ids.shape
# Embeddings
x = self.embedding(input_ids) + self.pos_embedding(mx.arange(T))
x = self.embed_dropout(x)
# Encode
for layer in self.layers:
x = layer(x, mask)
x = self.norm(x)
# Pool: mean of non-padding tokens
if mask is not None:
mask_expanded = mask[:, :, None].astype(mx.float32)
x = mx.sum(x * mask_expanded, axis=1) / mx.sum(mask_expanded, axis=1)
else:
x = mx.mean(x, axis=1)
# Classify
return self.classifier(x)11.4 Data Preparation
def prepare_data():
"""Load and prepare the IMDB dataset."""
tokenizer = AutoTokenizer.from_pretrained("gpt2")
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
dataset = load_dataset("imdb")
def tokenize(examples):
return tokenizer(
examples["text"],
padding="max_length",
truncation=True,
max_length=MAX_SEQ_LEN,
)
tokenized = dataset.map(tokenize, batched=True)
train_data = []
for example in tokenized["train"]:
train_data.append({
"input_ids": np.array(example["input_ids"], dtype=np.int32),
"attention_mask": np.array(example["attention_mask"], dtype=np.int32),
"label": np.array(example["label"], dtype=np.int32),
})
test_data = []
for example in tokenized["test"]:
test_data.append({
"input_ids": np.array(example["input_ids"], dtype=np.int32),
"attention_mask": np.array(example["attention_mask"], dtype=np.int32),
"label": np.array(example["label"], dtype=np.int32),
})
return train_data, test_data, tokenizer11.5 Training with Accuracy Tracking
def train_classifier():
"""Train the sentiment classifier."""
train_data, test_data, tokenizer = prepare_data()
model = SentimentClassifier(vocab_size=tokenizer.vocab_size, num_classes=2)
def loss_fn(model, batch):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["label"]
logits = model(input_ids, mask=None)
return nn.losses.cross_entropy(logits, labels, reduction="mean")
loss_and_grad = nn.value_and_grad(model, loss_fn)
optimizer = optim.AdamW(learning_rate=LEARNING_RATE, weight_decay=0.01)
model.train()
for epoch in range(NUM_EPOCHS):
np.random.shuffle(train_data)
total_loss = 0.0
correct = 0
total = 0
for i in tqdm(range(0, len(train_data), BATCH_SIZE),
desc=f"Epoch {epoch+1}"):
batch_data = train_data[i:i + BATCH_SIZE]
batch = {
"input_ids": mx.stack([mx.array(d["input_ids"]) for d in batch_data]),
"attention_mask": mx.stack([mx.array(d["attention_mask"]) for d in batch_data]),
"label": mx.array([d["label"] for d in batch_data]),
}
loss, grads = loss_and_grad(model, batch)
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state, loss)
total_loss += loss.item()
# Compute accuracy
logits = model(batch["input_ids"])
predictions = mx.argmax(logits, axis=-1)
mx.eval(predictions)
correct += mx.sum(predictions == batch["label"]).item()
total += len(batch_data)
avg_loss = total_loss / (len(train_data) / BATCH_SIZE)
accuracy = correct / total
print(f"Epoch {epoch+1}: loss = {avg_loss:.4f}, "
f"accuracy = {accuracy:.4f}")
return model11.6 Understanding Loss Functions for Classification
import mlx.core as mx
import mlx.nn as nn
# Cross-entropy loss: the standard for classification
# Measures how far the predicted probability distribution is from the true one
# For a single example:
logits = mx.array([[2.0, -1.0]]) # Model output (2 classes)
target = mx.array([0]) # True class is 0
loss = nn.losses.cross_entropy(logits, target)
mx.eval(loss)
print(f"Cross-entropy loss: {loss.item():.4f}")
# The loss is lower when the model assigns higher probability to the correct class
good_logits = mx.array([[10.0, -10.0]]) # Very confident correct prediction
bad_logits = mx.array([[-10.0, 10.0]]) # Very confident wrong prediction
print(f"Good prediction loss: {nn.losses.cross_entropy(good_logits, target).item():.4f}")
print(f"Bad prediction loss: {nn.losses.cross_entropy(bad_logits, target).item():.4f}")12: Fine-Tuning Large Language Models with LoRA
12.1 Why Fine-Tuning?
Training a large language model from scratch requires enormous compute and data. Fine-tuning adapts a pre-trained model to a specific task using much less data and compute. It’s like how a web developer who knows JavaScript can quickly learn TypeScript — the foundation is already there.
12.2 LoRA Low-Rank Adaptation
LoRA is a parameter-efficient fine-tuning technique. Instead of updating all model weights, LoRA adds small “adapter” matrices that are trained while the original weights stay frozen:
Original: y = W @ x (W is large, e.g., 4096 x 4096)
LoRA: y = W @ x + B @ A @ x (A is 4096 x r, B is r x 4096, r << 4096)Where r (rank) is typically 8-64. This reduces trainable parameters by 90%+ while maintaining quality.
import mlx.core as mx
import mlx.nn as nn
class LoRALinear(nn.Module):
"""Linear layer with LoRA adaptation."""
def __init__(self, base_linear: nn.Linear, rank: int = 8, alpha: float = 16.0):
super().__init__()
self.base_linear = base_linear
self.rank = rank
self.alpha = alpha
self.scale = alpha / rank
# Original dimensions
in_features = base_linear.weight.shape[1]
out_features = base_linear.weight.shape[0]
# LoRA matrices
# A: (in_features, rank) -- initialized with Kaiming
# B: (rank, out_features) -- initialized with zeros
self.lora_A = mx.random.normal(shape=(in_features, rank)) * 0.01
self.lora_B = mx.zeros((rank, out_features))
def __call__(self, x):
# Original forward pass (frozen weights)
base_output = self.base_linear(x)
# LoRA adaptation
lora_output = (x @ self.lora_A @ self.lora_B) * self.scale
return base_output + lora_output12.3 Applying LoRA to a Model
def apply_lora(model, rank=8, alpha=16.0, target_modules=None):
"""
Apply LoRA to specific layers of a model.
Args:
model: The nn.Module to modify
rank: LoRA rank
alpha: LoRA alpha (scaling factor)
target_modules: List of module name patterns to apply LoRA to
"""
if target_modules is None:
# Default: apply to attention projection layers
target_modules = ["w_q", "w_k", "w_v", "w_o"]
def _apply_to_module(name, module):
for attr_name in dir(module):
child = getattr(module, attr_name)
if isinstance(child, nn.Linear):
# Check if this module name matches any target
should_apply = any(t in attr_name for t in target_modules)
if should_apply:
lora_layer = LoRALinear(child, rank=rank, alpha=alpha)
setattr(module, attr_name, lora_layer)
# Apply LoRA and freeze base weights
model.apply_to_modules(lambda n, m: _apply_to_module(n, m))
# Freeze all parameters except LoRA
model.freeze()
# Unfreeze LoRA parameters
def _unfreeze_lora(module):
if isinstance(module, LoRALinear):
module.lora_A.requires_grad = True
module.lora_B.requires_grad = True
model.apply_to_modules(lambda n, m: _unfreeze_lora(m))
return model12.4 Using MLX-LM for Fine-Tuning
The mlx-lm package provides a ready-to-use fine-tuning pipeline.
# Fine-tune with LoRA using mlx-lm
# Prepare your data in JSONL format:
# {"text": "prompt and response text"}
# {"text": "another example"}
# Run fine-tuning
mlx_lm.lora \
--model mlx-community/Qwen3-4B-Instruct-2507-4bit \
--data ./my_data \
--batch-size 4 \
--lora-layers 16 \
--iters 1000 \
--learning-rate 1e-512.5 Fine-Tuning Data Format
For supervised fine-tuning (SFT), organize your data as JSONL:
{ "messages": [{"role": "system", "content": "You are a helpful coding assistant."}, {"role": "user", "content": "Write a Python function to reverse a string."}, {"role": "assistant", "content": "def reverse_string(s):\n return s[::-1]"}]}
{"messages": [{"role": "system", "content": "You are a helpful coding assistant."}, {"role": "user", "content": "Explain REST APIs."}, {"role": "assistant", "content": "REST (Representational State Transfer) is an architectural style for web services..."}]}12.6 A Complete LoRA Fine-Tuning Script
The full project code is in projects/lora_finetune/. Here’s the core training loop.
"""
LoRA fine-tuning script for MLX.
"""
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from transformers import AutoTokenizer, AutoModelForCausalLM
import json
import numpy as np
from tqdm import tqdm
# Configuration
MODEL_NAME = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
LORA_RANK = 8
LORA_ALPHA = 16.0
LEARNING_RATE = 1e-5
BATCH_SIZE = 4
NUM_ITERS = 1000
SEQ_LEN = 256
def load_data(path):
"""Load JSONL training data."""
data = []
with open(path, "r") as f:
for line in f:
if line.strip():
example = json.loads(line)
text = example.get("text", "")
if text:
data.append(text)
return data
def create_batches(data, tokenizer, batch_size, seq_len):
"""Tokenize and batch data."""
all_tokens = []
for text in data:
tokens = tokenizer.encode(text)[:seq_len]
if len(tokens) > 1:
all_tokens.append(tokens)
# Pad to uniform length
max_len = min(max(len(t) for t in all_tokens), seq_len)
padded = []
for tokens in all_tokens:
if len(tokens) < max_len:
tokens = tokens + [0] * (max_len - len(tokens))
padded.append(tokens[:max_len])
# Create batches
batches = []
for i in range(0, len(padded), batch_size):
batch = padded[i:i + batch_size]
if len(batch) == batch_size:
input_ids = mx.array(batch)[:, :-1]
targets = mx.array(batch)[:, 1:]
batches.append((input_ids, targets))
return batches
def train_lora():
"""Fine-tune a model with LoRA."""
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# Load data
train_data = load_data("train.jsonl")
val_data = load_data("val.jsonl")
train_batches = create_batches(train_data, tokenizer, BATCH_SIZE, SEQ_LEN)
val_batches = create_batches(val_data, tokenizer, BATCH_SIZE, SEQ_LEN)
print(f"Training batches: {len(train_batches)}")
print(f"Validation batches: {len(val_batches)}")
# Note: In practice, use mlx-lm's built-in LoRA support
# This shows the conceptual training loop
print("Use mlx_lm.lora for production fine-tuning")
print(f" mlx_lm.lora --model {MODEL_NAME} --data ./data --iters {NUM_ITERS}")12.7 Evaluating Fine-Tuned Models
def evaluate_generation(model, tokenizer, prompts, max_tokens=100):
"""Evaluate a fine-tuned model on generation quality."""
model.eval()
for prompt in prompts:
# Tokenize
input_ids = mx.array([tokenizer.encode(prompt)])
# Generate
output_ids = model.generate(input_ids, max_new_tokens=max_tokens)
# Decode
text = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True)
print(f"Prompt: {prompt}")
print(f"Output: {text}")
print("-" * 80)12.8 Quantization for Efficient Inference
After fine-tuning, you can quantize the model for even more efficient inference:
import mlx.core as mx
import mlx.nn as nn
# Quantize a model to 4-bit (reduces memory by ~4x)
def quantize_model(model, bits=4, group_size=64):
"""Quantize model weights for efficient inference."""
nn.quantize(model, bits=bits, group_size=group_size)
return model
# Or use mlx-lm from the command line:
# mlx_lm.convert --model ./my_finetuned_model -q --quantize --bits 412.9 Saving and Loading Fine-Tuned Models
import mlx.core as mx
# Save LoRA weights (only the adapter weights, not the full model)
def save_lora_weights(model, path):
"""Save only the LoRA adapter weights."""
lora_params = {}
for name, param in model.parameters().items():
if "lora" in name:
lora_params[name] = param
mx.save_safetensors(path, lora_params)
# Load weights
def load_weights(model, path):
"""Load weights from a safetensors file."""
weights = mx.load(path)
model.load_weights(list(weights.items()))
return model12.10 Training Tips and Best Practices
import mlx.optimizers as optim
# Warmup + cosine decay (standard for fine-tuning)
warmup_steps = 100
total_steps = 1000
# Linear warmup
warmup_schedule = optim.linear_schedule(0, LEARNING_RATE, warmup_steps)
# Cosine decay after warmup
decay_schedule = optim.cosine_decay(LEARNING_RATE, total_steps - warmup_steps, 1e-6)
# Combine
lr_schedule = optim.join_schedules([warmup_schedule, decay_schedule], [warmup_steps])
optimizer = optim.AdamW(learning_rate=lr_schedule, weight_decay=0.01)Gradient Clipping (prevent exploding gradients):
# Clip gradient norm to prevent instability
grads = optim.clip_grad_norm(grads, max_norm=1.0)Mixed Precision Training:
# Use float16 for forward pass, float32 for optimizer
model.set_dtype(mx.float16)
# MLX handles the mixed precision automatically in most casesCheckpointing (save training progress):
def save_checkpoint(model, optimizer, step, path):
"""Save training checkpoint."""
checkpoint = {
"step": step,
"model_state": model.parameters(),
"optimizer_state": optimizer.state,
}
mx.savez(path, checkpoint)
def load_checkpoint(model, optimizer, path):
"""Load training checkpoint."""
checkpoint = mx.load(path)
model.update(checkpoint["model_state"])
optimizer.state = checkpoint["optimizer_state"]
return checkpoint["step"]13: Model Optimization and Performance
13.1 Memory Management in MLX
Understanding memory is crucial for training large models. MLX provides tools to monitor and control memory usage.
import mlx.core as mx
# Monitor memory usage
print(f"Active memory: {mx.get_active_memory() / 1e9:.2f} GB")
print(f"Peak memory: {mx.get_peak_memory() / 1e9:.2f} GB")
print(f"Cache memory: {mx.get_cache_memory() / 1e9:.2f} GB")
# Reset peak memory tracking
mx.reset_peak_memory()
# Set memory limits
mx.set_memory_limit(8 * 1024 * 1024 * 1024) # 8 GB limit
mx.set_cache_limit(2 * 1024 * 1024 * 1024) # 2 GB cache limit
# Clear the cache manually
mx.clear_cache()13.2 Metal GPU Information
import mlx.core as mx
if mx.metal.is_available():
info = mx.metal.device_info()
print(f"GPU available: True")
print(f"Device name: {info.get('device_name', 'Unknown')}")
print(f"Memory: {info.get('memory_size', 0) / 1e9:.1f} GB")
print(f"Compute cores: {info.get('num_compute_cores', 'Unknown')}")
print(f"Max threadgroup memory: {info.get('max_threadgroup_memory', 'Unknown')}")13.3 Compilation for Speed
MLX’s compile() function can significantly speed up repetitive computations.
import mlx.core as mx
import time
@mx.compile
def compiled_forward(model, x):
return model(x)
@mx.compile
def compiled_loss_and_grad(model, x, y):
logits = model(x)
loss = mx.mean((logits - y) ** 2)
return loss
# First call: compiles the graph (slower)
# Subsequent calls: uses compiled graph (faster)
# The compilation is cached based on input shapes and typesWhen to use compile():
- Training loops where shapes don’t change between iterations
- Inference with fixed batch sizes
- Repetitive computations in data processing pipelines
When NOT to use compile():
- Rapidly changing input shapes
- One-off computations
- Debugging (compiled functions are harder to inspect)
13.4 Gradient Checkpointing
For models too large to fit in memory, gradient checkpointing trades compute for memory by recomputing activations during the backward pass.
import mlx.core as mx
import mlx.nn as nn
class CheckpointedTransformerBlock(nn.Module):
"""Transformer block with gradient checkpointing."""
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.attention = nn.MultiHeadAttention(d_model, num_heads)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model),
)
self.dropout = nn.Dropout(dropout)
def _attention_forward(self, x):
"""Attention sub-layer (will be checkpointed)."""
residual = x
x = self.norm1(x)
x = self.attention(x)
return residual + self.dropout(x)
def _ffn_forward(self, x):
"""FFN sub-layer (will be checkpointed)."""
residual = x
x = self.norm2(x)
return residual + self.ffn(x)
def __call__(self, x):
# Use checkpointing to save memory
x = mx.checkpoint(self._attention_forward)(x)
x = mx.checkpoint(self._ffn_forward)(x)
return x13.5 Quantization
Quantization reduces model size by storing weights in lower precision.
import mlx.core as mx
import mlx.nn as nn
# Quantize a model to 4-bit
model = MyLargeModel()
# Method 1: Use nn.quantize
nn.quantize(model, bits=4, group_size=64)
# Method 2: Manual quantization
weights = model.layers[0].attention.w_q.weight
quantized, scales, biases = mx.quantize(weights, bits=4, group_size=64)
# Dequantize when needed for computation:
dequantized = mx.dequantize(quantized, scales, biases, bits=4, group_size=64)
# QuantizedLinear layer (automatic quantization for linear layers)
ql = nn.QuantizedLinear(
input_dims=4096,
output_dims=4096,
bias=False,
bits=4,
group_size=64,
)13.6 KV-Cache for Efficient Generation
During autoregressive generation, caching previous key-value pairs avoids recomputing attention for all previous positions.
import mlx.core as mx
class KVCache:
"""Key-Value cache for efficient autoregressive generation."""
def __init__(self):
self.keys = None
self.values = None
def update(self, new_keys, new_values):
"""Append new KV pairs to cache."""
if self.keys is None:
self.keys = new_keys
self.values = new_values
else:
self.keys = mx.concatenate([self.keys, new_keys], axis=2)
self.values = mx.concatenate([self.values, new_values], axis=2)
return self.keys, self.values
def reset(self):
"""Clear the cache."""
self.keys = None
self.values = None
def generate_with_cache(model, input_ids, max_tokens=100, temperature=0.8):
"""Generate text using KV-cache for efficiency."""
model.eval()
batch_size = input_ids.shape[0]
# Initialize KV caches for each layer
caches = [KVCache() for _ in range(len(model.layers))]
# Prefill: process the entire input at once
logits = model(input_ids, caches=caches)
next_logits = logits[:, -1, :]
generated = [input_ids]
for _ in range(max_tokens):
# Sample next token
next_logits = next_logits / temperature
probs = mx.softmax(next_logits, axis=-1)
next_token = mx.random.categorical(mx.log(probs + 1e-10)[:, None, :])[:, 0]
generated.append(next_token[:, None])
# Process only the new token (not the whole sequence again!)
logits = model(next_token[:, None], caches=caches)
next_logits = logits[:, -1, :]
mx.eval(next_logits)
return mx.concatenate(generated, axis=1)Without KV-cache: Each generation step processes all previous tokens -> O(n^2) compute
With KV-cache: Each step processes only the new token -> O(n) compute
13.7 Speculative Decoding
For even faster generation, speculative decoding uses a small “draft” model to predict multiple tokens, then validates them with the large model.
# Conceptual outline (not full implementation)
def speculative_decode(draft_model, target_model, input_ids,
max_tokens=100, num_speculative=4):
"""
Use a small draft model to speculate multiple tokens,
then validate with the target model.
"""
for _ in range(max_tokens // num_speculative):
# 1. Draft model generates num_speculative tokens quickly
draft_tokens = draft_model.generate(input_ids, max_tokens=num_speculative)
# 2. Target model validates all draft tokens in parallel
target_logits = target_model(draft_tokens)
# 3. Accept tokens that match, reject and resample from first mismatch
accepted = 0
for i in range(num_speculative):
target_prob = mx.softmax(target_logits[:, i, :])
draft_prob = mx.softmax(draft_model(draft_tokens[:, :i+1])[:, -1, :])
# Accept if target probability >= draft probability
if True: # Simplified; real implementation compares distributions
accepted += 1
else:
break
input_ids = draft_tokens[:, :accepted + 1]
return input_ids14: Distributed Training and Advanced Patterns
14.1 Data Parallelism
MLX supports distributed training across multiple GPUs.
import mlx.core as mx
# Check distributed availability
if mx.distributed.is_available():
# Initialize distributed
group = mx.distributed.init()
# Get world size and rank
world_size = group.size()
rank = group.rank()
print(f"Process {rank}/{world_size}")
# All-reduce: sum gradients across all processes
gradients = compute_gradients()
summed = mx.distributed.all_sum(gradients)
averaged = summed / world_size14.2 Tensor Parallelism
Split large model layers across multiple devices.
import mlx.nn as nn
# Shard a linear layer across devices
# Instead of one (4096, 4096) matrix, use two (4096, 2048) matrices
sharded_linear = nn.ShardedToAllLinear(
input_dims=4096,
output_dims=4096,
bias=False,
)
# Or use the shard_linear utility
linear = nn.Linear(4096, 4096)
sharded = nn.layers.distributed.shard_linear(linear, group)14.3 Stream-Based Parallelism
On a single Mac, you can use streams to overlap CPU and GPU computation.
import mlx.core as mx
# Create separate streams for CPU and GPU
cpu_stream = mx.new_stream(mx.cpu)
gpu_stream = mx.new_stream(mx.gpu)
# Run data preprocessing on CPU while GPU processes previous batch
with mx.stream(cpu_stream):
# Data preprocessing
batch = preprocess_data(raw_data)
with mx.stream(gpu_stream):
# Model forward pass
output = model(batch)
mx.synchronize() # Wait for both to complete14.4 Custom Metal Kernels
For performance-critical operations, you can write custom GPU kernels.
import mlx.core as mx
# Define a custom Metal kernel
kernel = mx.fast.metal_kernel(
name="my_kernel",
input_names=["input"],
output_names=["output"],
source="""
uint elem = thread_position_in_grid.x;
if (elem < input.shape[0]) {
output[elem] = input[elem] * 2.0f;
}
""",
grid=((1024,),),
threads=(256,),
)
# Use the kernel
x = mx.random.normal((1024,))
result = kernel(inputs=[x])
mx.eval(result)14.5 Profiling and Debugging
import mlx.core as mx
import time
# Simple timing
def time_operation(fn, *args, warmup=5, repeats=20):
"""Time an MLX operation."""
# Warmup
for _ in range(warmup):
result = fn(*args)
mx.eval(result)
# Time
times = []
for _ in range(repeats):
start = time.perf_counter()
result = fn(*args)
mx.eval(result)
end = time.perf_counter()
times.append(end - start)
avg_time = sum(times) / len(times)
std_time = (sum((t - avg_time)**2 for t in times) / len(times)) ** 0.5
print(f"Average time: {avg_time*1000:.2f} ms (±{std_time*1000:.2f} ms)")
return avg_time
# Metal capture for profiling
if mx.metal.is_available():
mx.metal.start_capture()
# ... run operations ...
mx.metal.stop_capture()
# Opens Xcode Instruments with the trace14.6 Exporting Models
import mlx.core as mx
# Export a function as a computation graph
@mx.export_function
def my_model(x):
return model(x)
# Export and save
mx.export_to_dot(my_model, "model_graph.dot") # GraphViz visualization
# Import an exported function
imported_fn = mx.import_function("exported_model.mlxfn")Continuous Learning and Next Steps
What you’ve learned in this guide:
- MLX fundamentals: arrays, operations, lazy evaluation
- Function transformations: grad, vmap, compile
- NLP basics: tokenization, embeddings, positional encoding
- Transformer architecture: attention, encoders, decoders
- Training: from character models to LoRA fine-tuning
- Optimization: quantization, KV-cache, compilation
Where to go next:
MLX Examples: https://github.com/ml-explore/mlx-examples
Study production-quality implementations
MLX LM: The mlx-lm package for more advanced LLM usage
HuggingFace MLX Community: https://huggingface.co/mlx-community Pre-converted models
MLX Documentation: https://ml-explore.github.io/mlx/build/html/index.html Full API reference
Advanced topics to explore:
- Vision-Language Models (LLaVA, CLIP) with
mlx-vlm - Speech models (Whisper) with MLX
- Diffusion models (Stable Diffusion, FLUX) with MLX
- Custom Metal kernels for performance-critical operations
- Multi-GPU distributed training
