microgpt
A few months ago, I watched a talk from Arnaud Bailly where he described his journey, reimplementing Andrej Karpathy's microgpt across 7 programming languages to understand how LLMs are working (the talk in French).
First, I would like to say that I praise his approach as, on my side, to be at ease with something. I need to somehow have an understanding 1 to 2 levels deeper than the concepts I'm dealing with.
The issue I have, probably my pride, is that the final Haskell implementation is 14 times slower than Common Lisp.
See the following table:
| Implementation | Time (s) | Ratio |
|---|---|---|
| Python | 60 | 1 |
| Haskell (naive) | 49 | 0.82 |
| Scheme | 27 | 0.45 |
| JavaScript | 17 | 0.28 |
| Forth | 12 | 0.2 |
| Haskell (optimized) | 10 | 0.16 |
| Idris | 10 | 0.16 |
| Common Lisp | 0.69 | 0.0115 |
| OCaml | 0.67 | 0.0112 |
I tried to re-implement myself shortly after I have seen seen his talk, as it takes a lot of mathematics to understand it, and I'm not good at it.
Few weeks ago, I was attending MuniHac and there was a talk about Accelerate, I told myself it was a good opportunity to use it in microgpt.
Note: unlike most of the other logs, I will not include the full code in it, but I'll add a link to each version.
I have used Gemini as a teacher in order to understand the maths behind microgpt and to review the code I was doing.
I started with a simple "translation" of microgpt, replacing mutation by pure computation see V0.
To illustrate how naive it is, I used only builtin data-structures such as lists, which I expect to be extremely slow given the random access patterns.
Look at the following examples:
data Node a = Node
{ val :: a,
deps :: [(NodeId, LocalPartialDerivative a)]
}
newtype NodeIdxMap a
= NodeIdxMap (IntMap.IntMap a)
deriving newtype (NFData)
type Vector = [NodeIdx Double]
type Matrix = [Vector]
-- | Parameters of a transformer layer, matrices are [nout][nin]
data Layer a = Layer
{ attnWq :: [[a]],
attnWk :: [[a]],
attnWv :: [[a]],
attnWo :: [[a]],
mlpFc1 :: [[a]],
mlpFc2 :: [[a]]
}
deriving stock (Generic, Functor, Foldable, Traversable)
deriving anyclass (NFData)
-- | The root microGPT structure.
-- The 'Foldable' order matches microgpt's flattened @params@ list.
data Weights a = Weights
{ wte :: [[a]],
wpe :: [[a]],
lmHead :: [[a]],
layers :: [Layer a]
}
deriving stock (Generic, Functor, Foldable, Traversable)
deriving anyclass (NFData)Since it would be mostly about matrices multiplications, I already have some optimizations ideas.
Arnaud used enough documents to reach 1 minute runtime in Python, as I want to iterate quickly, I have only used 100 documents, which already took more than three minutes to run.
See the following Python results:
benchmarking microGPT/forward (100 docs)
time 3.166 s (3.097 s .. 3.299 s)
std dev 56.846 ms (10 runs)
benchmarking microGPT/forward+backward (100 docs)
time 7.337 s (7.161 s .. 7.533 s)
std dev 128.941 ms (10 runs)
benchmarking microGPT/train (100 steps)
time 7.110 s (6.914 s .. 7.305 s)
std dev 120.801 ms (10 runs)And the following Haskell results:
benchmarking microGPT/forward (100 docs)
time 2.567 s (2.519 s .. 2.627 s)
1.000 R² (1.000 R² .. 1.000 R²)
mean 2.632 s (2.605 s .. 2.653 s)
std dev 33.20 ms (15.00 ms .. 43.60 ms)
variance introduced by outliers: 19% (moderately inflated)
benchmarking microGPT/forward+backward (100 docs)
time 5.251 s (4.808 s .. 5.456 s)
0.999 R² (0.998 R² .. 1.000 R²)
mean 5.299 s (5.238 s .. 5.344 s)
std dev 68.12 ms (48.88 ms .. 81.92 ms)
variance introduced by outliers: 19% (moderately inflated)
benchmarking microGPT/train (10 steps)
time 5.359 s (5.193 s .. 5.479 s)
1.000 R² (1.000 R² .. 1.000 R²)
mean 5.338 s (5.306 s .. 5.371 s)
std dev 40.72 ms (20.07 ms .. 53.98 ms)
variance introduced by outliers: 19% (moderately inflated)If we compare forward+backward results, we have a ratio of 0.71 (x1.4), which is lower than the 0.82 (x1.2) naive version of Arnaud, but close enough to be a starting point.
Note: I have tested with 10, 20, 500, the ratio is stable
As expected, I replaced list by Vector which is an efficient structure for random access with fusion optimizations, see the source, which yields the following results:
benchmarking microGPT/forward (100 docs)
time 2.932 s (2.718 s .. 3.101 s)
0.999 R² (0.999 R² .. 1.000 R²)
mean 2.990 s (2.933 s .. 3.089 s)
std dev 93.18 ms (8.340 ms .. 117.6 ms)
variance introduced by outliers: 19% (moderately inflated)
benchmarking microGPT/forward+backward (100 docs)
time 6.074 s (5.700 s .. 6.439 s)
1.000 R² (0.998 R² .. 1.000 R²)
mean 6.024 s (5.962 s .. 6.082 s)
std dev 69.28 ms (58.51 ms .. 74.69 ms)
variance introduced by outliers: 19% (moderately inflated)
benchmarking microGPT/train (10 steps)
time 6.368 s (6.072 s .. 6.639 s)
1.000 R² (0.999 R² .. 1.000 R²)
mean 6.391 s (6.344 s .. 6.463 s)
std dev 68.47 ms (6.394 ms .. 86.74 ms)
variance introduced by outliers: 19% (moderately inflated)
The ratio is now 0.82 (x1.2), which is actually slower, a good reminder that optimizing without profiling is a waste of time.
Now comes the time of introducing accelerate, see the source, which, only on my CPU gives the following results:
benchmarking microGPT/forward (100 docs)
time 33.04 ms (32.01 ms .. 34.28 ms)
0.996 R² (0.991 R² .. 0.999 R²)
mean 33.75 ms (33.34 ms .. 34.35 ms)
std dev 1.237 ms (779.2 μs .. 1.801 ms)
variance introduced by outliers: 11% (moderately inflated)
benchmarking microGPT/forward+backward (100 docs)
time 204.4 ms (194.5 ms .. 212.7 ms)
0.998 R² (0.994 R² .. 1.000 R²)
mean 204.3 ms (201.1 ms .. 208.3 ms)
std dev 5.098 ms (3.442 ms .. 6.908 ms)
variance introduced by outliers: 14% (moderately inflated)
benchmarking microGPT/train (10 steps)
time 108.1 ms (105.1 ms .. 111.3 ms)
0.998 R² (0.994 R² .. 1.000 R²)
mean 107.8 ms (106.1 ms .. 109.6 ms)
std dev 2.882 ms (2.183 ms .. 4.119 ms)The ratio is now 0.0278 (x36), which is close to the 0.0112 (x89.3) of OCaml.
When I have described my project to Gemini, it told me directly that my implementation would be slower due to the allocations, at that I should use ST, which is a Monad enable local and safe mutations.
I took also the opportunity to switch from Vector to UVector, its unboxed version.
Note: as a reminder, by default, typed are boxed in Haskell, which mean that, a constructor does not holds the values of its parameters, but pointers to the boxed value; it's essentially to support laziness in GHC.
See the source code, which gives the following results:
benchmarking microGPT/forward (100 docs)
time 4.310 ms (4.254 ms .. 4.411 ms)
0.998 R² (0.997 R² .. 0.999 R²)
mean 4.297 ms (4.269 ms .. 4.341 ms)
std dev 103.1 μs (64.91 μs .. 152.5 μs)
benchmarking microGPT/forward+backward (100 docs)
time 36.79 ms (35.87 ms .. 37.90 ms)
0.996 R² (0.989 R² .. 0.999 R²)
mean 36.67 ms (36.04 ms .. 37.50 ms)
std dev 1.432 ms (1.137 ms .. 1.958 ms)
variance introduced by outliers: 12% (moderately inflated)
benchmarking microGPT/train (10 steps)
time 13.24 ms (13.05 ms .. 13.46 ms)
0.998 R² (0.996 R² .. 0.999 R²)
mean 13.52 ms (13.41 ms .. 13.68 ms)
std dev 331.6 μs (243.3 μs .. 442.7 μs)The ratio is now 0.005 (x200), which beyond the Claude Code-generated implementations.
At this point I run out of ideas, then I recalled than Gemini was worrying about allocations.
After looking at the profiles, it seems that I did two obvious mistakes:
- Reading and dealing with
String, instead ofByteStringwhich are backed byVector, more efficient - Defining
NodeIdxMapasIntMapwhile most of indices are contiguous, which is another good fit forVector
See the sources, yielding the following results:
benchmarking microGPT/forward (100 docs)
time 4.309 ms (4.237 ms .. 4.393 ms)
0.997 R² (0.995 R² .. 0.999 R²)
mean 4.321 ms (4.282 ms .. 4.408 ms)
std dev 184.4 μs (94.26 μs .. 313.2 μs)
variance introduced by outliers: 22% (moderately inflated)
benchmarking microGPT/forward+backward (100 docs)
time 9.773 ms (9.523 ms .. 10.02 ms)
0.997 R² (0.994 R² .. 0.999 R²)
mean 9.693 ms (9.598 ms .. 9.845 ms)
std dev 332.1 μs (231.4 μs .. 541.9 μs)
variance introduced by outliers: 12% (moderately inflated)
benchmarking microGPT/train (10 steps)
time 12.97 ms (12.72 ms .. 13.16 ms)
0.999 R² (0.997 R² .. 0.999 R²)
mean 13.15 ms (13.03 ms .. 13.55 ms)
std dev 500.7 μs (182.7 μs .. 934.1 μs)
variance introduced by outliers: 14% (moderately inflated)The ratio is now 0.0013 (x751), nearly a x4 improvement, allocations are definitively always good to take care of.
Then, as a last resort, I have asked Gemini for ideas, in short, I had to remove parametric types which prevents optimizations, change some Float and Double to their unbowed versions FloatX4 and DoubleX8, adjust compilation flags.
See the sources, yielding the following results:
benchmarking microGPT/forward (100 docs)
time 1.144 ms (1.119 ms .. 1.171 ms)
0.994 R² (0.990 R² .. 0.997 R²)
mean 1.190 ms (1.163 ms .. 1.230 ms)
std dev 100.8 μs (69.00 μs .. 149.7 μs)
variance introduced by outliers: 65% (severely inflated)
benchmarking microGPT/forward+backward (100 docs)
time 3.675 ms (3.588 ms .. 3.756 ms)
0.992 R² (0.987 R² .. 0.996 R²)
mean 3.649 ms (3.522 ms .. 3.736 ms)
std dev 332.5 μs (270.3 μs .. 485.8 μs)
variance introduced by outliers: 59% (severely inflated)
benchmarking microGPT/forward/multicore (100 docs)
time 270.8 μs (254.7 μs .. 289.7 μs)
0.979 R² (0.965 R² .. 0.993 R²)
mean 265.9 μs (255.1 μs .. 279.3 μs)
std dev 43.02 μs (25.39 μs .. 69.49 μs)
variance introduced by outliers: 91% (severely inflated)
benchmarking microGPT/forward+backward/multicore (100 docs)
time 821.8 μs (785.3 μs .. 863.2 μs)
0.987 R² (0.982 R² .. 0.995 R²)
mean 821.7 μs (797.3 μs .. 861.5 μs)
std dev 101.2 μs (66.06 μs .. 177.4 μs)
variance introduced by outliers: 81% (severely inflated)
benchmarking microGPT/forward/fp32 (100 docs)
time 256.6 μs (250.0 μs .. 261.7 μs)
0.995 R² (0.992 R² .. 0.998 R²)
mean 261.4 μs (255.3 μs .. 281.2 μs)
std dev 33.50 μs (13.36 μs .. 65.22 μs)
variance introduced by outliers: 86% (severely inflated)
benchmarking microGPT/forward+backward/fp32 (100 docs)
time 608.4 μs (592.1 μs .. 630.9 μs)
0.991 R² (0.984 R² .. 0.996 R²)
mean 608.9 μs (597.3 μs .. 627.0 μs)
std dev 44.67 μs (31.82 μs .. 61.61 μs)
variance introduced by outliers: 62% (severely inflated)
benchmarking microGPT/train (10 steps)
time 3.942 ms (3.881 ms .. 4.003 ms)
0.998 R² (0.996 R² .. 0.999 R²)
mean 3.998 ms (3.956 ms .. 4.059 ms)
std dev 168.2 μs (98.84 μs .. 290.3 μs)
variance introduced by outliers: 23% (moderately inflated)Note: Gemini proposed multiple versions (FP64/FP32/Multicore)
Finally, ratios are between 0.005 (x200) and 0,00008 (x12,060), but results are so small, they are not reliable.
To have more accurate results, we should test against 1000 documents, let's start with the following results for Python:
benchmarking microGPT/forward (1000 docs)
time 34.728 s (34.097 s .. 35.544 s)
std dev 515.328 ms (10 runs)
benchmarking microGPT/forward+backward (1000 docs)
time 97.175 s (96.713 s .. 97.893 s)
std dev 391.146 ms (10 runs)
benchmarking microGPT/train (1000 steps)
time 77.638 s (76.969 s .. 78.634 s)
std dev 533.232 ms (10 runs)And we can repeat the process for the latest version, as follows:
benchmarking microGPT/forward (1000 docs)
time 12.10 ms (11.83 ms .. 12.42 ms)
0.996 R² (0.993 R² .. 0.999 R²)
mean 12.00 ms (11.88 ms .. 12.22 ms)
std dev 399.3 μs (273.5 μs .. 566.1 μs)
variance introduced by outliers: 10% (moderately inflated)
benchmarking microGPT/forward+backward (1000 docs)
time 38.32 ms (37.64 ms .. 38.79 ms)
0.999 R² (0.997 R² .. 1.000 R²)
mean 38.66 ms (38.25 ms .. 39.13 ms)
std dev 941.9 μs (691.5 μs .. 1.350 ms)
benchmarking microGPT/forward/multicore (1000 docs)
time 4.034 ms (3.909 ms .. 4.167 ms)
0.991 R² (0.983 R² .. 0.995 R²)
mean 3.996 ms (3.924 ms .. 4.118 ms)
std dev 273.3 μs (196.4 μs .. 381.7 μs)
variance introduced by outliers: 45% (moderately inflated)
benchmarking microGPT/forward+backward/multicore (1000 docs)
time 7.655 ms (7.394 ms .. 7.989 ms)
0.987 R² (0.977 R² .. 0.995 R²)
mean 7.738 ms (7.580 ms .. 7.976 ms)
std dev 578.2 μs (461.0 μs .. 766.3 μs)
variance introduced by outliers: 43% (moderately inflated)
benchmarking microGPT/forward/fp32 (1000 docs)
time 2.106 ms (2.065 ms .. 2.142 ms)
0.995 R² (0.992 R² .. 0.998 R²)
mean 2.120 ms (2.091 ms .. 2.165 ms)
std dev 112.5 μs (83.59 μs .. 176.1 μs)
variance introduced by outliers: 38% (moderately inflated)
benchmarking microGPT/forward+backward/fp32 (1000 docs)
time 6.583 ms (6.338 ms .. 6.802 ms)
0.993 R² (0.988 R² .. 0.996 R²)
mean 6.505 ms (6.417 ms .. 6.615 ms)
std dev 332.9 μs (261.9 μs .. 430.9 μs)
variance introduced by outliers: 28% (moderately inflated)
benchmarking microGPT/train (10 steps)
time 36.59 ms (36.35 ms .. 36.83 ms)
1.000 R² (0.999 R² .. 1.000 R²)
mean 36.59 ms (36.46 ms .. 36.80 ms)
std dev 348.4 μs (187.4 μs .. 489.6 μs)Which sets the ratios between 0,0004 (x2,500) to 0,00007 (x14,762).
In summary, here are the ratios for each version, as follows:
| Implementation | forward | forward+backward | train |
|---|---|---|---|
| Python | 1 | 1 | 1 |
| OCaml* | - | 0.0112 | - |
| Haskell V0 (State) | 0.81 | 0.71 | 0.75 |
| Haskell V1 (Vector) | 0.93 | 0.82 | 0.90 |
| Haskell V2 (Accelerate) | 0.0104 | 0.0278 | 0.0152 |
| Haskell V3 (ST) | 0.0014 | 0.0050 | 0.0019 |
| Haskell V4 (Alloc) | 0.0014 | 0.0013 | 0.0018 |
| Haskell V5 (LLVM) | 0.00036 | 0.00050 | 0.00055 |
| Haskell V5 (Multicore) | 0.00009 | 0.00011 | - |
| Haskell V5 (FP32) | 0.00008 | 0.00008 | - |
| Haskell V5 (1000 docs) | 0.00035 | 0.00039 | 0.00047 |
| Haskell V5 Multicore (1000 docs) | 0.00012 | 0.00008 | - |
| Haskell V5 FP32 (1000 docs) | 0.00006 | 0.00007 | - |
| Haskell V6 (1000 docs) | 0.00035 | 0.00044 | 0.00047 |
| Haskell V6 Multicore (1000 docs) | 0.00011 | 0.00007 | - |
| Haskell V6 FP32 (1000 docs) | 0.00006 | 0.00006 | - |
* Overall ratio from Arnaud Bailly's benchmark (not broken down by phase).
To highlight a bit of the speedups, as follows:
| Implementation | forward | forward+backward | train |
|---|---|---|---|
| Python | x1.0 | x1.0 | x1.0 |
| OCaml* | - | x89.3 | - |
| Haskell V0 (State) | x1.2 | x1.4 | x1.3 |
| Haskell V1 (Vector) | x1.1 | x1.2 | x1.1 |
| Haskell V2 (Accelerate) | x96 | x36 | x66 |
| Haskell V3 (ST) | x735 | x199 | x537 |
| Haskell V4 (Alloc) | x735 | x751 | x548 |
| Haskell V5 (LLVM) | x2,768 | x1,997 | x1,804 |
| Haskell V5 (Multicore) | x11,691 | x8,928 | - |
| Haskell V5 (FP32) | x12,338 | x12,060 | - |
| Haskell V5 (1000 docs) | x2,870 | x2,536 | x2,122 |
| Haskell V5 Multicore (1000 docs) | x8,609 | x12,694 | - |
| Haskell V5 FP32 (1000 docs) | x16,490 | x14,762 | - |
| Haskell V6 (1000 docs) | x2,894 | x2,265 | x2,132 |
| Haskell V6 Multicore (1000 docs) | x8,737 | x14,362 | - |
| Haskell V6 FP32 (1000 docs) | x15,930 | x16,877 | - |
* Overall speedup from Arnaud Bailly's benchmark.
We should also consider the size of the code, as follows:
| Implementation | Lines of code | Ratio (vs Python) |
|---|---|---|
| Python | 224 | 1.0x |
| Haskell V0 (State) | 576 | 2.6x |
| Haskell V1 (Vector) | 597 | 2.7x |
| Haskell V2 (Accelerate) | 597 | 2.7x |
| Haskell V3 (ST) | 2,887 | 12.9x |
| Haskell V4 (Alloc) | 2,988 | 13.3x |
| Haskell V5 (LLVM) | 4,402 | 19.7x |
| Haskell V6 (Bare) | 3,439 | 15.4x |
It's not surprising, but the more rewarding the optimization is, the more effort it takes.
It also confirms that, well-tuned, the programming language is not what determines the performance of the program.
PS: at this point, it is legit to criticize the methodology as the original implementation was in pure Python.
Removing Accelerate gives the following results:
benchmarking microGPT/forward (1000 docs)
time 12.00 ms (11.77 ms .. 12.27 ms)
0.998 R² (0.996 R² .. 0.999 R²)
mean 11.96 ms (11.87 ms .. 12.07 ms)
std dev 248.7 μs (203.0 μs .. 316.7 μs)
benchmarking microGPT/forward+backward (1000 docs)
time 42.91 ms (41.29 ms .. 44.46 ms)
0.997 R² (0.994 R² .. 0.999 R²)
mean 42.90 ms (42.26 ms .. 43.99 ms)
std dev 1.594 ms (822.8 μs .. 2.648 ms)
benchmarking microGPT/forward/multicore (1000 docs)
time 3.975 ms (3.837 ms .. 4.132 ms)
0.992 R² (0.988 R² .. 0.996 R²)
mean 4.013 ms (3.952 ms .. 4.084 ms)
std dev 224.4 μs (178.7 μs .. 275.0 μs)
variance introduced by outliers: 34% (moderately inflated)
benchmarking microGPT/forward+backward/multicore (1000 docs)
time 6.766 ms (6.452 ms .. 7.044 ms)
0.990 R² (0.983 R² .. 0.994 R²)
mean 6.897 ms (6.756 ms .. 7.635 ms)
std dev 674.5 μs (338.6 μs .. 1.460 ms)
variance introduced by outliers: 57% (severely inflated)
benchmarking microGPT/forward/fp32 (1000 docs)
time 2.180 ms (2.124 ms .. 2.234 ms)
0.996 R² (0.994 R² .. 0.997 R²)
mean 2.137 ms (2.110 ms .. 2.162 ms)
std dev 91.99 μs (76.57 μs .. 110.7 μs)
variance introduced by outliers: 27% (moderately inflated)
benchmarking microGPT/forward+backward/fp32 (1000 docs)
time 5.758 ms (5.638 ms .. 5.862 ms)
0.993 R² (0.987 R² .. 0.998 R²)
mean 5.964 ms (5.863 ms .. 6.340 ms)
std dev 542.8 μs (277.3 μs .. 1.021 ms)
variance introduced by outliers: 54% (severely inflated)
benchmarking microGPT/train (10 steps)
time 36.41 ms (36.22 ms .. 36.52 ms)
1.000 R² (1.000 R² .. 1.000 R²)
mean 36.38 ms (36.31 ms .. 36.48 ms)
std dev 182.6 μs (140.1 μs .. 242.1 μs)Which sets the ratio to 0,000078 in FP32 (x12,821), slightly better than the previous version (0,00009 / x11,145).
Less is more, especially in optimizations.