This project is a port of Andrej Karpathy's llm.c to Mojo. Visit llm.c for a detailed explanation of the original project.
Notes:
- llm.mojo is a port of an earlier (2024) version of llm.c (train_gpt2.c @ 2d43e5b), not the current upstream. llm.c has since optimized its CPU matmul (register tiling in
matmul_forward). - We are preparing this repo for the upcoming Mojo 1.0 release — work in progress; it currently builds against the Mojo 1.0.0b2 beta.
Before using llm.mojo for the first time, please run the following preparatory commands in a virtual environment:
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
python prepro_tinyshakespeare.py
python train_gpt2.pyIf you don't have it, install pixi:
curl -fsSL https://pixi.sh/install.sh | shStart the virtual environment and execute the training program:
pixi shell
mojo train_gpt2.mojoNote: The first time you run
pixi shell, it will automatically install all necessary dependencies defined inpixi.toml.
For a more detailed step-by-step guide including additional setup details and options, please refer to our detailed usage instructions.
Basic benchmark results: (M2 MacBook Pro)
-
Below are average training loop times, observed across the various implementations. Please note that these results are intended to provide a general comparison rather than precise, repeatable metrics.
-
We are running the OpenMP-enabled train_gpt2.c with 10 threads. (
OMP_NUM_THREADS=10 ./train_gpt2)
| Implementation | Average Training Loop Time | Throughput |
|---|---|---|
| train_gpt2.mojo | 1704 ms | 150 tok/s |
| train_gpt2.c (with OpenMP) | 1987 ms | 128 tok/s |
| train_gpt2.c (no OpenMP) | 7405 ms | 34 tok/s |
We ported test_gpt2.c from the original repository to Mojo to validate our port's functionality. For instructions on how to run this test and insights into the results it yields, please see our guide here.
llm.mojo began in 2024 as a response to Karpathy's then newly released llm.c. Its purpose was to show that Mojo could implement the same low-level, C-style program — raw pointers, manual memory management — while matching its performance.
Sustained interest in the repo has kept us updating it to track new Mojo releases, without expanding its scope. In that same spirit, our next planned milestone is a port to Mojo 1.0. Beyond 1.0 we have no concrete plans to develop the project further.
We occasionally weighed adding GPU support. That ground is now well covered by ulmentflam/llm.mojo, a GPU-accelerated port built on hand-written CUDA and Metal kernels via MAX.
Inspired by ulmentflam/llm.mojo, we are looking into speeding up CPU throughput by utilizing MAX's linalg GEMM kernel (linalg.matmul) in place of our hand-written matmul loops. See variants.md for the variants we have built and their speed comparison.
See changelog.md.
MIT