|
|
.. _compile: |
|
|
|
|
|
Compilation |
|
|
=========== |
|
|
|
|
|
.. currentmodule:: mlx.core |
|
|
|
|
|
MLX has a :func:`compile` function transformation which compiles computation |
|
|
graphs. Function compilation results in smaller graphs by merging common work |
|
|
and fusing certain operations. In many cases this can lead to big improvements |
|
|
in run-time and memory use. |
|
|
|
|
|
Getting started with :func:`compile` is simple, but there are some edge cases |
|
|
that are good to be aware of for more complex graphs and advanced usage. |
|
|
|
|
|
Basics of Compile |
|
|
----------------- |
|
|
|
|
|
Let's start with a simple example: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def fun(x, y): |
|
|
return mx.exp(-x) + y |
|
|
|
|
|
x = mx.array(1.0) |
|
|
y = mx.array(2.0) |
|
|
|
|
|
# Regular call, no compilation |
|
|
# Prints: array(2.36788, dtype=float32) |
|
|
print(fun(x, y)) |
|
|
|
|
|
# Compile the function |
|
|
compiled_fun = mx.compile(fun) |
|
|
|
|
|
# Prints: array(2.36788, dtype=float32) |
|
|
print(compiled_fun(x, y)) |
|
|
|
|
|
The output of both the regular function and the compiled function is the same |
|
|
up to numerical precision. |
|
|
|
|
|
The first time you call a compiled function, MLX will build the compute |
|
|
graph, optimize it, and generate and compile code. This can be relatively |
|
|
slow. However, MLX will cache compiled functions, so calling a compiled |
|
|
function multiple times will not initiate a new compilation. This means you |
|
|
should typically compile functions that you plan to use more than once. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def fun(x, y): |
|
|
return mx.exp(-x) + y |
|
|
|
|
|
x = mx.array(1.0) |
|
|
y = mx.array(2.0) |
|
|
|
|
|
compiled_fun = mx.compile(fun) |
|
|
|
|
|
# Compiled here |
|
|
compiled_fun(x, y) |
|
|
|
|
|
# Not compiled again |
|
|
compiled_fun(x, y) |
|
|
|
|
|
# Not compiled again |
|
|
mx.compile(fun)(x, y) |
|
|
|
|
|
There are some important cases to be aware of that can cause a function to |
|
|
be recompiled: |
|
|
|
|
|
* Changing the shape or number of dimensions |
|
|
* Changing the type of any of the inputs |
|
|
* Changing the number of inputs to the function |
|
|
|
|
|
In certain cases only some of the compilation stack will be rerun (for |
|
|
example when changing the shapes) and in other cases the full compilation |
|
|
stack will be rerun (for example when changing the types). In general you |
|
|
should avoid compiling functions too frequently. |
|
|
|
|
|
Another idiom to watch out for is compiling functions which get created and |
|
|
destroyed frequently. This can happen, for example, when compiling an anonymous |
|
|
function in a loop: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
a = mx.array(1.0) |
|
|
# Don't do this, compiles lambda at each iteration |
|
|
for _ in range(5): |
|
|
mx.compile(lambda x: mx.exp(mx.abs(x)))(a) |
|
|
|
|
|
Example Speedup |
|
|
--------------- |
|
|
|
|
|
The :func:`mlx.nn.gelu` is a nonlinear activation function commonly used with |
|
|
Transformer-based models. The implementation involves several unary and binary |
|
|
element-wise operations: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def gelu(x): |
|
|
return x * (1 + mx.erf(x / math.sqrt(2))) / 2 |
|
|
|
|
|
If you use this function with small arrays, it will be overhead bound. If you |
|
|
use it with large arrays it will be memory bandwidth bound. However, all of |
|
|
the operations in the ``gelu`` are fusible into a single kernel with |
|
|
:func:`compile`. This can speedup both cases considerably. |
|
|
|
|
|
Let's compare the runtime of the regular function versus the compiled |
|
|
function. We'll use the following timing helper which does a warm up and |
|
|
handles synchronization: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
import time |
|
|
|
|
|
def timeit(fun, x): |
|
|
# warm up |
|
|
for _ in range(10): |
|
|
mx.eval(fun(x)) |
|
|
|
|
|
tic = time.perf_counter() |
|
|
for _ in range(100): |
|
|
mx.eval(fun(x)) |
|
|
toc = time.perf_counter() |
|
|
tpi = 1e3 * (toc - tic) / 100 |
|
|
print(f"Time per iteration {tpi:.3f} (ms)") |
|
|
|
|
|
|
|
|
Now make an array, and benchmark both functions: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
x = mx.random.uniform(shape=(32, 1000, 4096)) |
|
|
timeit(gelu, x) |
|
|
timeit(mx.compile(gelu), x) |
|
|
|
|
|
On an M1 Max the times are 15.5 and 3.1 milliseconds. The compiled ``gelu`` is |
|
|
five times faster. |
|
|
|
|
|
Debugging |
|
|
--------- |
|
|
|
|
|
When a compiled function is first called, it is traced with placeholder |
|
|
inputs. This means you can't evaluate arrays (for example to print their |
|
|
contents) inside compiled functions. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
@mx.compile |
|
|
def fun(x): |
|
|
z = -x |
|
|
print(z) # Crash |
|
|
return mx.exp(z) |
|
|
|
|
|
fun(mx.array(5.0)) |
|
|
|
|
|
For debugging, inspecting arrays can be helpful. One way to do that is to |
|
|
globally disable compilation using the :func:`disable_compile` function or |
|
|
``MLX_DISABLE_COMPILE`` flag. For example the following is okay even though |
|
|
``fun`` is compiled: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
@mx.compile |
|
|
def fun(x): |
|
|
z = -x |
|
|
print(z) # Okay |
|
|
return mx.exp(z) |
|
|
|
|
|
mx.disable_compile() |
|
|
fun(mx.array(5.0)) |
|
|
|
|
|
|
|
|
Pure Functions |
|
|
-------------- |
|
|
|
|
|
Compiled functions are intended to be *pure*; that is they should not have side |
|
|
effects. For example: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
state = [] |
|
|
|
|
|
@mx.compile |
|
|
def fun(x, y): |
|
|
z = x + y |
|
|
state.append(z) |
|
|
return mx.exp(z) |
|
|
|
|
|
fun(mx.array(1.0), mx.array(2.0)) |
|
|
# Crash! |
|
|
print(state) |
|
|
|
|
|
After the first call of ``fun``, the ``state`` list will hold a placeholder |
|
|
array. The placeholder does not have any data; it is only used to build the |
|
|
computation graph. Printing such an array results in a crash. |
|
|
|
|
|
You have two options to deal with this. The first option is to simply return |
|
|
``state`` as an output: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
state = [] |
|
|
|
|
|
@mx.compile |
|
|
def fun(x, y): |
|
|
z = x + y |
|
|
state.append(z) |
|
|
return mx.exp(z), state |
|
|
|
|
|
_, state = fun(mx.array(1.0), mx.array(2.0)) |
|
|
# Prints [array(3, dtype=float32)] |
|
|
print(state) |
|
|
|
|
|
In some cases returning updated state can be pretty inconvenient. Hence, |
|
|
:func:`compile` has a parameter to capture implicit outputs: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
from functools import partial |
|
|
|
|
|
state = [] |
|
|
|
|
|
# Tell compile to capture state as an output |
|
|
@partial(mx.compile, outputs=state) |
|
|
def fun(x, y): |
|
|
z = x + y |
|
|
state.append(z) |
|
|
return mx.exp(z) |
|
|
|
|
|
fun(mx.array(1.0), mx.array(2.0)) |
|
|
# Prints [array(3, dtype=float32)] |
|
|
print(state) |
|
|
|
|
|
This is particularly useful for compiling a function which includes an update |
|
|
to a container of arrays, as is commonly done when training the parameters of a |
|
|
:class:`mlx.nn.Module`. |
|
|
|
|
|
Compiled functions will also treat any inputs not in the parameter list as |
|
|
constants. For example: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
state = [mx.array(1.0)] |
|
|
|
|
|
@mx.compile |
|
|
def fun(x): |
|
|
return x + state[0] |
|
|
|
|
|
# Prints array(2, dtype=float32) |
|
|
print(fun(mx.array(1.0))) |
|
|
|
|
|
# Update state |
|
|
state[0] = mx.array(5.0) |
|
|
|
|
|
# Still prints array(2, dtype=float32) |
|
|
print(fun(mx.array(1.0))) |
|
|
|
|
|
In order to have the change of state reflected in the outputs of ``fun`` you |
|
|
again have two options. The first option is to simply pass ``state`` as input |
|
|
to the function. In some cases this can be pretty inconvenient. Hence, |
|
|
:func:`compile` also has a parameter to capture implicit inputs: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
from functools import partial |
|
|
state = [mx.array(1.0)] |
|
|
|
|
|
# Tell compile to capture state as an input |
|
|
@partial(mx.compile, inputs=state) |
|
|
def fun(x): |
|
|
return x + state[0] |
|
|
|
|
|
# Prints array(2, dtype=float32) |
|
|
print(fun(mx.array(1.0))) |
|
|
|
|
|
# Update state |
|
|
state[0] = mx.array(5.0) |
|
|
|
|
|
# Prints array(6, dtype=float32) |
|
|
print(fun(mx.array(1.0))) |
|
|
|
|
|
|
|
|
Compiling Training Graphs |
|
|
------------------------- |
|
|
|
|
|
This section will step through how to use :func:`compile` with a simple example |
|
|
of a common setup: training a model with :obj:`mlx.nn.Module` using an |
|
|
:obj:`mlx.optimizers.Optimizer` with state. We will show how to compile the |
|
|
full forward, backward, and update with :func:`compile`. |
|
|
|
|
|
To start, here is the simple example without any compilation: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
import mlx.core as mx |
|
|
import mlx.nn as nn |
|
|
import mlx.optimizers as optim |
|
|
|
|
|
# 4 examples with 10 features each |
|
|
x = mx.random.uniform(shape=(4, 10)) |
|
|
|
|
|
# 0, 1 targets |
|
|
y = mx.array([0, 1, 0, 1]) |
|
|
|
|
|
# Simple linear model |
|
|
model = nn.Linear(10, 1) |
|
|
|
|
|
# SGD with momentum |
|
|
optimizer = optim.SGD(learning_rate=0.1, momentum=0.8) |
|
|
|
|
|
def loss_fn(model, x, y): |
|
|
logits = model(x).squeeze() |
|
|
return nn.losses.binary_cross_entropy(logits, y) |
|
|
|
|
|
loss_and_grad_fn = nn.value_and_grad(model, loss_fn) |
|
|
|
|
|
# Perform 10 steps of gradient descent |
|
|
for it in range(10): |
|
|
loss, grads = loss_and_grad_fn(model, x, y) |
|
|
optimizer.update(model, grads) |
|
|
mx.eval(model.parameters(), optimizer.state) |
|
|
|
|
|
To compile the update we can put it all in a function and compile it with the |
|
|
appropriate input and output captures. Here's the same example but compiled: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
import mlx.core as mx |
|
|
import mlx.nn as nn |
|
|
import mlx.optimizers as optim |
|
|
from functools import partial |
|
|
|
|
|
# 4 examples with 10 features each |
|
|
x = mx.random.uniform(shape=(4, 10)) |
|
|
|
|
|
# 0, 1 targets |
|
|
y = mx.array([0, 1, 0, 1]) |
|
|
|
|
|
# Simple linear model |
|
|
model = nn.Linear(10, 1) |
|
|
|
|
|
# SGD with momentum |
|
|
optimizer = optim.SGD(learning_rate=0.1, momentum=0.8) |
|
|
|
|
|
def loss_fn(model, x, y): |
|
|
logits = model(x).squeeze() |
|
|
return nn.losses.binary_cross_entropy(logits, y) |
|
|
|
|
|
# The state that will be captured as input and output |
|
|
state = [model.state, optimizer.state] |
|
|
|
|
|
@partial(mx.compile, inputs=state, outputs=state) |
|
|
def step(x, y): |
|
|
loss_and_grad_fn = nn.value_and_grad(model, loss_fn) |
|
|
loss, grads = loss_and_grad_fn(model, x, y) |
|
|
optimizer.update(model, grads) |
|
|
return loss |
|
|
|
|
|
# Perform 10 steps of gradient descent |
|
|
for it in range(10): |
|
|
loss = step(x, y) |
|
|
# Evaluate the model and optimizer state |
|
|
mx.eval(state) |
|
|
print(loss) |
|
|
|
|
|
|
|
|
.. note:: |
|
|
|
|
|
If you are using a module which performs random sampling such as |
|
|
:func:`mlx.nn.Dropout`, make sure you also include ``mx.random.state`` in the |
|
|
``state`` captured by :func:`compile`, i.e. ``state = [model.state, |
|
|
optimizer.state, mx.random.state]``. |
|
|
|
|
|
|
|
|
.. note:: |
|
|
|
|
|
For more examples of compiling full training graphs checkout the `MLX |
|
|
Examples <https: |
|
|
|
|
|
Transformations with Compile |
|
|
---------------------------- |
|
|
|
|
|
In MLX function transformations are composable. You can apply any function |
|
|
transformation to the output of any other function transformation. For more on |
|
|
this, see the documentation on :ref:`function transforms |
|
|
<function_transforms>`. |
|
|
|
|
|
Compiling transformed functions works just as expected: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
grad_fn = mx.grad(mx.exp) |
|
|
|
|
|
compiled_grad_fn = mx.compile(grad_fn) |
|
|
|
|
|
# Prints: array(2.71828, dtype=float32) |
|
|
print(grad_fn(mx.array(1.0))) |
|
|
|
|
|
# Also prints: array(2.71828, dtype=float32) |
|
|
print(compiled_grad_fn(mx.array(1.0))) |
|
|
|
|
|
.. note:: |
|
|
|
|
|
In order to compile as much as possible, a transformation of a compiled |
|
|
function will not by default be compiled. To compile the transformed |
|
|
function simply pass it through :func:`compile`. |
|
|
|
|
|
You can also compile functions which themselves call compiled functions. A |
|
|
good practice is to compile the outer most function to give :func:`compile` |
|
|
the most opportunity to optimize the computation graph: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
@mx.compile |
|
|
def inner(x): |
|
|
return mx.exp(-mx.abs(x)) |
|
|
|
|
|
def outer(x): |
|
|
inner(inner(x)) |
|
|
|
|
|
# Compiling the outer function is good to do as it will likely |
|
|
# be faster even though the inner functions are compiled |
|
|
fun = mx.compile(outer) |
|
|
|
|
|
|
|
|
|
|
|
.. _shapeless_compile: |
|
|
|
|
|
Shapeless Compilation |
|
|
--------------------- |
|
|
|
|
|
When the shape of an input to a compiled function changes, the function is |
|
|
recompiled. You can compile a function once and run it on inputs with |
|
|
variable shapes by specifying ``shapeless=True`` to :func:`compile`. In this |
|
|
case changes to the shapes of the inputs do not cause the function to be |
|
|
recompiled. |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def fun(x, y): |
|
|
return mx.abs(x + y) |
|
|
|
|
|
compiled_fun = mx.compile(fun, shapeless=True) |
|
|
|
|
|
x = mx.array(1.0) |
|
|
y = mx.array(-2.0) |
|
|
|
|
|
# Firt call compiles the function |
|
|
print(compiled_fun(x, y)) |
|
|
|
|
|
# Second call with different shapes |
|
|
# does not recompile the function |
|
|
x = mx.array([1.0, -6.0]) |
|
|
y = mx.array([-2.0, 3.0]) |
|
|
print(compiled_fun(x, y)) |
|
|
|
|
|
|
|
|
Use shapeless compilations carefully. Since compilation is not triggered when |
|
|
shapes change, any graphs which are conditional on the input shapes will not |
|
|
work as expected. Shape-dependent computations are common and sometimes subtle |
|
|
to detect. For example: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def fun(x): |
|
|
return x.reshape(x.shape[0] * x.shape[1], -1) |
|
|
|
|
|
compiled_fun = mx.compile(fun, shapeless=True) |
|
|
|
|
|
x = mx.random.uniform(shape=(2, 3, 4)) |
|
|
|
|
|
out = compiled_fun(x) |
|
|
|
|
|
x = mx.random.uniform(shape=(5, 5, 3)) |
|
|
|
|
|
# Error, can't reshape (5, 5, 3) to (6, -1) |
|
|
out = compiled_fun(x) |
|
|
|
|
|
The second call to the ``compiled_fun`` fails because of the call to |
|
|
:func:`reshape` which uses the static shape of ``x`` in the first call. We can |
|
|
fix this by using :func:`flatten` to avoid hardcoding the shape of ``x``: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
def fun(x): |
|
|
return x.flatten(0, 1) |
|
|
|
|
|
compiled_fun = mx.compile(fun, shapeless=True) |
|
|
|
|
|
x = mx.random.uniform(shape=(2, 3, 4)) |
|
|
|
|
|
out = compiled_fun(x) |
|
|
|
|
|
x = mx.random.uniform(shape=(5, 5, 3)) |
|
|
|
|
|
# Ok |
|
|
out = compiled_fun(x) |
|
|
|
