Docs / Tutorials / Build a GPT in Motoko

Build a GPT in Motokoalpha

Write microGPT, a complete GPT with its own autograd, as one Motoko actor, train it on 32,000 names and sample new ones.

In this tutorial you build a GPT from nothing: a tokenizer, automatic differentiation, a transformer, the Adam optimizer, a training loop and sampling, with no libraries beyond Motoko's core. It follows Andrej Karpathy's microgpt.py, the whole algorithm in about 200 lines of Python, section by section. At the end you have an actor that learns from 32,000 names and invents new ones.

The model is tiny: one layer, 16-wide embeddings, 4,192 parameters. Every number is a scalar node in a computation graph, the way microgpt.py does it, so you can follow each multiplication from the input to the gradient.

Prerequisites:

moxzi and moxzid, and mops
curl and jq
The first actor tutorial, or familiarity with calling an actor over HTTP

The finished project:

mops.toml       one dependency: core
actors.json     tells moxzid to host gpt.wasm as "gpt"
names.txt       the dataset
gpt.mo          the model, the training loop and sampling

1. Set up the project#

Create a directory with a mops.toml that depends on Motoko's core library, and install it:

mops.toml

[dependencies]
core = "2.6.0"

mops install

Download the dataset, the list of names from Karpathy's makemore project, one name per line:

curl -LO https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt
wc -l names.txt

32032 names.txt

The last line has no newline, so the file holds 32,033 names.

Everything else goes in one file, gpt.mo, an actor. The steps below fill it in from top to bottom, in the same order as microgpt.py. The complete file is in step 9.

2. Randomness#

microgpt.py starts with random.seed(42). Training needs random numbers in three places: to shuffle the dataset, to initialize the parameters and to sample names. Write a seeded generator, xorshift64*, and derive a uniform draw and a normal draw from it:

gpt.mo

// Let there be order among chaos: a seeded random number generator (xorshift64*).
var rng : Nat64 = 42;

func next() : Nat64 {
  rng ^= rng >> 12;
  rng ^= rng << 25;
  rng ^= rng >> 27;
  rng *% 0x2545F4914F6CDD1D
};

// A uniform draw in [0, 1), and a normal draw: the sum of 12 uniforms is close to a bell curve.
func uniform() : Float = Float.fromInt(Nat64.toNat(next() >> 11)) * 0x1p-53;

func gauss(std : Float) : Float {
  var acc = 0.0;
  for (_ in Nat.range(0, 12)) acc += uniform();
  std * (acc - 6.0)
};

gauss adds 12 uniform draws: their sum has mean 6 and variance 1, close enough to a bell curve for initializing weights.

The generator is integer arithmetic. Float.exp and Float.log come from moxzi's runtime, not from your operating system's math library, and Float.sqrt is a single wasm instruction with an exactly specified result. So every run, on every machine, produces the same numbers: the losses and names on this page are the ones you get, to the last digit.

They differ from the ones microgpt.py prints. Python's random is a different generator, so it shuffles the names and draws the weights differently even with the same seed.

3. The dataset#

The dataset is a list of documents, one per line, shuffled once:

gpt.mo

// Let there be a dataset: one document per line (here, one name per line).
var docs : [Text] = [];

func shuffle(xs : [var Text]) {
  var i = xs.size();
  while (i > 1) {
    i -= 1;
    let j = Nat64.toNat(next() % Nat64.fromNat(i + 1));
    let x = xs[i];
    xs[i] := xs[j];
    xs[j] := x;
  };
};

docs is a variable of the actor. Actors keep their variables between calls, so you load the dataset once and every later call sees it.

4. The tokenizer#

A tokenizer turns text into a sequence of integers. Here each distinct character is a token, and one extra token, BOS (beginning of sequence), marks where a name starts and ends:

gpt.mo

// Let there be a tokenizer: each distinct character gets a token id, plus one special
// Beginning Of Sequence token (BOS) that marks where a name starts and ends.
var uchars : [Char] = [];
var BOS = 0;
var vocabSize = 0;

func buildVocabulary() {
  let seen = List.empty<Char>();
  for (doc in docs.values()) {
    for (c in doc.chars()) if (not seen.contains(c)) seen.add(c);
  };
  uchars := seen.toArray().sort();
  BOS := uchars.size();
  vocabSize := uchars.size() + 1;
};

func encode(doc : Text) : [Nat] {
  let cs = doc.toArray();
  Array.tabulate<Nat>(cs.size() + 2, func i {
    if (i == 0 or i == cs.size() + 1) return BOS;
    switch (uchars.indexOf(cs[i - 1])) {
      case (?id) id;
      case null Runtime.trap("character not in the vocabulary");
    }
  })
};

The names use the 26 letters a to z, so the vocabulary has 27 tokens and BOS is 26. encode("emma") returns [26, 4, 12, 12, 0, 26].

seen.contains(c), .sort() and uchars.indexOf(...) need a way to compare characters. You don't pass one: Motoko finds Char.equal and Char.compare in the imported Char module.

5. Autograd#

A neural network learns by gradient descent: compute a loss, then nudge every parameter in the direction that lowers it. The direction comes from the derivative of the loss with respect to each parameter, and autograd computes those derivatives by applying the chain rule backward through every operation of the forward pass.

microgpt.py does this with a Value class that overloads Python's operators:

def __mul__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    return Value(self.data * other.data, (self, other), (other.data, self.data))

In Motoko a node is a record. Its data and grad fields are var, because the backward pass and the optimizer update them:

gpt.mo

// Let there be autograd. A Node holds a number from the forward pass, the derivative of the
// loss with respect to it (filled in by the backward pass), the nodes it was computed from,
// and its local derivative with respect to each of them.
type Node = {
  var data : Float;
  var grad : Float;
  children : [Node];
  localGrads : [Float];
};

// Every node computed during a step, in the order it was created.
transient var tape = List.empty<Node>();

func leaf(data : Float) : Node = { var data; var grad = 0.0; children = []; localGrads = [] };

func node(data : Float, children : [Node], localGrads : [Float]) : Node {
  let n : Node = { var data; var grad = 0.0; children; localGrads };
  tape.add(n);
  n
};

Each operation computes its result and records its local derivatives: for a * b, the derivative with respect to a is b, and with respect to b it's a. Motoko doesn't overload operators, so a * b is written mul(a, b):

gpt.mo

func add(a : Node, b : Node) : Node = node(a.data + b.data, [a, b], [1.0, 1.0]);
func mul(a : Node, b : Node) : Node = node(a.data * b.data, [a, b], [b.data, a.data]);
func addF(a : Node, c : Float) : Node = node(a.data + c, [a], [1.0]);
func subF(a : Node, c : Float) : Node = node(a.data - c, [a], [1.0]);
func mulF(a : Node, c : Float) : Node = node(a.data * c, [a], [c]);
func neg(a : Node) : Node = mulF(a, -1.0);
func recip(a : Node) : Node = node(1.0 / a.data, [a], [-1.0 / (a.data * a.data)]);
func div(a : Node, b : Node) : Node = mul(a, recip(b));
func rsqrt(a : Node) : Node {
  let r = 1.0 / Float.sqrt(a.data);
  node(r, [a], [-0.5 * r / a.data])
};
func exp(a : Node) : Node { let e = Float.exp(a.data); node(e, [a], [e]) };
func log(a : Node) : Node = node(Float.log(a.data), [a], [1.0 / a.data]);
func relu(a : Node) : Node =
  node(if (a.data > 0.0) a.data else 0.0, [a], [if (a.data > 0.0) 1.0 else 0.0]);

func sum(xs : [Node]) : Node {
  var s = xs[0];
  for (i in Nat.range(1, xs.size())) s := add(s, xs[i]);
  s
};

Division is multiplication by a reciprocal, and rsqrt is 1/sqrt(x), both as microgpt.py computes them with **-1 and **-0.5.

The backward pass needs the nodes in topological order, so that a node's gradient is complete before it is passed on to its children. microgpt.py sorts them with a recursive walk and a visited set. In this version every new node goes on tape as it's created, and a node is always created after its children, so the tape is already in topological order:

gpt.mo

// A node is always created after the nodes it is computed from, so the tape is already in
// topological order: walk it backwards and apply the chain rule.
func backward(loss : Node) {
  loss.grad := 1.0;
  var i = tape.size();
  while (i > 0) {
    i -= 1;
    let v = tape.at(i);
    for (j in v.children.keys()) {
      v.children[j].grad += v.localGrads[j] * v.grad;
    };
  };
};

tape is declared transient. A training step builds between 22,000 and 110,000 nodes, depending on the length of the name, and throws them away when the step ends. transient means the variable doesn't need to survive an upgrade of the actor. Everything else you declare with var does, which matters in the next tutorial.

6. The parameters#

The parameters are matrices of leaf nodes, drawn from a normal distribution with standard deviation 0.08:

gpt.mo

// The parameters: the knowledge of the model.
let nLayer = 1; // depth: the number of transformer layers
let nEmbd = 16; // width: the embedding dimension
let blockSize = 16; // the longest context the model attends to (the longest name has 15 characters)
let nHead = 4; // attention heads per layer
let headDim = nEmbd / nHead;

type Matrix = [[Node]];
type Layer = { wq : Matrix; wk : Matrix; wv : Matrix; wo : Matrix; fc1 : Matrix; fc2 : Matrix };

func matrix(nout : Nat, nin : Nat) : Matrix =
  Array.tabulate<[Node]>(nout, func _ = Array.tabulate<Node>(nin, func _ = leaf(gauss(0.08))));

var wte : Matrix = []; // token embeddings
var wpe : Matrix = []; // position embeddings
var lmHead : Matrix = []; // from the final vector back to a score per token
var layers : [Layer] = [];
var params : [Node] = []; // every parameter in one list, for the optimizer

func initParams() {
  wte := matrix(vocabSize, nEmbd);
  wpe := matrix(blockSize, nEmbd);
  lmHead := matrix(vocabSize, nEmbd);
  layers := Array.tabulate<Layer>(nLayer, func _ {
    let wq = matrix(nEmbd, nEmbd);
    let wk = matrix(nEmbd, nEmbd);
    let wv = matrix(nEmbd, nEmbd);
    let wo = matrix(nEmbd, nEmbd);
    let fc1 = matrix(4 * nEmbd, nEmbd);
    let fc2 = matrix(nEmbd, 4 * nEmbd);
    { wq; wk; wv; wo; fc1; fc2 }
  });
  let all = List.empty<Node>();
  func addMatrix(m : Matrix) { for (row in m.values()) for (p in row.values()) all.add(p) };
  addMatrix(wte);
  addMatrix(wpe);
  addMatrix(lmHead);
  for (l in layers.values()) {
    for (m in [l.wq, l.wk, l.wv, l.wo, l.fc1, l.fc2].values()) addMatrix(m);
  };
  params := all.toArray();
};

Matrix	Shape	What it holds
`wte`	27 × 16	an embedding vector for each token
`wpe`	16 × 16	an embedding vector for each position
`wq`, `wk`, `wv`, `wo`	16 × 16	attention: queries, keys, values and the output projection
`fc1`, `fc2`	64 × 16, 16 × 64	the MLP
`lmHead`	27 × 16	turns the final vector into a score for each token

params lists all 4,192 parameters in one array for the optimizer.

7. The model#

Three building blocks: linear multiplies a matrix by a vector, softmax turns scores into probabilities, and rmsnorm scales a vector to unit root-mean-square:

gpt.mo

// The model's building blocks: a matrix-vector product, softmax, and RMS normalization.
func linear(x : [Node], w : Matrix) : [Node] =
  w.map(func(row : [Node]) : Node = sum(Array.tabulate<Node>(x.size(), func i = mul(row[i], x[i]))));

func softmax(logits : [Node]) : [Node] {
  var max = logits[0].data;
  for (l in logits.values()) if (l.data > max) max := l.data;
  let exps = logits.map(func(l : Node) : Node = exp(subF(l, max)));
  let total = sum(exps);
  exps.map(func(e : Node) : Node = div(e, total))
};

func rmsnorm(x : [Node]) : [Node] {
  let ms = mulF(sum(x.map(func(xi : Node) : Node = mul(xi, xi))), 1.0 / Float.fromInt(x.size()));
  let scale = rsqrt(addF(ms, 1e-5));
  x.map(func(xi : Node) : Node = mul(xi, scale))
};

softmax subtracts the largest score before exponentiating. That doesn't change the result and keeps exp from overflowing.

The model is GPT-2 with microgpt.py's simplifications: RMSNorm instead of LayerNorm, no biases, and ReLU instead of GeLU. It processes one token at a time and keeps each position's keys and values in a cache, so attention at position n can look at positions 0 to n:

gpt.mo

// The model: a token and its position in, a score for every possible next token out.
// Each layer keeps the keys and values of the positions it has seen (the KV cache).
type Cache = { keys : List.List<[Node]>; values : List.List<[Node]> };

func newCaches() : [Cache] =
  Array.tabulate<Cache>(nLayer, func _ = { keys = List.empty<[Node]>(); values = List.empty<[Node]>() });

func gpt(tokenId : Nat, posId : Nat, caches : [Cache]) : [Node] {
  let tokEmb = wte[tokenId];
  let posEmb = wpe[posId];
  var x = rmsnorm(Array.tabulate<Node>(nEmbd, func i = add(tokEmb[i], posEmb[i])));

  for (li in layers.keys()) {
    let layer = layers[li];
    let cache = caches[li];

    // 1) Multi-head attention
    let xResidual = x;
    x := rmsnorm(x);
    let q = linear(x, layer.wq);
    cache.keys.add(linear(x, layer.wk));
    cache.values.add(linear(x, layer.wv));
    let seen = cache.keys.size();
    let heads = Array.tabulate<[Node]>(nHead, func h {
      let hs = h * headDim;
      let attnLogits = Array.tabulate<Node>(seen, func t {
        let k = cache.keys.at(t);
        mulF(sum(Array.tabulate<Node>(headDim, func j = mul(q[hs + j], k[hs + j]))), 1.0 / Float.sqrt(Float.fromInt(headDim)))
      });
      let attnWeights = softmax(attnLogits);
      Array.tabulate<Node>(headDim, func j =
        sum(Array.tabulate<Node>(seen, func t = mul(attnWeights[t], cache.values.at(t)[hs + j]))))
    });
    let xAttn = Array.tabulate<Node>(nEmbd, func i = heads[i / headDim][i % headDim]);
    let attnOut = linear(xAttn, layer.wo);
    x := Array.tabulate<Node>(nEmbd, func i = add(attnOut[i], xResidual[i]));

    // 2) MLP
    let mlpResidual = x;
    let hidden = linear(rmsnorm(x), layer.fc1).map(relu);
    let mlpOut = linear(hidden, layer.fc2);
    x := Array.tabulate<Node>(nEmbd, func i = add(mlpOut[i], mlpResidual[i]));
  };

  linear(x, lmHead)
};

The layer has two blocks, each added back onto its input (a residual connection):

Attention. Each of the 4 heads compares this position's query with the keys of every position so far, turns the scores into weights with softmax, and takes the weighted sum of their values.
MLP. A 64-wide hidden layer with ReLU.

8. Training#

Adam keeps two running averages per parameter, of its gradient and of its gradient squared:

gpt.mo

// Let there be Adam, the optimizer, and its buffers.
let learningRate = 0.01;
let beta1 = 0.85;
let beta2 = 0.99;
let epsAdam = 1e-8;
var m : [var Float] = [var]; // first moment
var v : [var Float] = [var]; // second moment
var beta1Pow = 1.0; // beta1 ** (step + 1), kept as a running product
var beta2Pow = 1.0;

Loading the dataset shuffles it, builds the vocabulary and draws the parameters, in that order, as microgpt.py does:

gpt.mo

public func load(text : Text) : async Text {
  let lines = text.split(#char '\n').filter(func(line : Text) : Bool = line != "").toArray();
  let shuffled = lines.toVarArray();
  shuffle(shuffled);
  docs := shuffled.toArray();
  buildVocabulary();
  initParams();
  m := VarArray.repeat<Float>(0.0, params.size());
  v := VarArray.repeat<Float>(0.0, params.size());
  "num docs: " # docs.size().toText() # ", vocab size: " # vocabSize.toText() # ", num params: " # params.size().toText()
};

A training step takes one name, predicts each next character, and measures the loss: the negative log of the probability the model gave the right character, averaged over the name. Then it runs the backward pass and lets Adam move every parameter:

gpt.mo

let numSteps = 1000;
var step = 0;

func trainStep() : Float {
  tape := List.empty<Node>(); // a fresh graph for this step

  // Take one document, tokenize it, and surround it with BOS on both sides.
  let tokens = encode(docs[step % docs.size()]);
  let n = Nat.min(blockSize, tokens.size() - 1);

  // Forward the tokens through the model, building the graph all the way to the loss.
  let caches = newCaches();
  let losses = List.empty<Node>();
  for (pos in Nat.range(0, n)) {
    let probs = softmax(gpt(tokens[pos], pos, caches));
    losses.add(neg(log(probs[tokens[pos + 1]])));
  };
  let loss = mulF(sum(losses.toArray()), 1.0 / Float.fromInt(n)); // may yours be low

  // Backward: the gradient of the loss with respect to every parameter.
  backward(loss);

  // Adam update, with linear learning-rate decay.
  let lr = learningRate * (1.0 - Float.fromInt(step) / Float.fromInt(numSteps));
  beta1Pow *= beta1;
  beta2Pow *= beta2;
  for (i in params.keys()) {
    let p = params[i];
    m[i] := beta1 * m[i] + (1.0 - beta1) * p.grad;
    v[i] := beta2 * v[i] + (1.0 - beta2) * (p.grad * p.grad);
    let mHat = m[i] / (1.0 - beta1Pow);
    let vHat = v[i] / (1.0 - beta2Pow);
    p.data -= (lr * mHat) / (Float.sqrt(vHat) + epsAdam);
    p.grad := 0.0;
  };

  step += 1;
  loss.data
};

public func train(steps : Nat) : async Text {
  var loss = 0.0;
  for (_ in Nat.range(0, steps)) {
    if (step < numSteps) loss := trainStep();
    await async {}; // end this message: the step is saved, and its graph can be collected
  };
  "step " # step.toText() # " / " # numSteps.toText() # " | loss " # loss.format(#fix 4)
};

train awaits after every step. An await ends the current message, and the rest of the loop continues in a new one. That does three things:

Each step is saved when it finishes. A message's changes are committed when it ends, so if the server stops, the completed steps are kept.
The garbage collector runs between steps. It runs between messages. Without the await, a call of 100 steps keeps all 5 million nodes it creates in memory until it returns, and the actor grows to about 1.5 GB; with it, the actor stays under 400 MB.
No message runs too long. moxzid and the Internet Computer stop a message after 40 billion instructions. A step takes about 165 million on average, so 1,000 steps in one message would need four times the limit.

9. Load, train and sample#

The last section of gpt.mo samples names. It starts from BOS, picks each next character at random according to the model's probabilities, and stops at the next BOS. The temperature divides the scores first: below 1, likely characters get likelier.

gpt.mo

// Inference: may the model babble back to us.
func choose(weights : [Float]) : Nat {
  let u = uniform();
  var cumulative = 0.0;
  for (i in weights.keys()) {
    cumulative += weights[i];
    if (cumulative > u) return i;
  };
  weights.size() - 1
};

public func sample(count : Nat, temperature : Float) : async [Text] {
  Array.tabulate<Text>(count, func _ {
    tape := List.empty<Node>();
    let caches = newCaches();
    var tokenId = BOS;
    var name = "";
    label generate for (pos in Nat.range(0, blockSize)) {
      let logits = gpt(tokenId, pos, caches);
      let probs = softmax(logits.map(func(l : Node) : Node = mulF(l, 1.0 / temperature)));
      tokenId := choose(probs.map(func(p : Node) : Float = p.data));
      if (tokenId == BOS) break generate;
      name #= Text.fromChar(uchars[tokenId]);
    };
    name
  })
};

Here's the complete file:

gpt.mo

// microGPT in Motoko: train and run a GPT in one dependency-free actor.
// Follows Andrej Karpathy's microgpt.py (https://gist.github.com/karpathy/8627fe009c40f57531cb18360106ce95)
// section by section. SPDX-License-Identifier: Apache-2.0
import Array "mo:core/Array";
import Char "mo:core/Char";
import Float "mo:core/Float";
import Iter "mo:core/Iter";
import List "mo:core/List";
import Nat "mo:core/Nat";
import Nat64 "mo:core/Nat64";
import Runtime "mo:core/Runtime";
import Text "mo:core/Text";
import VarArray "mo:core/VarArray";

actor {
  // Let there be order among chaos: a seeded random number generator (xorshift64*).
  var rng : Nat64 = 42;

  func next() : Nat64 {
    rng ^= rng >> 12;
    rng ^= rng << 25;
    rng ^= rng >> 27;
    rng *% 0x2545F4914F6CDD1D
  };

  // A uniform draw in [0, 1), and a normal draw: the sum of 12 uniforms is close to a bell curve.
  func uniform() : Float = Float.fromInt(Nat64.toNat(next() >> 11)) * 0x1p-53;

  func gauss(std : Float) : Float {
    var acc = 0.0;
    for (_ in Nat.range(0, 12)) acc += uniform();
    std * (acc - 6.0)
  };

  // Let there be a dataset: one document per line (here, one name per line).
  var docs : [Text] = [];

  func shuffle(xs : [var Text]) {
    var i = xs.size();
    while (i > 1) {
      i -= 1;
      let j = Nat64.toNat(next() % Nat64.fromNat(i + 1));
      let x = xs[i];
      xs[i] := xs[j];
      xs[j] := x;
    };
  };

  // Let there be a tokenizer: each distinct character gets a token id, plus one special
  // Beginning Of Sequence token (BOS) that marks where a name starts and ends.
  var uchars : [Char] = [];
  var BOS = 0;
  var vocabSize = 0;

  func buildVocabulary() {
    let seen = List.empty<Char>();
    for (doc in docs.values()) {
      for (c in doc.chars()) if (not seen.contains(c)) seen.add(c);
    };
    uchars := seen.toArray().sort();
    BOS := uchars.size();
    vocabSize := uchars.size() + 1;
  };

  func encode(doc : Text) : [Nat] {
    let cs = doc.toArray();
    Array.tabulate<Nat>(cs.size() + 2, func i {
      if (i == 0 or i == cs.size() + 1) return BOS;
      switch (uchars.indexOf(cs[i - 1])) {
        case (?id) id;
        case null Runtime.trap("character not in the vocabulary");
      }
    })
  };

  // Let there be autograd. A Node holds a number from the forward pass, the derivative of the
  // loss with respect to it (filled in by the backward pass), the nodes it was computed from,
  // and its local derivative with respect to each of them.
  type Node = {
    var data : Float;
    var grad : Float;
    children : [Node];
    localGrads : [Float];
  };

  // Every node computed during a step, in the order it was created.
  transient var tape = List.empty<Node>();

  func leaf(data : Float) : Node = { var data; var grad = 0.0; children = []; localGrads = [] };

  func node(data : Float, children : [Node], localGrads : [Float]) : Node {
    let n : Node = { var data; var grad = 0.0; children; localGrads };
    tape.add(n);
    n
  };

  func add(a : Node, b : Node) : Node = node(a.data + b.data, [a, b], [1.0, 1.0]);
  func mul(a : Node, b : Node) : Node = node(a.data * b.data, [a, b], [b.data, a.data]);
  func addF(a : Node, c : Float) : Node = node(a.data + c, [a], [1.0]);
  func subF(a : Node, c : Float) : Node = node(a.data - c, [a], [1.0]);
  func mulF(a : Node, c : Float) : Node = node(a.data * c, [a], [c]);
  func neg(a : Node) : Node = mulF(a, -1.0);
  func recip(a : Node) : Node = node(1.0 / a.data, [a], [-1.0 / (a.data * a.data)]);
  func div(a : Node, b : Node) : Node = mul(a, recip(b));
  func rsqrt(a : Node) : Node {
    let r = 1.0 / Float.sqrt(a.data);
    node(r, [a], [-0.5 * r / a.data])
  };
  func exp(a : Node) : Node { let e = Float.exp(a.data); node(e, [a], [e]) };
  func log(a : Node) : Node = node(Float.log(a.data), [a], [1.0 / a.data]);
  func relu(a : Node) : Node =
    node(if (a.data > 0.0) a.data else 0.0, [a], [if (a.data > 0.0) 1.0 else 0.0]);

  func sum(xs : [Node]) : Node {
    var s = xs[0];
    for (i in Nat.range(1, xs.size())) s := add(s, xs[i]);
    s
  };

  // A node is always created after the nodes it is computed from, so the tape is already in
  // topological order: walk it backwards and apply the chain rule.
  func backward(loss : Node) {
    loss.grad := 1.0;
    var i = tape.size();
    while (i > 0) {
      i -= 1;
      let v = tape.at(i);
      for (j in v.children.keys()) {
        v.children[j].grad += v.localGrads[j] * v.grad;
      };
    };
  };

  // The parameters: the knowledge of the model.
  let nLayer = 1; // depth: the number of transformer layers
  let nEmbd = 16; // width: the embedding dimension
  let blockSize = 16; // the longest context the model attends to (the longest name has 15 characters)
  let nHead = 4; // attention heads per layer
  let headDim = nEmbd / nHead;

  type Matrix = [[Node]];
  type Layer = { wq : Matrix; wk : Matrix; wv : Matrix; wo : Matrix; fc1 : Matrix; fc2 : Matrix };

  func matrix(nout : Nat, nin : Nat) : Matrix =
    Array.tabulate<[Node]>(nout, func _ = Array.tabulate<Node>(nin, func _ = leaf(gauss(0.08))));

  var wte : Matrix = []; // token embeddings
  var wpe : Matrix = []; // position embeddings
  var lmHead : Matrix = []; // from the final vector back to a score per token
  var layers : [Layer] = [];
  var params : [Node] = []; // every parameter in one list, for the optimizer

  func initParams() {
    wte := matrix(vocabSize, nEmbd);
    wpe := matrix(blockSize, nEmbd);
    lmHead := matrix(vocabSize, nEmbd);
    layers := Array.tabulate<Layer>(nLayer, func _ {
      let wq = matrix(nEmbd, nEmbd);
      let wk = matrix(nEmbd, nEmbd);
      let wv = matrix(nEmbd, nEmbd);
      let wo = matrix(nEmbd, nEmbd);
      let fc1 = matrix(4 * nEmbd, nEmbd);
      let fc2 = matrix(nEmbd, 4 * nEmbd);
      { wq; wk; wv; wo; fc1; fc2 }
    });
    let all = List.empty<Node>();
    func addMatrix(m : Matrix) { for (row in m.values()) for (p in row.values()) all.add(p) };
    addMatrix(wte);
    addMatrix(wpe);
    addMatrix(lmHead);
    for (l in layers.values()) {
      for (m in [l.wq, l.wk, l.wv, l.wo, l.fc1, l.fc2].values()) addMatrix(m);
    };
    params := all.toArray();
  };

  // The model's building blocks: a matrix-vector product, softmax, and RMS normalization.
  func linear(x : [Node], w : Matrix) : [Node] =
    w.map(func(row : [Node]) : Node = sum(Array.tabulate<Node>(x.size(), func i = mul(row[i], x[i]))));

  func softmax(logits : [Node]) : [Node] {
    var max = logits[0].data;
    for (l in logits.values()) if (l.data > max) max := l.data;
    let exps = logits.map(func(l : Node) : Node = exp(subF(l, max)));
    let total = sum(exps);
    exps.map(func(e : Node) : Node = div(e, total))
  };

  func rmsnorm(x : [Node]) : [Node] {
    let ms = mulF(sum(x.map(func(xi : Node) : Node = mul(xi, xi))), 1.0 / Float.fromInt(x.size()));
    let scale = rsqrt(addF(ms, 1e-5));
    x.map(func(xi : Node) : Node = mul(xi, scale))
  };

  // The model: a token and its position in, a score for every possible next token out.
  // Each layer keeps the keys and values of the positions it has seen (the KV cache).
  type Cache = { keys : List.List<[Node]>; values : List.List<[Node]> };

  func newCaches() : [Cache] =
    Array.tabulate<Cache>(nLayer, func _ = { keys = List.empty<[Node]>(); values = List.empty<[Node]>() });

  func gpt(tokenId : Nat, posId : Nat, caches : [Cache]) : [Node] {
    let tokEmb = wte[tokenId];
    let posEmb = wpe[posId];
    var x = rmsnorm(Array.tabulate<Node>(nEmbd, func i = add(tokEmb[i], posEmb[i])));

    for (li in layers.keys()) {
      let layer = layers[li];
      let cache = caches[li];

      // 1) Multi-head attention
      let xResidual = x;
      x := rmsnorm(x);
      let q = linear(x, layer.wq);
      cache.keys.add(linear(x, layer.wk));
      cache.values.add(linear(x, layer.wv));
      let seen = cache.keys.size();
      let heads = Array.tabulate<[Node]>(nHead, func h {
        let hs = h * headDim;
        let attnLogits = Array.tabulate<Node>(seen, func t {
          let k = cache.keys.at(t);
          mulF(sum(Array.tabulate<Node>(headDim, func j = mul(q[hs + j], k[hs + j]))), 1.0 / Float.sqrt(Float.fromInt(headDim)))
        });
        let attnWeights = softmax(attnLogits);
        Array.tabulate<Node>(headDim, func j =
          sum(Array.tabulate<Node>(seen, func t = mul(attnWeights[t], cache.values.at(t)[hs + j]))))
      });
      let xAttn = Array.tabulate<Node>(nEmbd, func i = heads[i / headDim][i % headDim]);
      let attnOut = linear(xAttn, layer.wo);
      x := Array.tabulate<Node>(nEmbd, func i = add(attnOut[i], xResidual[i]));

      // 2) MLP
      let mlpResidual = x;
      let hidden = linear(rmsnorm(x), layer.fc1).map(relu);
      let mlpOut = linear(hidden, layer.fc2);
      x := Array.tabulate<Node>(nEmbd, func i = add(mlpOut[i], mlpResidual[i]));
    };

    linear(x, lmHead)
  };

  // Let there be Adam, the optimizer, and its buffers.
  let learningRate = 0.01;
  let beta1 = 0.85;
  let beta2 = 0.99;
  let epsAdam = 1e-8;
  var m : [var Float] = [var]; // first moment
  var v : [var Float] = [var]; // second moment
  var beta1Pow = 1.0; // beta1 ** (step + 1), kept as a running product
  var beta2Pow = 1.0;

  public func load(text : Text) : async Text {
    let lines = text.split(#char '\n').filter(func(line : Text) : Bool = line != "").toArray();
    let shuffled = lines.toVarArray();
    shuffle(shuffled);
    docs := shuffled.toArray();
    buildVocabulary();
    initParams();
    m := VarArray.repeat<Float>(0.0, params.size());
    v := VarArray.repeat<Float>(0.0, params.size());
    "num docs: " # docs.size().toText() # ", vocab size: " # vocabSize.toText() # ", num params: " # params.size().toText()
  };

  let numSteps = 1000;
  var step = 0;

  func trainStep() : Float {
    tape := List.empty<Node>(); // a fresh graph for this step

    // Take one document, tokenize it, and surround it with BOS on both sides.
    let tokens = encode(docs[step % docs.size()]);
    let n = Nat.min(blockSize, tokens.size() - 1);

    // Forward the tokens through the model, building the graph all the way to the loss.
    let caches = newCaches();
    let losses = List.empty<Node>();
    for (pos in Nat.range(0, n)) {
      let probs = softmax(gpt(tokens[pos], pos, caches));
      losses.add(neg(log(probs[tokens[pos + 1]])));
    };
    let loss = mulF(sum(losses.toArray()), 1.0 / Float.fromInt(n)); // may yours be low

    // Backward: the gradient of the loss with respect to every parameter.
    backward(loss);

    // Adam update, with linear learning-rate decay.
    let lr = learningRate * (1.0 - Float.fromInt(step) / Float.fromInt(numSteps));
    beta1Pow *= beta1;
    beta2Pow *= beta2;
    for (i in params.keys()) {
      let p = params[i];
      m[i] := beta1 * m[i] + (1.0 - beta1) * p.grad;
      v[i] := beta2 * v[i] + (1.0 - beta2) * (p.grad * p.grad);
      let mHat = m[i] / (1.0 - beta1Pow);
      let vHat = v[i] / (1.0 - beta2Pow);
      p.data -= (lr * mHat) / (Float.sqrt(vHat) + epsAdam);
      p.grad := 0.0;
    };

    step += 1;
    loss.data
  };

  public func train(steps : Nat) : async Text {
    var loss = 0.0;
    for (_ in Nat.range(0, steps)) {
      if (step < numSteps) loss := trainStep();
      await async {}; // end this message: the step is saved, and its graph can be collected
    };
    "step " # step.toText() # " / " # numSteps.toText() # " | loss " # loss.format(#fix 4)
  };

  // Inference: may the model babble back to us.
  func choose(weights : [Float]) : Nat {
    let u = uniform();
    var cumulative = 0.0;
    for (i in weights.keys()) {
      cumulative += weights[i];
      if (cumulative > u) return i;
    };
    weights.size() - 1
  };

  public func sample(count : Nat, temperature : Float) : async [Text] {
    Array.tabulate<Text>(count, func _ {
      tape := List.empty<Node>();
      let caches = newCaches();
      var tokenId = BOS;
      var name = "";
      label generate for (pos in Nat.range(0, blockSize)) {
        let logits = gpt(tokenId, pos, caches);
        let probs = softmax(logits.map(func(l : Node) : Node = mulF(l, 1.0 / temperature)));
        tokenId := choose(probs.map(func(p : Node) : Float = p.data));
        if (tokenId == BOS) break generate;
        name #= Text.fromChar(uchars[tokenId]);
      };
      name
    })
  };
};

Create actors.json, then build and serve:

actors.json

{
  "actors": [
    { "name": "gpt", "wasm": "gpt.wasm" }
  ]
}

moxzi build gpt.mo

gpt.mo -> gpt.wasm + gpt.did (368912 bytes, 20 file(s), 235 steps, 2.71s)
compiled by ~/.local/lib/moxzi/compiler.wasm (sha256 9294817c6e9069cbd0e712d0e9dde0ec1b6893d0e7b42d798fa3dd4c8cd772cb), moxzi 0.1.0-alpha.10

moxzid -m actors.json -s ./state

In another terminal, send the dataset. jq -Rs '[.]' wraps the whole file as one JSON string argument:

jq -Rs '[.]' names.txt | curl -s -w '\n' --data-binary @- http://127.0.0.1:7000/call-json/gpt/load

["num docs: 32033, vocab size: 27, num params: 4192"]

Train for 100 steps:

curl -s -w '\n' -d '[100]' http://127.0.0.1:7000/call-json/gpt/train

["step 100 / 1000 | loss 2.8319"]

An untrained model gives each of the 27 tokens about the same probability, a loss of about ln 27 = 3.3. After 100 steps it's 2.83. The actor keeps the model between calls, so the next call continues from step 100. Train the remaining 900 steps:

curl -s -w '\n' -d '[900]' http://127.0.0.1:7000/call-json/gpt/train

["step 1000 / 1000 | loss 2.2349"]

Sample 20 names at temperature 0.5:

curl -s -d '[20, 0.5]' http://127.0.0.1:7000/call-json/gpt/sample | jq -r '.[0][]'

ahyana
aran
ralen
amon
anni
amanga
radan
aanana
tanari
karian
erinai
enalin
alide
ein
rale
rian
kala
aorly
alidin
aman

Six of them (aran, amon, anni, rian, kala, aman) are in the dataset. The other 14 are new names: the model learned which letters follow which, how names start and how long they run.

All 1,000 steps take about 35 seconds on an Apple M1 Max. microgpt.py takes about 105 seconds on the same machine under CPython 3.11.

Limitations#

Limitation	Detail
Scalar autograd	Every multiplication is a separate node, so a step builds tens of thousands of them. That makes every step inspectable, and slow. Real models work on whole tensors: see GPU math for actors
Names of 15 characters or fewer	`blockSize` is 16 positions: `BOS` plus 15 characters. A longer name runs past `wpe` and traps
Characters outside the dataset	`encode` traps with `character not in the vocabulary`
One name per step	microgpt.py trains on one document per step too. To train on several, run the forward pass for each and average their losses before `backward`

Next steps#

Keep training through upgrades and crashes: train in the background, change the code mid-training, and kill the server.
Train a Shakespeare GPT on the GPU: the same model on whole tensors, about 190 times larger.
Choosing a runtime: where else the same actor can run, including the Internet Computer.