Train a Shakespeare GPT on the GPUalpha
Move from scalars to tensors, train nanoGPT's character-level Shakespeare model from a Motoko actor on moxzid's tensor device, and generate text.
In Build a GPT in Motoko every number was a node, and a step of a 4,192-parameter model built tens of thousands of them. Real models work on whole tensors, on a GPU. In this tutorial an actor trains an 800,000-parameter GPT on 1.1 MB of Shakespeare's plays on moxzid's tensor device, then writes new Shakespeare.
It follows the character-level Shakespeare model of Andrej Karpathy's nanoGPT, with nanoGPT's settings for a laptop: 4 layers, 4 attention heads, 128-wide embeddings, 64 characters of context, 12 sequences per step, 2,000 steps. The actor builds the model and decides what to train on; the device holds the weights and does the arithmetic.
Prerequisites:
- moxzi and moxzid, mops,
curland jq - Build a GPT in Motoko, for the model and the training loop
- GPU math for actors, skimmed, for how an actor calls the device
1. Set up the project#
The project has three files: mops.toml, actors.json and gpt.mo.
[dependencies]
core = "2.6.0"{
"actors": [
{ "name": "gpt", "wasm": "gpt.wasm" }
]
}mops installDownload tiny Shakespeare, the dataset nanoGPT uses:
curl -sLO https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
wc -c input.txt1115394 input.txt2. Describe the device#
The device is the management canister, aaaaa-aa, with methods named moxzi_device_*. A program sent to moxzi_device_run is a list of operations. Declare the ones you use:
// The tensor device is the management canister, aaaaa-aa. Declare the operations you use.
type Op = {
#input : Nat32;
#param : Text;
#narrow : { x : Nat32; dim : Nat32; start : Nat32; len : Nat32 };
#reshape : { x : Nat32; shape : [Nat32] };
#transpose : { x : Nat32; a : Nat32; b : Nat32 };
#embedding : { table : Nat32; ids : Nat32 };
#linear : { x : Nat32; w : Nat32 };
#rms_norm : Nat32;
#rope : Nat32;
#causal_attention : { q : Nat32; k : Nat32; v : Nat32 };
#relu : Nat32;
#add : { a : Nat32; b : Nat32 };
#cross_entropy : { logits : Nat32; targets : Nat32 };
};
type Device = actor {
moxzi_device_params : ([{ name : Text; shape : [Nat32]; init : { #values : Blob } }]) -> async ();
moxzi_device_optimizer : ({
name : Text; params : [Text]; lr : Float;
beta1 : Float; beta2 : Float; eps : Float; weight_decay : Float; kind : ?{ #adamw };
}) -> async ();
moxzi_device_run : ({
inputs : [{ data : Blob; elem : { #u8 }; shape : [Nat32] }];
program : [Op];
train : ?{ loss : Nat32; optimizer : Text; clip : ?Float; lr : ?Float; also : ?[Text] };
scalars : [Nat32];
tensors : [Nat32];
}) -> async { scalars : [Float]; tensors : [Blob] };
};
transient let device : Device = actor "aaaaa-aa";The settings come straight from nanoGPT's config/train_shakespeare_char.py, with its laptop overrides:
// nanoGPT's shakespeare_char settings for a laptop.
let nLayer = 4;
let nHead = 4;
let nEmbd = 128;
let blockSize = 64; // characters of context
let batchSize = 12; // sequences per step
let maxIters = 2000;
let learningRate = 1e-3;
let minLr = 1e-4;
let warmupIters = 100;3. The data#
The tokenizer is microGPT's, with bytes instead of characters: each distinct byte of the text gets a token id. Tiny Shakespeare has 65 of them. The actor keeps the text as a Blob, 1 byte per character, and looks each byte's id up in a 256-entry table:
// The text's bytes, and a token id for each distinct byte. The first 90% trains the model;
// the rest checks it on text it hasn't seen.
var data : Blob = "";
var vocab : [Nat8] = []; // token id -> byte
var idOf : [Nat8] = []; // byte -> token id
var nTrain = 0;
public func load(text : Text) : async Text {
data := text.encodeUtf8();
let seen = VarArray.repeat<Bool>(false, 256);
for (b in data.values()) seen[b.toNat()] := true;
vocab := Array.tabulate<Nat>(256, func b = b).filter(func(b : Nat) : Bool = seen[b]).map(Nat8.fromNat);
if (vocab[Nat.sub(vocab.size(), 1)] >= 128) Runtime.trap("this tokenizer handles ASCII text only");
let ids = VarArray.repeat<Nat8>(0, 256);
for (i in vocab.keys()) ids[vocab[i].toNat()] := Nat8.fromNat(i);
idOf := ids.toArray();
nTrain := data.size() * 9 / 10;
await initModel();
"characters: " # data.size().toText() # ", vocabulary: " # vocab.size().toText()
};
func token(i : Nat) : Nat8 = idOf[data[i].toNat()];The first 90% of the text trains the model. The last 10% is held out to measure how well it predicts text it hasn't seen, as nanoGPT does.
4. Starting weights#
The random generator is microGPT's:
// The same seeded generator as microGPT: it draws the starting weights, picks training
// windows and samples text, so every run on the CPU gives the same numbers.
var rng : Nat64 = 1337;
func next() : Nat64 {
rng ^= rng >> 12;
rng ^= rng << 25;
rng ^= rng >> 27;
rng *% 0x2545F4914F6CDD1D
};
func uniform() : Float = Float.fromInt(Nat64.toNat(next() >> 11)) * 0x1p-53;
func below(n : Nat) : Nat = Nat64.toNat(next() % Nat64.fromNat(n));
func gauss(std : Float) : Float {
var acc = 0.0;
for (_ in Nat.range(0, 12)) acc += uniform();
std * (acc - 6.0)
};The actor draws the starting weights itself and sends them to the device. The device can draw them too (init = #normal std), but from a generator you can't seed, so every run would start from a different model. Drawn here, a run on the CPU device gives the same numbers every time.
Tensors travel as float32 bytes, and Motoko has no bit cast between a Float and its bits, so two small functions convert by hand:
// The device exchanges tensors as little-endian float32 bytes. Motoko has no bit cast
// between Float and its bits, so convert by hand.
func f32Bits(x : Float) : Nat32 {
if (x == 0.0) return 0;
let a = Float.abs(x);
var e = Float.toInt(Float.floor(Float.log(a) / Float.log(2.0)));
var p = Float.pow(2.0, Float.fromInt(e));
if (p > a) { e -= 1; p /= 2.0 } else if (p * 2.0 <= a) { e += 1; p *= 2.0 };
let m = Float.toInt(Float.nearest((a / p - 1.0) * 0x1p23));
let bits = Nat32.fromIntWrap((e + 127) * 0x800000 + m);
if (x < 0.0) bits | 0x80000000 else bits
};
func f32(b : Blob, at : Nat) : Float {
let bits = b[at].toNat32() | (b[at + 1].toNat32() << 8) | (b[at + 2].toNat32() << 16) | (b[at + 3].toNat32() << 24);
let sign = if (bits >> 31 == 1) -1.0 else 1.0;
let e = (bits >> 23) & 0xFF;
let m = Float.fromInt((bits & 0x7FFFFF).toNat());
if (e == 0) sign * m * 0x1p-149 else sign * (1.0 + m * 0x1p-23) * Float.pow(2.0, Float.fromInt(e.toNat()) - 127.0)
};Each parameter goes to the device in its own message:
// Parameters live on the device: the actor only knows their names.
func layer(i : Nat, w : Text) : Text = "h" # i.toText() # "." # w;
func n32(n : Nat) : Nat32 = Nat32.fromNat(n);
var paramNames : [Text] = [];
// Draw `rows` x `cols` weights and send them to the device, one parameter per message.
func param(name : Text, rows : Nat, cols : Nat, std : Float) : async () {
let bytes = VarArray.repeat<Nat8>(0, 4 * rows * cols);
for (i in Nat.range(0, rows * cols)) {
let bits = f32Bits(gauss(std));
bytes[4 * i] := Nat8.fromNat((bits & 0xFF).toNat());
bytes[4 * i + 1] := Nat8.fromNat(((bits >> 8) & 0xFF).toNat());
bytes[4 * i + 2] := Nat8.fromNat(((bits >> 16) & 0xFF).toNat());
bytes[4 * i + 3] := Nat8.fromNat((bits >> 24).toNat());
};
await device.moxzi_device_params([{ name; shape = [n32(rows), n32(cols)]; init = #values(Blob.fromVarArray(bytes)) }]);
};
func initModel() : async () {
let (v, d) = (vocab.size(), nEmbd);
let std = 0.02;
let projStd = std / Float.sqrt(2.0 * Float.fromInt(nLayer)); // nanoGPT scales the residual projections
let names = List.empty<Text>();
func add(name : Text) : Text { names.add(name); name };
await param(add "wte", v, d, std);
for (i in Nat.range(0, nLayer)) {
await param(add(layer(i, "qkv")), 3 * d, d, std);
await param(add(layer(i, "proj")), d, d, projStd);
await param(add(layer(i, "fc")), 4 * d, d, std);
await param(add(layer(i, "out")), d, 4 * d, projStd);
};
await param(add "head", v, d, std);
paramNames := names.toArray();
await device.moxzi_device_optimizer({
name = "adamw"; params = paramNames; lr = learningRate;
beta1 = 0.9; beta2 = 0.99; eps = 1e-8; weight_decay = 0.1; kind = ?#adamw;
});
};await param(...) ends the message, and the garbage collector runs between messages. Drawing all 800,000 weights in one message would keep every intermediate value in memory until the message ends; one parameter at a time, the actor stays small. A small actor matters on moxzid: it saves each actor's memory after every message, so a bloated actor makes every training step slower.
The weights live on the device from now on. The actor refers to them by name: wte, h0.qkv, h0.proj and so on.
5. The model#
A program is a list of operations, and each operation's result is a register: its position in the list. gpt appends the model's operations and returns the register of the logits:
// A program is a list of operations. Each result is a register: its position in the list.
// gpt adds the model's operations: token ids [b, t] in, logits [b, t, vocab] out.
func gpt(ops : List.List<Op>, tokens : Nat32, b : Nat, t : Nat) : Nat32 {
func op(o : Op) : Nat32 { ops.add(o); n32(Nat.sub(ops.size(), 1)) };
let (bb, tt, d, hd) = (n32(b), n32(t), n32(nEmbd), n32(nEmbd / nHead));
// [b, t, d] -> [b, heads, t, head size]
func heads(x : Nat32) : Nat32 = op(#transpose { x = op(#reshape { x; shape = [bb, tt, n32(nHead), hd] }); a = 1; b = 2 });
var x = op(#embedding { table = op(#param "wte"); ids = tokens });
for (i in Nat.range(0, nLayer)) {
// 1) Causal self-attention, with rotary position embeddings on the queries and keys
let qkv = op(#linear { x = op(#rms_norm x); w = op(#param(layer(i, "qkv"))) });
let q = op(#rope(heads(op(#narrow { x = qkv; dim = 2; start = 0; len = d }))));
let k = op(#rope(heads(op(#narrow { x = qkv; dim = 2; start = d; len = d }))));
let v = heads(op(#narrow { x = qkv; dim = 2; start = 2 * d; len = d }));
let att = op(#causal_attention { q; k; v });
let merged = op(#reshape { x = op(#transpose { x = att; a = 1; b = 2 }); shape = [bb, tt, d] });
x := op(#add { a = x; b = op(#linear { x = merged; w = op(#param(layer(i, "proj"))) }) });
// 2) MLP
let hidden = op(#relu(op(#linear { x = op(#rms_norm x); w = op(#param(layer(i, "fc"))) })));
x := op(#add { a = x; b = op(#linear { x = hidden; w = op(#param(layer(i, "out"))) }) });
};
op(#linear { x = op(#rms_norm x); w = op(#param "head") })
};
// The loss on a batch: input 0 is [batch, block + 1] token ids.
func lossProgram() : ([Op], Nat32) {
let ops = List.empty<Op>();
func op(o : Op) : Nat32 { ops.add(o); n32(Nat.sub(ops.size(), 1)) };
let (b, t) = (n32(batchSize), n32(blockSize));
let seq = op(#input 0);
let inputs = op(#narrow { x = seq; dim = 1; start = 0; len = t });
let targets = op(#reshape { x = op(#narrow { x = seq; dim = 1; start = 1; len = t }); shape = [b * t] });
let logits = gpt(ops, inputs, batchSize, blockSize);
let flat = op(#reshape { x = logits; shape = [b * t, n32(vocab.size())] });
let loss = op(#cross_entropy { logits = flat; targets });
(ops.toArray(), loss)
};Compare it with microGPT's gpt function. The structure is the same: embed, then per layer attention and an MLP, each added onto its input, then a final projection to a score per token. Two things changed:
- Whole tensors.
#linearmultiplies every position of every sequence in the batch at once.#causal_attentioncomputes all heads at every position, each looking only at the positions before it. - Positions are rotations.
#roperotates each query and key by an angle that depends on its position, so attention sees how far apart two characters are. nanoGPT instead learns a vector per position, as microGPT did.
The model follows nanoGPT's GPT with three of Llama's choices, the ones the device provides: RMSNorm instead of LayerNorm, rotary positions instead of learned ones, and ReLU instead of GeLU. It has no dropout; nanoGPT's laptop settings turn dropout off too.
6. Training#
A training step sends 12 random windows of 65 characters and the program, and asks the device to train on the loss: compute gradients, clip them to a norm of 1, and take an AdamW step at this step's learning rate. Only the loss comes back.
var step = 0;
// Learning rate: linear warmup, then cosine decay to minLr (nanoGPT's get_lr).
func lrAt(it : Nat) : Float {
if (it < warmupIters) return learningRate * Float.fromInt(it + 1) / Float.fromInt(warmupIters + 1);
let ratio = Float.fromInt(Nat.sub(it, warmupIters)) / Float.fromInt(Nat.sub(maxIters, warmupIters));
minLr + 0.5 * (1.0 + Float.cos(Float.pi * ratio)) * (learningRate - minLr)
};
// Windows of block + 1 tokens starting at the given offsets: [batch, block + 1].
func batch(starts : [Nat]) : Blob =
Blob.fromArray(Array.tabulate<Nat8>(starts.size() * (blockSize + 1), func i =
token(starts[i / (blockSize + 1)] + i % (blockSize + 1))));
func run(starts : [Nat], train : Bool) : async Float {
let (program, loss) = lossProgram();
let r = await device.moxzi_device_run({
inputs = [{ data = batch(starts); elem = #u8; shape = [n32(batchSize), n32(blockSize + 1)] }];
program;
train = if train ?{ loss; optimizer = "adamw"; clip = ?1.0; lr = ?lrAt(step); also = null } else null;
scalars = [loss];
tensors = [];
});
r.scalars[0]
};
public func train(steps : Nat) : async Text {
var loss = 0.0;
for (_ in Nat.range(0, steps)) {
if (step < maxIters) {
let starts = Array.tabulate<Nat>(batchSize, func _ = below(Nat.sub(nTrain, blockSize + 1)));
loss := await run(starts, true);
step += 1;
};
};
"step " # step.toText() # " | loss " # loss.format(#fix 4)
};The learning rate warms up over 100 steps, then follows a cosine down to a tenth of its peak by step 2,000, as nanoGPT's does.
To measure progress, evaluate averages the loss over 20 fixed batches from each split, without training:
// Average loss over 20 fixed batches from each split, the numbers nanoGPT reports.
func estimate(from : Nat, to : Nat) : async Float {
let span = Nat.sub(to, from + blockSize + 1);
var total = 0.0;
for (k in Nat.range(0, 20)) {
let starts = Array.tabulate<Nat>(batchSize, func j = from + (k * batchSize + j) * span / (20 * batchSize));
total += await run(starts, false);
};
total / 20.0
};
public func evaluate() : async Text {
let trainLoss = await estimate(0, nTrain);
let valLoss = await estimate(nTrain, data.size());
"step " # step.toText() # ": train loss " # trainLoss.format(#fix 4) # ", val loss " # valLoss.format(#fix 4)
};7. Train it#
Build the actor and start moxzid with a device:
moxzi build gpt.mogpt.mo -> gpt.wasm + gpt.did (381880 bytes, 22 file(s), 235 steps, 2.84s)
compiled by ~/.local/lib/moxzi/compiler.wasm (sha256 9294817c6e9069cbd0e712d0e9dde0ec1b6893d0e7b42d798fa3dd4c8cd772cb), moxzi 0.1.0-alpha.10moxzid -m actors.json -s ./state --device burn:cpuThis page uses burn:cpu so the numbers match yours exactly. On a Mac, use burn:metal; on an NVIDIA GPU, burn:cuda. GPU results differ from these in the last digits.
Send the text. load builds the vocabulary and sends the starting weights:
jq -Rs '[.]' input.txt | curl -s -w '\n' --data-binary @- http://127.0.0.1:7000/call-json/gpt/load["characters: 1115394, vocabulary: 65"]Train 1,000 steps and evaluate:
curl -s -w '\n' -d '[1000]' http://127.0.0.1:7000/call-json/gpt/train
curl -s -w '\n' -d '[]' http://127.0.0.1:7000/call-json/gpt/evaluate["step 1000 | loss 1.7829"]
["step 1000: train loss 1.7230, val loss 1.8564"]And the remaining 1,000:
curl -s -w '\n' -d '[1000]' http://127.0.0.1:7000/call-json/gpt/train
curl -s -w '\n' -d '[]' http://127.0.0.1:7000/call-json/gpt/evaluate["step 2000 | loss 1.4615"]
["step 2000: train loss 1.5442, val loss 1.7135"]The validation loss, on text the model never trained on, falls from 4.17 (a uniform guess among 65 characters) to 1.71. nanoGPT reports 1.88 for the same settings; the rotary positions likely account for most of the difference. The training loss, 1.54, is lower than the validation loss: the model is starting to memorize its training text.
On an Apple M1 Max the 2,000 steps take about 4 minutes on burn:cpu and about 2 minutes on burn:metal.
8. Generate#
Generation feeds the model the last 64 characters and samples the next one from its output, 500 times. The device returns the logits as float32 bytes; the actor decodes them, applies the temperature and draws:
// Generate `count` characters after a newline, as nanoGPT's sample.py does.
public func generate(count : Nat, temperature : Float) : async Text {
let out = List.empty<Nat8>(); // token ids
out.add(idOf[10]); // "\n"
for (_ in Nat.range(0, count)) {
// The model sees the last blockSize characters.
let t = Nat.min(out.size(), blockSize);
let context = Array.tabulate<Nat8>(t, func i = out.at(Nat.sub(out.size(), t) + i));
let ops = List.empty<Op>();
func op(o : Op) : Nat32 { ops.add(o); n32(Nat.sub(ops.size(), 1)) };
let logits = gpt(ops, op(#input 0), 1, t);
let last = op(#narrow { x = logits; dim = 1; start = n32(Nat.sub(t, 1)); len = 1 });
let r = await device.moxzi_device_run({
inputs = [{ data = Blob.fromArray(context); elem = #u8; shape = [1, n32(t)] }];
program = ops.toArray();
train = null;
scalars = [];
tensors = [last];
});
// Softmax at the temperature, then draw a token.
let scores = Array.tabulate<Float>(vocab.size(), func v = f32(r.tensors[0], 4 * v) / temperature);
var max = scores[0];
for (s in scores.values()) if (s > max) max := s;
var total = 0.0;
for (s in scores.values()) total += Float.exp(s - max);
var u = uniform() * total;
var pick = Nat.sub(vocab.size(), 1);
label draw for (v in scores.keys()) {
let p = Float.exp(scores[v] - max);
if (u < p) { pick := v; break draw };
u -= p;
};
out.add(Nat8.fromNat(pick));
};
let bytes = Array.tabulate<Nat8>(Nat.sub(out.size(), 1), func i = vocab[out.at(i + 1).toNat()]);
switch (Text.decodeUtf8(Blob.fromArray(bytes))) { case (?text) text; case null "" }
};curl -s -d '[500, 0.8]' http://127.0.0.1:7000/call-json/gpt/generate | jq -r '.[0]'From hear thou say''t for unspaker by knower.
Which hid to resperion, aumberlant a head;
Me not a treathout of our not fair?
KING HENRY VI:
Nobly man.
Nurse:
I am that doth
Speak at let thou thank death.
Do this and know all
With desbyour negler, for the prince on turn,
And as I shall defit, the death great retoly be?
AUTOLYCUS:
May, though, he have man be to it is with your honour,
And to this now thy solding and that gave to before enmeo,
And him the grace of best mights to me.
MENENIUS:
AIt has learned the shape of a play: speakers' names in capitals, a colon, lines of verse, and mostly real English words, from 2,000 steps on 1 MB of text.
Limitations#
| Limitation | What to do instead |
|---|---|
| ASCII text only | The tokenizer gives each byte an id. For other text, use characters, as microGPT does, or a subword tokenizer |
| Generation recomputes the whole context for every character | Keep a key/value cache on the device: the rope_from and attention operations, with store and load |
| GPU results vary in the last digits | Use burn:cpu when you need exact repeatability |
| One actor, one device | See the next tutorial: several actors can train and read the same model |
Next steps#
- Train, serve and survive with actors: serve text from the model while it trains, kill the server mid-run, and train with several actors at once.
- GPU math for actors: every device method, sharing weights between machines, and paying for device time.