Serve and train with several actorsalpha

Let one actor write Shakespeare from another's weights while it trains, then train with four actors that average their weights.

In Train a Shakespeare GPT on the GPU one actor trained and used its model. In this tutorial the work is split between actors that share the device:

• A trainer trains in the background and publishes its weights as it goes.
• A writer generates text with the trainer's latest weights, while training continues. It has no weights of its own.
• Four islands each train on their own random batches and average their weights every 10 steps.

The weights never pass through a message. Actors publish them on the device and read each other's publications by name.

Prerequisites:

• The project from Train a Shakespeare GPT on the GPU

The finished project has two actors, trainer.mo and writer.mo, and a manifest that hosts the trainer four more times as islands:

{
  "actors": [
    { "name": "trainer", "wasm": "trainer.wasm" },
    { "name": "writer", "wasm": "writer.wasm" },
    { "name": "island0", "wasm": "trainer.wasm" },
    { "name": "island1", "wasm": "trainer.wasm" },
    { "name": "island2", "wasm": "trainer.wasm" },
    { "name": "island3", "wasm": "trainer.wasm" }
  ]
}

1. Train in the background and publish#

Start from the previous tutorial's gpt.mo, renamed trainer.mo. Name the actor Self (actor Self { ... }) so it can report its own principal, and add publish and mix to the device type:

// 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] };
  moxzi_device_publish : ({ names : [Text]; prefix : Text; readers : ?{ #anyone } }) -> async ();
  moxzi_device_mix : ([{ into : Text; terms : [{ owner : ?Principal; name : Text; weight : Float }] }]) -> async ();
};

transient let device : Device = actor "aaaaa-aa";

Then add background training. A timer runs 10 steps per callback, the way Keep training through upgrades and crashes did, and after each batch publishes every parameter under the prefix live/, readable by anyone:

// Train in the background, 10 steps per timer callback, and publish the weights after each
// batch so other actors can read them.
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) {
      let starts = Array.tabulate<Nat>(batchSize, func _ = below(Nat.sub(nTrain, blockSize + 1)));
      lastLoss := await run(starts, true);
      step += 1;
      k += 1;
    };
    await device.moxzi_device_publish({ names = paramNames; prefix = "live/"; readers = ?#anyone });
    schedule<system>();
  });
};

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

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

public query func vocabulary() : async Blob { Blob.fromArray(vocab) };

public query func id() : async Principal { Principal.fromActor(Self) };

// After an upgrade the timer is gone: the actor's body runs again and re-arms it.
schedule<system>();

Publishing copies the parameters on the device. The copy stays as it was while training changes the parameters, until the next batch publishes again.

2. Write from the trainer's weights#

The writer has the same model, with one change: every weight is a shared_ operation, which reads another actor's publication, instead of param. In Motoko the operation is spelled shared_ because shared is a keyword; the device sees shared:

// The same operations as the trainer's, plus `shared`: a weight another actor published.
type Op = {
  #input : Nat32;
  #shared_ : { owner : Principal; name : 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 };
};

type Device = actor {
  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";

follow remembers which trainer to read and fetches its vocabulary:

type Trainer = actor { vocabulary : query () -> async Blob };

var trainer : ?Principal = null;
var vocab : [Nat8] = [];

public func follow(t : Principal) : async () {
  trainer := ?t;
  vocab := (await (actor (t.toText()) : Trainer).vocabulary()).toArray();
};

// The trainer's model, with every weight read from its publication: `live/` + the name.
func gpt(ops : List.List<Op>, owner : Principal, tokens : Nat32, t : Nat) : Nat32 {
  func op(o : Op) : Nat32 { ops.add(o); n32(Nat.sub(ops.size(), 1)) };
  func w(name : Text) : Nat32 = op(#shared_ { owner; name = "live/" # name });
  let (tt, d, hd) = (n32(t), n32(nEmbd), n32(nEmbd / nHead));
  func heads(x : Nat32) : Nat32 = op(#transpose { x = op(#reshape { x; shape = [1, tt, n32(nHead), hd] }); a = 1; b = 2 });

  var x = op(#embedding { table = w "wte"; ids = tokens });
  for (i in Nat.range(0, nLayer)) {
    let qkv = op(#linear { x = op(#rms_norm x); w = w(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 = [1, tt, d] });
    x := op(#add { a = x; b = op(#linear { x = merged; w = w(layer(i, "proj")) }) });
    let hidden = op(#relu(op(#linear { x = op(#rms_norm x); w = w(layer(i, "fc")) })));
    x := op(#add { a = x; b = op(#linear { x = hidden; w = w(layer(i, "out")) }) });
  };
  op(#linear { x = op(#rms_norm x); w = w "head" })
};

write is the previous tutorial's generate, with the trainer's weights:

public func write(count : Nat, temperature : Float) : async Text {
  let owner = switch (trainer) { case (?t) t; case null Runtime.trap("call follow first") };
  let newline = switch (vocab.indexOf(10 : Nat8)) { case (?i) i; case null 0 };
  let out = List.empty<Nat8>();
  out.add(Nat8.fromNat(newline));
  for (_ in Nat.range(0, count)) {
    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, owner, op(#input 0), 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];
    });
    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 "" }
};

3. Run them#

Build both actors and start moxzid:

moxzi build trainer.mo
moxzi build writer.mo

trainer.mo -> trainer.wasm + trainer.did (442304 bytes, 25 file(s), 267 steps, 3.50s)
compiled by ~/.local/lib/moxzi/compiler.wasm (sha256 9294817c6e9069cbd0e712d0e9dde0ec1b6893d0e7b42d798fa3dd4c8cd772cb), moxzi 0.1.0-alpha.10
writer.mo -> writer.wasm + writer.did (318947 bytes, 23 file(s), 240 steps, 2.91s)
compiled by ~/.local/lib/moxzi/compiler.wasm (sha256 9294817c6e9069cbd0e712d0e9dde0ec1b6893d0e7b42d798fa3dd4c8cd772cb), moxzi 0.1.0-alpha.10

moxzid -m actors.json -s ./state --device burn:cpu

Load the text into the trainer, point the writer at it, and train to step 250 in the background:

jq -Rs '[.]' input.txt | curl -s -w '\n' --data-binary @- http://127.0.0.1:7000/call-json/trainer/load
TRAINER=$(curl -s -d '[]' http://127.0.0.1:7000/query-json/trainer/id | jq -r '.[0]')
curl -s -w '\n' -d "[\"$TRAINER\"]" http://127.0.0.1:7000/call-json/writer/follow
curl -s -w '\n' -d '[250]' http://127.0.0.1:7000/call-json/trainer/trainInBackground

["characters: 1115394, vocabulary: 65"]
[]
[]

When it gets there, check it:

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

["step 250 | loss 2.0896"]
["step 250: train loss 2.1469, val loss 2.1385"]

Ask the writer for text. It runs the model on the trainer's step-250 weights:

curl -s -d '[200, 0.8]' http://127.0.0.1:7000/call-json/writer/write | jq -r '.[0]'

Iserow mase that hy mine gore.

Hing.
'Sernowill word wiss made you matte:
Ind caild dive blond thee wind but, pee tois,
And whe so shey stree the as this not ine Peaor,
And fe the the whas thing the 

Resume training to step 1,000, and ask the writer again straight away. It answers with text from whichever batch the trainer published last, while the trainer keeps training:

curl -s -w '\n' -d '[1000]' http://127.0.0.1:7000/call-json/trainer/trainInBackground
curl -s -d '[200, 0.8]' http://127.0.0.1:7000/call-json/writer/write | jq -r '.[0]'

When training is done, evaluate and write once more:

curl -s -w '\n' -d '[]' http://127.0.0.1:7000/call-json/trainer/evaluate
curl -s -d '[200, 0.8]' http://127.0.0.1:7000/call-json/writer/write | jq -r '.[0]'

["step 1000: train loss 1.7230, val loss 1.8564"]
Fied, go wild o' that life all a from a not.

PLADY CALINGBEY:
Whive god thee she thousison to his sward,
Which my speaks that God herey griepe in train,
Or will I till heart: thou hast with me hands;

The losses are the ones the previous tutorial printed at step 1,000. Publishing and the writer's reads didn't change what the trainer computed.

4. Train with islands#

Islands are a way to train one model on several devices, or several machines, without sending gradients after every step. Each island trains on its own batches. Every few steps, they average their weights and carry on from the average.

Add three methods to the trainer: setIsland gives each island its own stream of batches, share publishes its weights under island/, and average sets each parameter to the mean of its own and its peers' published copies with one moxzi_device_mix call. lead runs the rounds: every island trains every steps, they share, and each averages with the others.

// Islands: several trainers learn from different random batches and average their weights
// every few steps. Each one starts from the same weights (the same seed); `island` changes
// which batches it draws.
public func setIsland(island : Nat) : async () {
  rng ^= Nat64.fromNat(island) *% 0x9E3779B97F4A7C15;
};

public func share() : async () {
  await device.moxzi_device_publish({ names = paramNames; prefix = "island/"; readers = ?#anyone });
};

// Set every parameter to the mean of this trainer's and its peers' published copies.
public func average(peers : [Principal]) : async () {
  let w = 1.0 / Float.fromInt(peers.size() + 1);
  await device.moxzi_device_mix(paramNames.map(func(name : Text) : { into : Text; terms : [{ owner : ?Principal; name : Text; weight : Float }] } = {
    into = name;
    terms = [{ owner = null; name; weight = w }].concat(
      peers.map(func(p : Principal) : { owner : ?Principal; name : Text; weight : Float } = { owner = ?p; name = "island/" # name; weight = w }));
  }));
};

// Run `rounds` rounds: every island trains `every` steps, shares, and averages with the others.
public func lead(peers : [Principal], rounds : Nat, every : Nat) : async Text {
  type Island = actor { train : Nat -> async Text; share : () -> async (); average : [Principal] -> async () };
  let me = Principal.fromActor(Self);
  let all = [me].concat(peers);
  for (_ in Nat.range(0, rounds)) {
    let trained = List.empty<async Text>();
    for (p in all.values()) trained.add((actor (p.toText()) : Island).train(every));
    for (f in trained.values()) ignore await f;
    for (p in all.values()) await (actor (p.toText()) : Island).share();
    for (p in all.values()) {
      await (actor (p.toText()) : Island).average(all.filter(func(q : Principal) : Bool = q != p));
    };
  };
  await evaluate()
};

The islands start from the same weights, because they draw them with the same seed. Load the text into each, give each its island number, and collect the peers' principals:

for i in 0 1 2 3; do
  jq -Rs '[.]' input.txt | curl -s -o /dev/null --data-binary @- http://127.0.0.1:7000/call-json/island$i/load
  curl -s -o /dev/null -d "[$i]" http://127.0.0.1:7000/call-json/island$i/setIsland
done
PEERS=$(for i in 1 2 3; do curl -s -d '[]' http://127.0.0.1:7000/query-json/island$i/id; done | jq -s -c 'map(.[0])')

Run 25 rounds of 10 steps: 250 steps on each island, 1,000 in all.

curl -s -w '\n' -d "[$PEERS, 25, 10]" http://127.0.0.1:7000/call-json/island0/lead

["step 250: train loss 2.0222, val loss 2.0416"]

At the same number of steps, the islands' average beats a single actor: it has learned from four times as many windows. It doesn't match a single actor that trains four times as long. Averaging only every 10 steps loses some of what the extra batches taught, in exchange for communicating a tenth as often.

On one machine the islands share one device, so they finish no sooner than one actor training 1,000 steps. Islands pay off when each one has its own GPU. Across machines, an island turns its publication into a one-time URL with moxzi_device_export, and its peers' devices fetch the bytes directly with moxzi_device_fetch. See GPU math for actors.

Limitations#

Next steps#

• GPU math for actors: publishing, mixing, moving weights between machines and paying for device time.
• Build a GPT in Motoko: the same ideas with every number visible.