Mnemonic

A chat model I wrote in assembly

Mnemonic is a small language model that I wrote by hand in x86-64 assembly and in PTX, the assembly language of NVIDIA's GPUs. That covers everything it took to make it: reading the training data, the tokenizer, the training on the GPU, shrinking the weights afterwards, and the program you chat with. None of it is C or Python, and it doesn't use the CUDA toolkit. The only outside code it calls is the Windows API and the NVIDIA driver.

I trained it at home on an RTX 5070 Ti, which is a graphics card made for games. It's nowhere near the models people usually chat with, but it knows its name, keeps up a simple conversation and explains everyday things reasonably well.

You can try it further down this page. I built it in eleven stages and wrote a chapter about each one, starting with the toolchain.

  1. DataFineWeb-Edu and smol-smoltalk come as Parquet files, read by my own Parquet, snappy and zstd code.
  2. TokenizerIt is byte-level BPE with 32,768 tokens, trained on a sample and then run over billions of tokens.
  3. TrainingThe forward pass, the backward pass and AdamW run in PTX on an RTX 5070 Ti.
  4. QuantizationThe trained weights are shrunk to int8 or int4.
  5. ChatOne engine runs on the CPU with AVX2 and VNNI, and a second one in WebAssembly runs in the browser.
From the raw data to a chat. Every step on this line is assembly I wrote.

Try it

It runs in your browser on your own computer, so the model has to download before it can answer, and nothing downloads until you press the button. It handles small talk and simple explanations, but it often gets facts and arithmetic wrong, and it does better with one question at a time than with a long back and forth.

  • The newest version, with 24 layers. It was pretrained on 10 billion tokens of FineWeb-Edu and then tuned for chat on smol-smoltalk. The weights are int8, which loses almost nothing against full precision at a quarter of the size.

How it was built

  • The toolchainNASM, the linker, a build script, and a hello world with no C runtime behind it.
  • FoundationsFiles, memory, printing floats, timers, random numbers, threads and a progress bar.
  • The data pipelineReading Parquet files, with their snappy and zstd compression, to get the training text out.
  • The tokenizerByte-level BPE, trained in assembly on a sample of the data and then run over all of it.
  • Talking to the GPUCalling the CUDA driver API from assembly and running the first PTX kernel.
  • A transformer in PTXThe forward and backward passes, gradient checks, tensor core matmuls and flash attention.
  • PretrainingTraining the 126M and then the 304M model on FineWeb-Edu with one RTX 5070 Ti.
  • Teaching it to chatFine-tuning on smol-smoltalk, giving it an identity, and two bugs that made it look dumb.
  • Quantization and the CPUint8 and int4 weights, AVX2 and VNNI matrix code, a KV cache and the console chat.
  • Making it fasterA profiler, a thread pool that spins, a second CUDA stream and better int4 scales.
  • In the browserThe same engine written again in WebAssembly, which is what runs the demo above.

If you'd rather run it on your own machine, Ollama has the 300M:

ollama run supergoatscriptguy/mnemonic

The weights for its own CPU engine are on Hugging Face, and the code is on GitHub.