Keep training through upgrades and crashesalpha

Train microGPT in the background, change its code mid-training, kill the server, and end with the same model as an uninterrupted run.

In this tutorial you train the GPT from Build a GPT in Motoko in the background. While it trains, you add a method to the actor and upgrade it, twice, then kill the server with kill -9. Training carries on each time, and the finished model is the one the uninterrupted run produced: the same final loss and the same 20 names, to the last digit.

A Python training script that has to survive a restart saves checkpoints: it writes the parameters, both Adam buffers, the random generator's state and the step counter to a file, and reads them back on start. The actor in this tutorial has no checkpoint code. Its variables are its state, and they survive upgrades and crashes.

Prerequisites:

• The project from Build a GPT in Motoko

1. Train in the background#

Add a timer that trains 10 steps per message until it reaches a target step, and a query that reports progress. Put this after the train function in gpt.mo, and add import Timer "mo:core/Timer"; to the imports:

// Train in the background: a timer runs a few steps per message until `target`.
var target = 0;
var lastLoss = 0.0;
transient var scheduled = false;

func schedule<system>() {
  if (scheduled or step >= target) return;
  scheduled := true;
  ignore Timer.setTimer<system>(#seconds 0, func() : async () {
    scheduled := false;
    var k = 0;
    while (k < 10 and step < target) {
      lastLoss := trainStep();
      k += 1;
    };
    schedule<system>();
  });
};

public func trainInBackground(until : Nat) : async () {
  target := Nat.min(until, numSteps);
  schedule<system>();
};

public query func status() : async Text {
  "step " # step.toText() # " / " # numSteps.toText() # " | loss " # lastLoss.format(#fix 4)
};

// The model, the optimizer and `target` survive an upgrade. Timers don't, so the actor's body,
// which runs again after every upgrade, re-arms the loop.
schedule<system>();

Each timer callback is its own message, so each batch of 10 steps is saved when the callback returns. Then the callback schedules the next one.

The last line, schedule<system>();, sits in the actor's body. The body runs when the actor is installed and again after every upgrade. The model, the optimizer and target survive an upgrade, but timers don't, so the body re-arms the loop. On a fresh install target is 0 and schedule returns without doing anything.

scheduled is transient for the same reason: it describes a timer, so it starts again at false after an upgrade, when there's no timer.

2. Start training#

Build and serve the actor:

moxzi build gpt.mo

gpt.mo -> gpt.wasm + gpt.did (377208 bytes, 22 file(s), 242 steps, 2.95s)
compiled by ~/.local/lib/moxzi/compiler.wasm (sha256 9294817c6e9069cbd0e712d0e9dde0ec1b6893d0e7b42d798fa3dd4c8cd772cb), moxzi 0.1.0-alpha.10

moxzid -m actors.json -s ./state

Load the names and train to step 300 in the background:

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

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

The call returns at once. Check on it until it reaches step 300 (about 10 seconds):

curl -s -w '\n' -d '[]' http://127.0.0.1:7000/query-json/gpt/status

["step 300 / 1000 | loss 2.3417"]

3. Change the code#

Add a method: score measures how surprising a name is to the model, as its average loss per character. Lower means more name-like:

// How surprising is a name to the model? Its average loss per character: lower is more name-like.
public query func score(name : Text) : async Text {
  tape := List.empty<Node>();
  let tokens = encode(name);
  let n = Nat.sub(tokens.size(), 1);
  let caches = newCaches();
  var total = 0.0;
  for (pos in Nat.range(0, n)) {
    let probs = softmax(gpt(tokens[pos], pos, caches));
    total -= Float.log(probs[tokens[pos + 1]].data);
  };
  name # ": " # (total / Float.fromInt(n)).format(#fix 2)
};

score is a query. It reads the model, and whatever it changes, such as the nodes it adds to tape, is discarded when it returns.

Build and upgrade the running actor:

moxzi build gpt.mo
curl -s -w '\n' --data-binary @gpt.wasm http://127.0.0.1:7000/upgrade/gpt

gpt.mo -> gpt.wasm + gpt.did (379022 bytes, 22 file(s), 242 steps, 3.31s)
compiled by ~/.local/lib/moxzi/compiler.wasm (sha256 9294817c6e9069cbd0e712d0e9dde0ec1b6893d0e7b42d798fa3dd4c8cd772cb), moxzi 0.1.0-alpha.10
upgraded

The actor runs the new code with the old state. It's still at step 300, with the same loss, and it can score names:

curl -s -w '\n' -d '[]' http://127.0.0.1:7000/query-json/gpt/status
curl -s -w '\n' -d '["emma"]' http://127.0.0.1:7000/query-json/gpt/score
curl -s -w '\n' -d '["xqzzk"]' http://127.0.0.1:7000/query-json/gpt/score

["step 300 / 1000 | loss 2.3417"]
["emma: 3.36"]
["xqzzk: 5.18"]

Everything declared with var came through the upgrade: the dataset, the vocabulary, the 4,192 parameters, both Adam buffers, the random generator, step and target. You didn't write code to save or restore any of it.

4. Upgrade while it trains#

Resume training to step 1,000:

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

[]

While it runs, upgrade again:

curl -s -w '\n' --data-binary @gpt.wasm http://127.0.0.1:7000/upgrade/gpt

upgraded

Call status a few times: the step keeps climbing. The upgrade cancelled the timer, and the actor's body scheduled a new one.

5. Crash the server#

Before it reaches step 1,000, kill moxzid and start it again:

pkill -9 moxzid
moxzid -m actors.json -s ./state

moxzid records every message in ./state before running it, and saves a snapshot of the actor each time a timer callback finishes. When it starts, it restores the actor and resumes its timers. The batch of steps that was running when you killed it never finished, so the restored actor contains none of it, and that timer callback runs again from the start.

6. Check the result#

When training finishes, look at the loss and sample 20 names at temperature 0.5:

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

["step 1000 / 1000 | loss 2.2349"]
ahyana
aran
ralen
amon
anni
amanga
radan
aanana
tanari
karian
erinai
enalin
alide
ein
rale
rian
kala
aorly
alidin
aman

These are the final loss and the 20 names from Build a GPT in Motoko, where training ran without interruption. Two upgrades and a crash changed when the steps ran, never what they computed.

What survives what#

An upgrade is refused if the new code declares a kept variable with an incompatible type, for example step changed from Nat to Text. To change a variable's type, write a migration.

Limitations#

Next steps#

• Persistence: where moxzid keeps state, and how to back it up.
• Migrations: change the types of an actor's variables across an upgrade.
• Train a Shakespeare GPT on the GPU: a larger model on a tensor device.