I built something on 1.0.0b3 for the occasion: a port of Karpathy’s llm.c that extends @dorjeduck’s llm.
with hand-written CUDA and Metal GPU kernels, training GPT-2 124M end to end in pure Mojo. On an M4 Max the Metal path runs at 498.92 ms/step in bf16, 1.72x faster than PyTorch MPS (llm.c has no Metal port, so MPS is the baseline there). On a GB10 it reached bf16 parity with llm.c CUDA. Comptime vendor dispatch and MAX’s linalg.matmul did a lot of heavy lifting; the main thing I hit was a SHARED-to-GENERIC address-space cast that silently reads zeros on Metal AIR (details and a probe in the repo). Full writeup: Llm.🔥: GPT-2 training in pure Mojo, with hand-written CUDA and Metal GPU kernels. Correctness is gated against PyTorch gradients.
1 Like