# 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.

```bash
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`](https://sx.041.io/docs/defrun.md#weight-averaging), 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:

```console
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

```console
$ 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](https://sx.041.io/docs/checkpoints.md#locations), or a directory
[`export`](#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.

| flag                                   | meaning                                                                                           |
| -------------------------------------- | ------------------------------------------------------------------------------------------------- |
| `--socket <path>`                      | the Unix socket to listen on                                                                      |
| `--batch N`                            | the most calls run together on the device; 1 without it                                           |
| `--output NAME`                        | only this output, repeatable, in the order given; every output without it                         |
| `--at`                                 | a last input `at`, an `i32` of shape `[1]`, and each output only at that place of its second axis |
| `--precision bf16`                     | run the model under the `bf16` policy, whatever it trained under                                  |
| `--trained`                            | the trained weights instead of the `defrun :ema` averages; a checkpoint's model only              |
| `--device cpu\|reference\|cuda\|metal` | where the model runs; `SEXPGPU_DEVICE`, then `cpu`, without it                                    |
| `--cpu-cores N`                        | `SEXPGPU_CPU_CORES`; maximum host worker cores for `--device cpu`, positive integer               |

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`:

```python
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:

```console
$ 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
```

| flag                               | meaning                                                                                  |
| ---------------------------------- | ---------------------------------------------------------------------------------------- |
| `--checkpoint <location>/step-<n>` | the checkpoint whose weights to export, any kind of [location](https://sx.041.io/docs/checkpoints.md#locations) |
| `--checkpoint <location>`          | the location's newest checkpoint                                                         |
| `-o <dir>`                         | where to write; a directory that exists is refused                                       |
| `--trained`                        | write the trained weights instead of the `defrun :ema` averages                          |

The file is compiled under the variant and knobs the checkpoint records, as
[`evaluate`](https://sx.041.io/docs/evaluation.md#sexpgpu-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`.

```text
/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](https://sx.041.io/docs/checkpoints.md), [evaluation passes](https://sx.041.io/docs/evaluation.md),
[devices](https://sx.041.io/docs/devices.md).

---

SexpGPU documentation. Every page: https://sx.041.io/llms.txt