You Should Not Need a CUDA Box to Look Inside a Model

The interpretability ecosystem grew up on NVIDIA hardware. The libraries that make activation capture and intervention pleasant — nnsight, nnterp — assume a CUDA stack underneath. If you own an Apple Silicon Mac, the practical answer has been to rent a GPU somewhere else before you can inspect anything at all.

That is a strange barrier for a field whose whole premise is that more people should be able to scrutinise these systems. Plenty of good interpretability work starts small: a hypothesis about a couple of layers, tested on a 1B model, on the machine you already own.

mlxterp exists to remove that first step. It brings mechanistic interpretability to MLX, Apple’s array framework, taking the API design of nnsight and nnterp into the MLX ecosystem.

Repository: github.com/coairesearch/mlxterp — open source, MIT licensed. Issues and contributions welcome.

Capture Everything, With a Context Manager

The core idea is that inspecting a model should not restructure your code. You wrap a model, you trace a forward pass, and the activations are there:

from mlxterp import InterpretableModel
from mlx_lm import load

base_model, tokenizer = load('mlx-community/Llama-3.2-1B-Instruct-4bit')
model = InterpretableModel(base_model, tokenizer=tokenizer)

with model.trace("Hello, how are you?") as trace:
    pass  # ~196 activations captured automatically

layer_5_attn = trace.activations['model.model.layers.5.self_attn']
q_proj_3     = trace.activations['model.model.layers.3.self_attn.q_proj']
mlp_7_gate   = trace.activations['model.model.layers.7.mlp.gate_proj']

A single forward pass captures roughly 196 activations at a useful granularity — embeddings, per-layer outputs, attention blocks, individual Q/K/V and output projections, RoPE, MLP gate/up/down projections, and both layer norms.

The capture is model-agnostic. There are no per-architecture implementations to write: it works with any MLX model, including ones you define yourself, because it hooks the module tree rather than pattern-matching known architectures.

Interventions

Reading activations is half of interpretability. The other half is changing them and seeing what happens — activation patching, steering, ablation. Those are the experiments that distinguish a component that matters from one that merely correlates, so the intervention API is a first-class part of the library rather than an add-on.

from mlxterp import interventions as iv

# Halve the contribution of one attention block
with model.trace("Test input",
                 interventions={'model.model.layers.5.self_attn': iv.scale(0.5)}) as trace:
    output = trace.activations['__model_output__']

# Add a steering vector
with model.trace("Test",
                 interventions={'model.model.layers.10': iv.add_vector(steering_vector)}):
    steered = model.output.save()

The built-in interventions cover the usual operations — zero_out, scale, add_vector, replace_with, clamp and noise — and compose() chains them, so multiple edits can be applied to the same hook point in one pass.

Tokenizer integration is included, which matters more than it sounds: most interpretability questions are about a specific token, so mapping between text, token IDs and positions needs to be trivial rather than something you reimplement each time.

Why Apple Silicon

MLX is built for Apple’s unified memory architecture, where CPU and GPU share the same pool. For interpretability that is a genuinely good fit, because capturing hundreds of activations per pass is memory-hungry in a way that discrete-GPU setups feel immediately. Metal acceleration handles the compute.

The practical consequence is that a reasonably specified MacBook becomes a usable interpretability workstation for small and mid-sized models — no cluster, no cloud account, no queue.

Where It Fits

mlxterp and eDIF solve the same accessibility problem from opposite ends. eDIF gives researchers remote access to frontier-scale models on shared GPU infrastructure. mlxterp makes the local end of the workflow work on hardware people already have.

Develop a method locally on a 1B model, then run it against a 70B model on shared infrastructure when it is ready. Neither step should require owning a datacentre.

Get It

mlxterp is open source under the MIT licence at github.com/coairesearch/mlxterp.

git clone https://github.com/coairesearch/mlxterp
cd mlxterp
uv sync
uv add mlx-lm      # to load real models

Or with pip:

pip install -e .
pip install mlx-lm

Requires Apple Silicon and MLX. Inspired by nnsight and nnterp. Issues, feature requests and pull requests are welcome on GitHub.