Porting microgpt to Futhark, Part III

This is Part III of porting Andrej Karpathy's microgpt to the data-parallel language Futhark. Part I covered the forward pass, and Part II covered training and inference. Now: the external driver program and some benchmarks for how the port scales.

At the end of Part II, the Futhark code wasn't quite ready to run on its own. Futhark compiles to standalone C or GPU kernels, but doesn't do arbitrary file I/O or string formatting. To actually run the pipeline, we need an external driver.

The driver program

The driver program is a Python script (demo.py). It downloads Karpathy's names.txt dataset (32,033 names), tokenizes characters into integers with a vocabulary of 27 (lowercase letters plus a BOS token), serializes the data to Futhark's standard text format on stdin, and decodes the output tokens back to characters.

On the Futhark side, a demo entry point bundles training and sampling into a single call:

entry demo [n_docs][max_doc_len]
  (vocab_size:  i64)
  (seed:        i32)
  (docs:       [n_docs][max_doc_len]i64)
  (doc_lens:   [n_docs]i64)
  (num_steps:   i64)
  (temperature: f32)
  (num_samples: i64)
  (infer_seed:  i32)
  : (f32, [num_samples][block_size]i64) =
  let (loss, wte, wpe, lmh, wq, wk, wv, wo, fc1, fc2) =
    train vocab_size seed docs doc_lens num_steps
  in (loss,
      generate wte wpe lmh wq wk wv wo fc1 fc2 temperature num_samples infer_seed)

In demo.py, this entry point is called by passing the formatted inputs to stdin and reading back stdout:

# Prepare serialized input in Futhark text format
data = "\n".join([
    f"{vocab_size}i64",
    "1337i32",
    "[" + ", ".join("[" + ", ".join(f"{t}i64" for t in row) + "]" for row in rows) + "]",
    "[" + ", ".join(f"{n}i64" for n in lens) + "]",
    f"{args.steps}i64",
    "0.5f32",
    f"{args.samples}i64",
    "42i32",
])

# Run compiled Futhark binary with the demo entry point
result = subprocess.run(
    [binary, "--entry-point", "demo"],
    input=data.encode(),
    capture_output=True,
)

# Parse output: loss (f32) followed by generated token IDs (i64)
output = result.stdout.decode()
loss = float(re.search(r"([-\d.e+]+)f32", output).group(1))
tokens = [int(m) for m in re.findall(r"(-?\d+)i64", output)]

Running demo.py against the sequential C backend (microgpt_c) runs 500 training steps and generates samples in a couple of seconds:

$ python3 demo.py --backend microgpt_c --steps 500 --samples 10
Downloading names.txt ...
num docs: 32033, vocab size: 27
expected initial loss: ~3.2958
training for 500 steps ...
final loss: 2.0173

--- inference (new, hallucinated names) ---
sample  1: hayeli
sample  2: heman
sample  3: aren
sample  4: kayli
sample  5: karin
sample  6: tynuin
sample  7: amans
sample  8: lelia
sample  9: koles
sample 10: alann

The loss drops steadily from the uniform baseline (-ln(1/27) ≈ 3.30) down to ~2.02, and it starts generating plausible new names. So the pipeline works end-to-end. Now for the real question: how does it scale?

Model configurations

To see how performance scales, I picked six model configurations bridging Karpathy's original tiny network up to a gpt2-small size architecture. All use a vocabulary of 27 characters. The parameter count formula (omitting biases and layer norm weights, as microgpt does) is:

params = n_embd * (2 * vocab_size + block_size) + 12 * n_layer * n_embd^2

The six benchmark configurations:

Note: gpt2-small is ~86M parameters here rather than the 117M of the original because microgpt omits biases, uses bias-free RMSNorm, and has a 27-token vocabulary instead of 50K.

Benchmark setup

I ran the benchmarks on one of our cluster machines with an Intel Xeon Platinum 8462Y+ CPU and an NVIDIA L40 GPU (Futhark 0.25.37, CUDA 12.2). Each run measured full training steps (forward pass, loss, backprop, and Adam update), with one warmup run and 10 measured repetitions.

Included in the comparison:

• Python: Karpathy's original implementation (1 thread, CPU).
• Futhark C: Compiled to sequential C (1 thread, CPU).
• Futhark CUDA: Compiled to CUDA (GPU).

I also tested Futhark's multicore CPU and OpenCL backends, but will leave those numbers out of the post. Multicore was slower than single-threaded C at almost every size. I have not rigorously benchmarked why, but I believe this is because the outer loops over steps, sequence positions, and layers are sequential dependencies, so Futhark can only parallelize the inner matrix operations via OpenMP. Below gpt2-small, thread synchronization overhead outweighs the extra cores, although 128-core multicore does pull ahead at the gpt2-small size. OpenCL worked, but was slightly slower than CUDA across the board and ran out of memory on gpt2-small, so reporting it doesn't add much to the results.

Median time per training step:

*: Failed with a Python recursion depth error

Scaling curves

Time per training step plotted against parameter count on a log-log scale:

The log-log scale obscures that in absolute terms the CUDA backend really does get substantially faster as networks get larger. Here are the two Futhark backends on a linear scale to illustrate that:

A few things I thought worth noting about these curves.

My initial motivation for choosing Futhark was that it can compile to parallel GPU code. But the single-threaded C backend is actually quite reasonable for this use-case at smaller network sizes. It gets about a 1700x speedup from the Python original on small networks, and more importantly, can scale to larger networks without crashing.

The scaling curves show a typical CPU vs. GPU crossover. On the tiny model, single-threaded C (0.20 ms) is actually 28x faster than CUDA (5.76 ms). But CUDA has a much shallower slope and pulls ahead around s128 (~800K parameters). By gpt2-small (86M parameters), single-threaded C takes about 9 minutes per step (542 s), while CUDA takes about 25 seconds, or about 22x faster.

Lines of code

Back in Part I, one question was whether we could port microgpt as 1-to-1 as possible and get much better scaling without losing too much concision. Did that work? Going by lines of code, the Futhark port is 2-2.5x as large, depending on whether you count the Python driver program. Excluding blank lines and comments:

• Karpathy's Python original (microgpt_orig.py): 160 lines
• Futhark kernels (microgpt.fut): 310 lines
• Python driver (demo.py): 87 lines

Breaking it down by component:

The biggest code-size improvement is in autodiff: using Futhark's built-in vjp2 saved 38 lines of custom code. On the other hand, including an explicit DIY autodiff was one of Karpathy's goals, so this may not be a fair comparison. In any case, that savings is quickly chipped away at by Futhark's purity and static type system. In Python, modifying weights in place during Adam takes 4 lines; in Futhark, threading immutable records through helper combinators (params_map2 and params_map3) adds about 80 lines of boilerplate. Another 50 lines go to flattening and unflattening records at the C/GPU entry point boundary.

* * *

Overall, I'm happy with the experiment. The Futhark port ended up about twice as long as Karpathy's Python original, mostly due to the bookkeeping of static array shapes, immutable parameter records, and entry-point marshaling. In return, we get fast automatic differentiation and the ability to scale on both CPU and GPU from the 4K toy model up to a GPT-2-small class network. The experience for a newcomer was overall positive enough that I can see using Futhark for future projects, though not exactly sure what yet.

The full code for the port is on GitHub.