Exporting without Python

A trained model often has to run where your Python code can’t: a server without your repository, a C++ application, a phone. PyTorch’s answer is torch.export, which records a function’s torch operations into a program that runs on its own. Its hard part is dynamic shapes: you must declare by hand which input sizes may change and how they relate, and you find out about mistakes when tracing fails.

ShapeScript’s types already say all of that. So shapescript export builds the dynamic shapes from the signature.

export fn

Mark a function export:

dim D = 768

## Splits each image into 16×16 patches.
export fn patchify(img: [b 3 16gh 16gw]) -> [b (gh gw) D]:
    return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c))

and export it:

shapescript export patches.ss -o patchify.pt2

Exporting runs torch.export in a Python with PyTorch and einops (--python, or the active virtualenv’s), but the program it writes needs neither ShapeScript nor your code:

import torch
patchify = torch.export.load("patchify.pt2").module()   # no ShapeScript needed
print(patchify(torch.randn(4, 3, 224, 160)).shape)
print(patchify(torch.randn(1, 3, 32, 64)).shape)
patchify(torch.randn(1, 3, 30, 64))

torch.Size([4, 140, 768])
torch.Size([1, 8, 768])
RuntimeError: Expected input *args[0].shape[2] = 30 to be of the form 16*s68, where s68 is an integer

The batch b stayed dynamic, and so did the image size, as a multiple of 16: 16gh became the derived dimension 16 * gh, and the program refuses a height of 30. You wrote none of that. Each generic axis becomes a torch.export.Dim, shared wherever it appears; dims and literal sizes are fixed.

--aoti writes an AOTInductor package instead, compiled ahead of time, which C++ can load too (it needs a C++ compiler). From Python, shapescript.export("patches.ss", "patchify", "patchify.pt2") does the same and returns the ExportedProgram.

What can be exported

torch.export records torch operations, so an exported function, and everything it calls, can’t use Python. The checker enforces it as you type:

import python torch

dim D = 768

export fn load_qkv() -> [D 3*D]:
    return torch.load("w_qkv.pt") as [D 3*D]

$ shapescript check export_bad.ss
error[E1701]: `torch.load` is Python, which an exported function can't use
 --> export_bad.ss:6:12
  |
6 |     return torch.load("w_qkv.pt") as [D 3*D]
  |            ^^^^^^^^^^^^^^^^^^^^^^ runs Python
  |
  = note: `torch.export` records torch operations only, and code from `import python` could be anything
  = help: pass the value in as a parameter, or declare a torch function with `extern fn`, as in `extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc"`

For more information about E1701, run `shapescript explain E1701`.

Also not allowed: models, modes, datasets, logs, gradients and optimizers, and the built-ins that use them. An exported program is for inference; training happens before it. An extern fn may be used if it names a function in torch (or math), so give a torch function a type with extern fn rather than calling it through import python.

What does carry over:

• as assertions become facts in the program. rows as [196 D] tells torch.export the size, and the program checks it when it runs.
• Unknown sizes (?k) become data-dependent sizes.
• if and for become torch.cond and torch.while_loop where they depend on values or dynamic sizes, and plain Python, unrolled, where they don’t.
• Iterator chains become torch.vmap or plain tensor operations, which torch.export records.

Serving a language model

A language model is served differently from how it’s trained. A serving engine such as vLLM or llama.cpp takes requests from many users at once, packs their tokens one after another along the sequence axis, and keeps the keys and values of past tokens in a cache, so each new token costs one step rather than a whole forward pass. To do that, an engine runs three parts of the model itself: the token embedding, attention and the output layer. Usually someone then writes the whole model again for each engine, and hopes it computes the same thing.

shapescript export --hf writes the model for vLLM instead, with everything but those three parts compiled from your source. Here is a small language model, which we’ll make servable:

dim D  = 64
dim H  = 4
dim Dh = D / H
dim V  = 500
dim L  = 2

model Tiny:
    param wte:   [V D]     = normal(std=0.02)
    param w_qkv: [L 3*D D] = normal(std=0.02)
    param w_o:   [L D D]   = normal(std=0.02)
    param norm:  [L D]     = ones

fn attention(x: [b n D], w_qkv: [3*D D], w_o: [D D]) -> [b n D]:
    q, k, v = linear(x, w_qkv) |> split(b n (3 H Dh) -> b H n Dh)
    att = einsum(q, k: b H n Dh, b H m Dh -> b H n m)
        |> scale(1 / sqrt(Dh))
        |> masked_fill(!causal[n m], -inf)
        |> softmax(over=m)
    return einsum(att, v: b H n m, b H m Dh -> b H n Dh)
        |> rearrange(b H n Dh -> b n (H Dh))
        |> linear(w_o)

fn forward(m: Tiny, ids: [b n] Index[V]) -> [b n V] Logit over V:
    x = embedding(ids, m.wte)
    for l in 0..L:
        x = x + attention(rms_norm(x, over=D, weight=m.norm[l]), m.w_qkv[l], m.w_o[l])
    return linear(x, m.wte)

It checks, and it trains. shapescript export tiny.ss --hf tiny_hf writes it as a Hugging Face folder, which transformers loads like any other model; any model can be exported that way, a classifier or a speech recognizer too, with its function as the model’s forward. But it notes that vLLM can’t serve this one, and --vllm, which asks for a model vLLM can serve, says why:

$ shapescript export tiny.ss --hf tiny_hf --vllm
error[E3304]: a serving engine can't run this: this works along `n`, which mixes its positions
  --> tiny.ss:18:12
   |
18 |         |> softmax(over=m)
   |            ^^^^^^^ works across `n`
   |
   = note: a serving engine packs the tokens of many requests along `n`, one request after another, so everything but attention must treat each position on its own; it replaces `scaled_dot_product_attention(…, attn_mask=causal[n n])` with its own, which keeps the requests apart
   = help: write attention as `scaled_dot_product_attention` with `attn_mask=causal[n n]` in the call, and take positions as a parameter `pos: [b n] Index[P]` rather than computing them from `n`

With requests packed one after another, a softmax over every position in the batch would mix one user’s tokens into another’s. The engine’s own attention keeps them apart, but it can only stand in for attention it recognizes. Attention written by hand, with einsum and softmax, could be anything. The checker follows n through every function forward passes it to, and the only thing allowed to work across positions is scaled_dot_product_attention with attn_mask=causal[n n] in the call. So write attention with it:

fn attention(x: [b n D], w_qkv: [3*D D], w_o: [D D]) -> [b n D]:
    q, k, v = linear(x, w_qkv) |> split(b n (3 H Dh) -> b H n Dh)
    return scaled_dot_product_attention(q, k, v, attn_mask=causal[n n])
        |> rearrange(b H n Dh -> b n (H Dh))
        |> linear(w_o)

It computes the same thing, and now the export succeeds:

$ shapescript export tiny.ss --hf tiny_hf --vllm --weights ckpt.pt
wrote tiny_hf: 2 attention layers of 4 heads, vocabulary of 500; vLLM can serve it

tiny_hf is an ordinary Hugging Face folder. Its modeling_shapescript.py is your forward, compiled, with three changes: the embedding and the output layer are left to the engine, and each attention calls whichever attention function the engine sets. transformers uses your attention by default, so it gives your model’s logits:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("tiny_hf", trust_remote_code=True)

and vLLM uses its own paged attention, with everything else running as your code says:

vllm serve tiny_hf --trust-remote-code --model-impl transformers --skip-tokenizer-init

There is nothing to keep in sync: change the model, export again.

The other rules are like this one. forward must take the model, the token ids and optionally their positions; embed the ids once, with embedding; and end by returning linear(h, m.W), the output layer, which the engine runs for the tokens it samples. A model that needs positions takes them as a parameter, pos: [b n] Index[P], rather than computing them from n: in a packed batch, the index along n isn’t a token’s position.

llama.cpp is different: a GGUF file holds weights and the name of an architecture, and llama.cpp runs its own implementation of that architecture. So --gguf doesn’t compile your model. It works out which of llama.cpp’s architectures your model is, by computing both on random weights, and refuses if none computes what yours does:

$ shapescript export tiny.ss --gguf tiny.gguf
error: `Tiny` isn't an architecture llama.cpp has: not `llama`: it has 3 fields stacked along its 2 blocks and 0 other fields, but `llama` has 9 tensors in each block and 1 other tensor; not `gpt2`: llama.cpp's architecture normalizes with `layer_norm` only, but the model uses rms_norm
error: exporting `forward` failed: `python3` exited with exit status: 1

The repository’s examples/llama.ss is a model --gguf writes as llama, and examples/gpt2.ss one it writes as gpt2.

Exercises

1. Export causal_attention from the Introduction. Which of its axes are dynamic in the program, and which are fixed?

2. Why is this an error, and how do you fix it while keeping it exportable?

   import python torch
   
   dim D = 64
   
   export fn act(x: [b D]) -> [b D]:
       return x |> torch.nn.functional.mish
   

3. Give Tiny learned position embeddings: a param wpe: [P D], and a parameter pos: [b n] Index[P] of forward whose embeddings are added to the tokens’. Why must the positions be a parameter, rather than computed with n |> map(i => i)? What does shapescript export --hf --vllm say if you try?

4. Make Tiny’s blocks those of examples/llama.ss: an RMSNorm and a SwiGLU feed-forward layer after attention, rotary embeddings, and a final norm. Export it with --gguf, and serve it with llama-server.