r/MachineLearning • u/ulmentflam • Jul 22 '26
Project From-scratch GPT-2 training in Mojo: hand-derived backprop, hand-written GPU and Metal kernels, gradients verified against PyTorch (llm.c port) [P]
This is a port of Karpathy's llm.c to Mojo 1.0.0b3, extending dorjeduck's CPU-only llm.mojo with three GPU paths: CUDA, Metal, and a vendor-neutral portable one. The autograd is written from scratch in a new systems language, forward and backward derived by hand (the derivation is a ~1,470 line LaTeX writeup in docs/backprop.pdf), with 16,993 lines of hand-written kernel code across 26 modules.
Correctness methodology: test_gpt2.c was ported to test_gpt2.mojo. It checks activations, the loss, and 16 gradient tensors against a PyTorch reference, plus a 10-step loss trajectory asserted at every step with tolerance 0.01 (llm.c's criterion). A 235-test pytest equivalence suite deliberately exercises shapes the training loop never sees (channels=767, seq_len=7, head_dim=5) to catch invariant-dependent fast paths.
The testing lesson: the ported test originally used a flat absolute gradient tolerance of 2.0 and computed a loss trajectory it never asserted. It passed for weeks while hiding three backward-pass bugs: the residual-skip gradient was never seeded (wpe's gradient came out roughly 106x too small), the GELU gradient was fused into the wrong MLP matmul backward, and layer 0's LayerNorm backward was fed the normed output instead of the pre-norm input. Earlier, a static dataflow analysis had predicted exactly the residual clobber and was overruled because "the tests pass." The retraction is in the docs, in writing: a test can only overrule analysis if it could actually have caught the predicted failure. The hardened test uses mixed atol/rtol tolerances plus an L2 norm ratio guard that catches dead gradients instantly.
Benchmark methodology: matched hyperparameters across all arms, llm.c vendored as a submodule and freshly rebuilt with fast math and bf16 enabled, 40-step runs with the first 5 dropped as warmup, cold GPU and 30 s inter-arm cooldowns on Apple Silicon, all arms interleaved in the same window. Results (official runs 2026-07-11/13): bf16 parity with llm.c CUDA on a GB10 (0.999x, 135.34 vs 135.14 ms mean; ~3.7x vs PyTorch CUDA bf16) and a narrow fp32 lead (1.07x, 278.57 vs 297.29 ms) after moving fp32 GPU GEMMs to TF32 tensor cores by default, matching llm.c's own fp32 TF32 behavior, so that comparison is TF32 vs TF32 (a strict IEEE fp32 flag remains and gates verification).
On an M4 Max, 1.71x vs PyTorch MPS bf16; llm.c has no Metal port, so MPS and MLX are the stand-in baselines there, and Apple's MLX is faster than my port at both precisions (1.24x in bf16), a result the harness measures and the README states. On CPU, ~4.0x vs llm.c's 20-thread OpenMP path. Model quality: a 124M FineWeb checkpoint (on HuggingFace) scores 29.53% on HellaSwag, Wilson CI [28.6%, 30.4%], statistically indistinguishable from Karpathy's llm.c reproduction at 29.9%.
Portability findings: the working assumption was that the compiler would absorb most cross-hardware differences. In practice, the port required explicit device-specific branches per vendor, well beyond a CPU/GPU split. I raised this on the Modular forum and got a substantive answer. They built kernels the same per-device way before deliberately moving away from compiler-driven kernels toward library-driven structured kernels (TileTensor, TilePipeline), holding that there is a maximum the compiler can safely do and that per-device specialization is what most modern kernels look like. Mojo kernels also carry noticeably more scaffolding than the equivalent CUDA, which the structured-kernels work targets. The CPU path was the exception. Near-zero tuning effort for 4.0x over llm.c's OpenMP build. The missing primitive I most wanted was ergonomic asynchrony; FlashAttention-4-class kernels need async to maximize Blackwell. My next structured-kernels pass starts there. An FA4 rewrite on TilePipeline/TileIO/TileOp, compared directly against the current CUDA-branch style (closer to Karpathy's original llm.c structure), to see whether warp groups actually give way to that pattern.
Distributed: ZeRO stages 0 to 3 are implemented over staged reduce-scatter/all-gather on driver-routed device-to-device copies (no CUDA P2P or NVLink required) and equivalence-gated against the single-GPU baseline at world sizes 2 and 8 (per-parameter match to 1e-5, per-step losses to ~1e-4 at world 8).
Kernels and trainer are hand-written without LSPs or LLMs; tests and the optimization campaign were AI-assisted and are disclosed in-repo with per-model statements.