33: Exporting models for Hugging Face and serving engines

E3301

Not a forward function a serving engine can run. A language model that serving engines run, as --vllm and --gguf require, must take the model, the token ids, [b n] Index[V] with V a dim, and optionally their positions, [b n] of an integer dtype, and give the next-token logits, [b n V]. Other parameters need defaults.

fn forward(m: Lm, ids: [b n C] Index[V]) -> [b n C V]:
error[E3301]: `forward` isn't a language model's forward function a serving engine can run

See Serving engines.

E3302

No output layer a serving engine can run. A serving engine runs the output layer itself, only for the tokens it samples, so the forward function must end with one return of linear(h, m.W) or h |> linear(m.W), with m.W a param of the model of shape [V D] and no bias.

return linear(x, m.emb, m.b_out)
error[E3302]: a serving engine needs the logits to come from an output layer it can run

Return the output layer’s result directly, and do anything else before it. See Serving engines.

E3303

No embedding a serving engine can run. A serving engine looks the tokens up in its own embedding layer, so the token ids must be used once, by embedding(ids, m.E) with m.E a param of shape [V D], and the embedding table and the output layer, which the engine holds, must be read nowhere else.

x = embedding(ids, m.emb) * mean(m.emb)
error[E3303]: `emb` is the embedding table, which a serving engine holds itself

See Serving engines.

E3304

Works across the positions a serving engine packs. A serving engine packs the tokens of many requests along the sequence axis n, one request after another, so everything the forward function does along n, in any function it passes n to, must treat each position on its own, except scaled_dot_product_attention with attn_mask=causal[n n] written in the call, which the engine replaces with its own. A sum, a softmax or a contraction across n, attention written by hand, attention with another mask, and positions computed from n are this error.

|> softmax(over=m)
error[E3304]: a serving engine can't run this: this works along `n`, which mixes its positions

Write attention as scaled_dot_product_attention(q, k, v, attn_mask=causal[n n]), and take positions as a parameter. See Serving engines.

E3305

Attention a serving engine can’t run. A serving engine keeps the keys and values of each attention in its cache, with one number of heads and one head width: every attention’s queries, keys and values must have one shape, [b H n Dh], with H and Dh dims, the same for every attention, and there must be at least one.

fn forward(m: Lm, ids: [b n] Index[V]) -> [b n V]:
error[E3305]: `forward` has no attention

Repeat keys and values with fewer heads to H before attention. See Serving engines.

E3306

Exported function that doesn’t take the model first. shapescript export --hf exports a model: its fields are the Hugging Face model’s weights, and the function it’s exported with becomes its forward, taking the function’s other parameters. So the function’s first parameter must be the model.

fn classify(x: [b D], m: Classifier) -> [b C]:
error[E3306]: `classify` must take the model first, to be exported as one

Make the model the first parameter. See Serving engines.