Benchmarking a native-MLX GPT (minGPT's gpt-mini) against PyTorch+MPS, partition by partition

[bahaehmimdi]

I've been benchmarking MLX vs. PyTorch+MPS on a couple of real model
ports (an M1 Pro), and wanted to share the GPT side since it validates
something transformer_lm already does, with numbers behind it, plus a
partition-level breakdown of why.

Setup

karpathy/minGPT's gpt-mini config (6 layer, 6 head, 192 embd, 2.71M
params), vendored unmodified, vs. a from-scratch mx.nn/mx.core port
matching the same architecture. Same compile pattern transformer_lm
already uses: mx.compile(step, inputs=state, outputs=state) (needed
because optimizer.update mutates state in place — plain
mx.compile(step) errors with "Attempting to eval an array without a
primitive").

Partition-by-partition (forward/backward/optimizer split apart, real shapes, batch 64 / block 128)

partition	PyTorch CPU	PyTorch MPS	native MLX	winner
embeddings	0.54ms	1.38ms	0.40ms	MLX
transformer block (fwd)	27.7ms	8.51ms	4.92ms	MLX
lm_head + loss (fwd)	0.79ms	0.56ms	0.42ms	MLX
forward+backward (whole model)	453.7ms	127.8ms	118.8ms	MLX
optimizer.step	10.6ms	7.66ms	4.52ms	MLX

MLX wins every partition tested here — including forward+backward
uncompiled (119ms vs. MPS's 128ms).

End to end

	ms/step	vs. MPS
PyTorch MPS	154.9ms	1.00x
native MLX, eager	124.7ms	1.25x
native MLX, compiled	88.7ms	1.76x

Stable across 5 separate runs (88.4–88.9ms).

The contrasting case, for balance

I ran the identical methodology against RIFE (video frame
interpolation, small-batch Conv2d + resize + backward-warp) on the same
machine, and got the opposite result: PyTorch MPS won every single
partition there, not just the aggregate. So "MLX wins" isn't a
backend-wide property — it's architecture-dependent, and worth actually
measuring per model rather than assuming either direction. One concrete
thing that did fall out of the RIFE side: nn.Upsample's bilinear
mode measured 2-4x slower than a matmul-based resize when upsampling —
filed as #4435 with a repro, since that looks like a real, fixable gap
rather than an architecture mismatch.

Happy to share the benchmark scripts if useful (they're a thin vendor of
minGPT + a from-scratch mlx_gpt.py, nothing that would make sense as
a PR on top of transformer_lm since the pattern's already there) —
mostly wanted to put real numbers behind it.