Running a trained model

Every checkpoint holds its model, so every run is ready to run as a model: sexpgpu serve takes a checkpoint and puts its model behind a Unix socket, bytes in and bytes out, for trying the model and prototyping the program around it: a sampling loop that runs it until an end token, a filter beside it, anything that needs the model's outputs and not a training run.

A library that runs a model inside your own serving stack, without a socket between them, is coming.

sexpgpu run my-run.sx --checkpoint-dir o41://my-team/my-run
sexpgpu serve o41://my-team/my-run --socket /tmp/my-model.sock --device cuda \
    --batch 8 --output logits --at

The model

A run with a checkpoint location compiles the model alone beside the run, and every checkpoint holds it as model.json beside its weights. The model is model, applied as a step applies it, to what eval-prepare, or prepare without one, returns as :inputs from a batch of one record of the training loader's fields.

• Inputs: each tensor :inputs holds becomes one input of the shape and dtype prepare makes it for one record, inputs when it is one tensor and inputs.0, inputs.1, ... when it is a list. What prepare computes to make them is not part of the model: a caller passes the model's inputs, not the loader's records. A value in :inputs that is not a tensor, a keyword say, is passed to the model as it is.
• Outputs: the tensors of the prediction, in lists and plists to any depth, each named by its path: output for one tensor, output.0, output.1, ... in a list, and a plist's keys, (list :logits l :aux (list a b)) giving logits, aux.0 and aux.1. Numbers, booleans, strings and nil in the prediction are left out; anything else is E-INFER-001. The model's metrics, counters and guards are left out too.
• Weights: under defrun :ema, the averages, as the passes read them; --trained gives the model the trained weights.
• Precision: the run's defrun :precision, as its evaluation runs it; each output keeps the dtype the model gives it.
• Shapes are fixed. A sequence model takes its whole window every call. To run a shorter prompt through a causal model, pad it after its end and read the outputs at its last position.

A model whose prediction holds anything but tensors cannot run alone. The run says so as it starts, with E-INFER-001, and trains and checkpoints as always, its checkpoints without a model:

note: the checkpoints will hold no model, so they run only from the file:
error[E-INFER-001]: a model run alone returns tensors, in lists and plists; `output.1` is a model

serve

$ sexpgpu serve /tmp/tiny-run --socket /tmp/tiny.sock --batch 4
serving /tmp/tiny-run/step-00000020 on /tmp/tiny.sock, up to 4 a call
  inputs  inputs [1, 16] i32
  outputs output [1, 16, 32] f32
^Csexpgpu serve: stopped by SIGINT, 256 calls in 135 runs

The operand is a checkpoint, <location>/step-<n> or a location for its newest one, any kind of location, or a directory export wrote. A frozen :load parameter, which a checkpoint names instead of holding, is read from the file its run read, a relative one from the current directory, and refused with E-LOAD-001 when that file has changed.

The three that shape the model are how a program asks for no more than it reads:

• --output leaves out what the others alone compute. A model that returns more than its prediction, a training-only head or the parts of a mixture, does none of that work, and its call moves only what was named.
• --at serves a sampling loop: it reads the logits after its prompt's last token, one row instead of the window's, so a call moves that row. Every output must have a second axis of one length; an at outside it is an error.
• --precision bf16 runs the products in bf16, faster and less exact: compare the model's evaluation numbers both ways before relying on it.

serve is for trying a model and building against it; a production stack embeds the model instead. The weights go to the device once and stay there; each call moves only its inputs and outputs. Clients connect side by side. Calls that arrive while the device is busy run together as one batch when it is free, up to --batch, so a call that arrives alone runs at once and waits for no other. A call's outputs are the same bits whatever it runs with.

SIGTERM, SIGINT (Ctrl-C) or SIGHUP stops it: the socket is removed, the calls it answered and the device runs they took are printed, and the process exits 128 + n. A socket a stopped process left behind is replaced; one another process still serves is refused, and so is a path that holds anything but a socket.

The protocol

Every message either way is a frame: its length in bytes as a little-endian u64, then the bytes.

1. On connecting, the client reads one frame: the checkpoint, its step and slug, the precision, and the model's inputs and outputs as [{"name": "inputs", "shape": [1, 16], "dtype": "i32"}], as JSON, as the flags left them.
2. Each call is one frame from the client holding every input of one record in the order inputs lists them.
3. Its answer is one frame starting with a status byte: 0 then every output in the order outputs lists them, or 1 then the error, UTF-8 text.

A tensor is its elements in row-major order, each little-endian at its dtype's width, with no header, since the shapes are fixed: f32 and i32 4 bytes, f64 and i64 8, bf16 2 (the high half of the f32, rounded to nearest even), bool 1, 0 or 1. A call frame of another length than the inputs take is answered with an error, and the connection closed.

A client

A greedy sampling loop in Python, over a model that takes one i32 window and returns its logits, served with --at:

import json, socket, struct
import numpy as np

def frame(s):
    (n,) = struct.unpack("<Q", s.recv(8, socket.MSG_WAITALL))
    return s.recv(n, socket.MSG_WAITALL)

s = socket.socket(socket.AF_UNIX)
s.connect("/tmp/my-model.sock")
about = json.loads(frame(s))
(window, _), (made,) = about["inputs"], about["outputs"]

def next_logits(tokens):
    x = np.zeros(window["shape"], np.int32)
    x[0, : len(tokens)] = tokens
    at = np.array([len(tokens) - 1], np.int32)
    s.sendall(struct.pack("<Q", x.nbytes + at.nbytes) + x.tobytes() + at.tobytes())
    reply = frame(s)
    if reply[0] != 0:
        raise RuntimeError(reply[1:].decode())
    return np.frombuffer(reply[1:], np.float32).reshape(made["shape"])[0, 0]

tokens = [5, 8]
while len(tokens) < window["shape"][1]:
    tokens.append(int(next_logits(tokens).argmax()))
print(tokens)

export

sexpgpu export writes a checkpoint's model and weights as one directory of their own, to copy where the checkpoint is not, or for a checkpoint that holds no model, written by an earlier version or by a run whose model did not run alone then:

$ sexpgpu export tiny.sx --checkpoint /tmp/tiny-run -o /tmp/tiny-model
export: /tmp/tiny-run/step-00000020
wrote /tmp/tiny-model: 32 weights, 109.8 KiB
  inputs  inputs [1, 16] i32
  outputs output [1, 16, 32] f32

The file is compiled under the variant and knobs the checkpoint records, as evaluate compiles it, and refused as evaluate refuses a checkpoint whose parameters or :load files are not the file's. Edited sources are allowed; the lines that moved are printed after export:. A model that cannot run alone is E-INFER-001.

/tmp/tiny-model/model.json           the model, as a checkpoint holds it, with the
                                     checkpoint, its step and slug
/tmp/tiny-model/params.safetensors   every weight the model reads, keyed by its path,
                                     a frozen :load parameter's included

As in a checkpoint, a bf16 weight is stored as f32 and model.json says what it means.

Related: checkpoints, evaluation passes, devices.