GPU math for actorsalpha

Give hosted actors a tensor device for training and running models on a GPU or the CPU.

Start moxzid with --device and every actor it hosts can train and run neural networks on a GPU. An actor reaches the device through the management canister, aaaaa-aa, with methods named moxzi_device_*. This page trains a small model from Motoko, then covers sharing weights, moving them between machines, paying for device time and surviving a crash.

Choose a device#

The examples on this page use burn:cpu so they run on any machine. Switch to a GPU device for real work: the actor code doesn't change.

Describe the device to Motoko#

Declare the methods you use as an actor type and bind it to aaaaa-aa. A program is a list of operations. Each operation's result is a register, numbered by its position in the list, that later operations refer to.

type Op = {
  #input : Nat32;
  #param : Text;
  #narrow : { x : Nat32; dim : Nat32; start : Nat32; len : Nat32 };
  #reshape : { x : Nat32; shape : [Nat32] };
  #embedding : { table : Nat32; ids : Nat32 };
  #linear : { x : Nat32; w : Nat32 };
  #cross_entropy : { logits : Nat32; targets : Nat32 };
};
type Device = actor {
  moxzi_device_params : ([{ name : Text; shape : [Nat32]; init : { #normal : Float } }]) -> 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] };
};
transient let device : Device = actor "aaaaa-aa";

The full interface, with every operation and option, is in interface.did in the moxzi repository under moxzi/device/.

Train a model#

This actor learns to predict the next letter of a phrase. It creates two parameters, an optimizer, and sends one training program per step. Tensors stay on the device: each step sends the text and gets the loss back as a number.

// 65 bytes of text: the model reads the first 64 and predicts each next one.
let text = Text.encodeUtf8("hello moxzi, hello actors, hello tensors on a device. hello again");

// Register numbers below are positions in `program`.
let program : [Op] = [
  #input 0,                                            // 0: the 65 bytes, [1, 65]
  #narrow { x = 0; dim = 1; start = 0; len = 64 },     // 1: what the model reads
  #narrow { x = 0; dim = 1; start = 1; len = 64 },     // 2: what it should predict
  #reshape { x = 2; shape = [64] },                    // 3
  #param "embed",                                      // 4: a vector per byte value
  #embedding { table = 4; ids = 1 },                   // 5: [1, 64, 16]
  #param "head",                                       // 6
  #linear { x = 5; w = 6 },                            // 7: a score per byte value, [1, 64, 256]
  #reshape { x = 7; shape = [64, 256] },               // 8
  #cross_entropy { logits = 8; targets = 3 },          // 9: the loss
];

public func learn(steps : Nat) : async { first : Float; last : Float } {
  await device.moxzi_device_params([
    { name = "embed"; shape = [256, 16]; init = #normal 0.1 },
    { name = "head"; shape = [256, 16]; init = #normal 0.1 },
  ]);
  await device.moxzi_device_optimizer({
    name = "adam"; params = ["embed", "head"]; lr = 0.05;
    beta1 = 0.9; beta2 = 0.99; eps = 1e-8; weight_decay = 0.0; kind = ?#adamw;
  });
  var first = 0.0;
  var last = 0.0;
  var i = 0;
  while (i < steps) {
    let r = await device.moxzi_device_run({
      inputs = [{ data = text; elem = #u8; shape = [1, 65] }];
      program;
      train = ?{ loss = 9; optimizer = "adam"; clip = null; lr = null; also = null };
      scalars = [9];
      tensors = [];
    });
    if (i == 0) first := r.scalars[0];
    last := r.scalars[0];
    i += 1;
  };
  { first; last };
};

Build it and start moxzid with a device:

moxzi build letters.mo

moxzid -m actors.json --device burn:cpu

Train for 200 steps. The loss starts near 5.5 (a uniform guess over 256 byte values) and falls as the model learns the phrase:

curl -s -w '\n' -d '[200]' http://127.0.0.1:7000/call-json/letters/learn

[{"first":5.547374725341797,"last":0.8760831356048584}]

Each actor has its own private space on the device. Parameters, optimizer state and stored values are keyed by the actor's principal, so two actors can both have a parameter named head without seeing each other's.

Share weights between actors#

Weights move between actors on the device, by name. They never pass through a message.

Averaging the weights of several training actors, or copying a model into a new actor, is one moxzi_device_mix call.

Move weights between machines#

Each moxzid drives its own device. To let a device on another machine fetch your publications, give moxzid an address to serve them on:

An actor calls moxzi_device_export to turn a publication into a one-time URL, and sends that URL to an actor on the other machine. That actor calls moxzi_device_fetch, and its device downloads the bytes directly. Only the URL travels in messages. A URL works once by default and expires after ten minutes. Exports have no other authentication, so put TLS in front of --device-listen on an untrusted network.

Pay for device time#

When moxzid charges cycles (--cycles greater than 0), each device call costs the calling actor --device-cycles-per-ms cycles per millisecond it took. The default is 100,000,000.

Survive a crash#

With -s DIR, the device's contents survive a crash along with the actors. moxzid records every device call that changes something, and saves a full copy of the device once the calls recorded since the last copy took --device-checkpoint-secs seconds of device time (default 30), or once the record passes --device-checkpoint-mb megabytes (default 64), whichever comes first. On startup it restores the device before any message runs. Each actor then replays its messages against the answers its device calls got the first time, so it comes back exactly as it was.

Startup re-runs every device call made since the last copy, which takes about as long as those calls took the first time: about --device-checkpoint-secs. A smaller value makes startup faster and copies the device more often. A copy of a large model takes time too, so moxzid never copies more often than every ten copy-durations. If a startup had more than that to re-run (the server ran with a larger setting before), it copies the device straight away, so the next startup does not re-run the same calls. --device-checkpoint-secs 0 copies by size only.

Limitations#