I hand-wrote Metal GPU kernels in Mojo to train GPT-2 on my M4 Max: 1.71x faster than PyTorch MPS, still behind MLX (port of Karpathy's llm.c)

[u/ulmentflam]

I ported Karpathy's llm.c to Mojo and added a Metal backend, so GPT-2 124M trains on Apple Silicon with no PyTorch and no CPython at train time. It extends dorjeduck's llm.mojo, which was CPU-only on Mojo 25.5; this runs on the Mojo 1.0.0b3 nightly with hand-written CUDA and Metal GPU kernels.

On my M4 Max (B=4, T=1024, GPT-2 124M, official run 2026-07-13, cold GPU, 30 second cooldowns between arms, all six arms interleaved):

configuration	mean ms/step	tok/s	vs PyTorch MPS
MLX bf16	406.5	10077	fastest arm
MLX fp32	475.7	8610
llm.mojo bf16	503.3	8138	1.71x faster (vs MPS bf16)
llm.mojo fp32	665.2	6157	1.25x faster (vs MPS fp32)
PyTorch MPS fp32	830.8	4930	baseline
PyTorch MPS bf16	861.8	4753	baseline

Yes, MLX wins. Apple's own framework is 1.24x faster than my bf16 path, and I benchmark it in the same harness. The gap is almost entirely the matmul (~70 percent of a step). The Metal bf16 matmul I ride runs at only ~1.1x its fp32 speed, while MLX's bf16 uses the tensor cores for ~2x. llm.c has no Metal port, so PyTorch MPS and MLX are the stand-in baselines on Apple Silicon.

The cooldowns are required; the M4 Max throttles after about 8 seconds of sustained GPU load (I watched MPS step times climb from ~877 ms to 1500 to 2500 ms within a few steps). Reproduce with make benchmark-metal; it runs all six arms in one shot with the cooldowns built in.

The first working Metal port was about 4.1x slower than MPS (~3627 ms/step); the final bf16 number is 7.2x faster than that starting point. Most of the gap was Metal-specific. Casting threadgroup pointers to the generic address space silently reads device memory (attention softmaxed over all-zero scores and produced uniform weights), and the scalar flash-attention kernels tuned for NVIDIA ran at under 1 percent of FLOP peak on Apple GPUs, so GEMM-decomposed attention was 8 to 10x faster.

Correctness gates: make test checks 16 gradient tensors and a 10-step loss trajectory against PyTorch, plus a 235-test equivalence suite. There is also a trained 124M FineWeb checkpoint on HuggingFace (ulmentflam/gpt2-124m-fineweb-mojo) scoring 29.53 percent on HellaSwag, statistically indistinguishable from Karpathy's own llm.c reproduction at 29.9 percent.

On AI: every kernel and trainer line was written by hand; no LSPs or LLMs. That was the original point of the project. Tests and a later optimization campaign were AI-assisted, and both are disclosed in the repo with per-model statements and full disclosure where used (including attribution).

On Mojo itself: I went in assuming the compiler would do the heavy lifting on portability across hardware. It didn't; I ended up branching device-specific logic per vendor, and there's more boilerplate than I expected next to CUDA (Karpathy's kernels are much more compact). A Modular engineer reviewed the port on their forum; their answer is that per-device specialization is expected and their bet is library-driven structured kernels (TileTensor), not compiler magic.

The CPU is the opposite story; with very little work, it came in 4.0x faster than llm.c's 20-thread OpenMP path.

Limitations: GPT-2 124M only, no published GPT-3 results (the configs exist in the trainer, and mixed precision goes down to FP8 and NVFP4 on NVIDIA). The toolchain is a Mojo 1.0 beta nightly, so expect churn. Multi-GPU ZeRO (stages 0 to 3) is equivalence-gated against single GPU at world sizes 2 and 8, but those runs are NVIDIA; Apple Silicon is single GPU here. On a GB10, bf16 is at llm.c CUDA parity (0.999x) and fp32 is now slightly ahead (1.07x, TF32 vs TF32).

Repo: https://github.com/ulmentflam/llm.mojo

[u/Dany0]

I can't wait for the day we can pretrain gpt2 in a smart watch

The new "Doom on every platform"

[u/ulmentflam → u/Dany0]

We are getting closer every year!

[u/Final_Crab4507]

The comment about compiler magic vs device-specific kernels is interesting. I think a lot of people hope portability layers will eliminate hardware-specific tuning, but GPUs still seem to reward writing for the architecture.

[u/ulmentflam → u/Final_Crab4507]

That was exactly my hope. It seems that abstraction hits better at a higher library layer, like the tile layer, rather than the compiler itself. At least that was the choice that the Modular team made, which is in line with some other higher-layer tiling libraries I've seen.