Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Export

shapescript export turns a function into a program for torch.export: a .pt2 file that torch.export.load reads and runs with torch alone, with no ShapeScript and none of the program’s Python, or an AOTInductor package compiled ahead of time, which C++ can load too. torch.export traces a function with example inputs and records the torch operations it runs; its hard part is saying which input sizes may change and how they relate. ShapeScript’s types already say that, so the example inputs and the dynamic shapes come from the function’s signature, and nothing is written by hand.

Export functions

export fn marks a function that is meant to be exported. The checker then checks it, and every function it calls, at any depth, for what can’t be traced into torch operations; shapescript check and the editor report these errors as for any other. An unmarked function can be exported too, with shapescript export --fn name, which checks it the same way first. A function that something exported calls isn’t marked itself: the errors in it name the exported function that reaches it.

In an exported function, and everything it calls:

  • Python is an error. A name from import python, or anything reached through one, such as torch.load(...), torch.pi or x |> torch.tanh, is E1701: torch.export can only record torch operations, and a Python call could do anything. In a notebook, so is a Python variable. Give a torch function a type with extern fn instead.
  • An extern fn must name a function in torch. Its path must be in the module torch or one of its submodules, or in math, which works on numbers and sizes; the prelude’s stubs into shapescript.rt, such as layer_norm, are torch code too. Any other module is E1702.
  • as on a Python value is an error (E1703), since such a value already is one.
  • Host-only values and operations are errors (E1704): models, modes, logs records, sinks, gradients, optimizer state, schedules, datasets and batches, the built-ins that make or use them (print, save, load, value_and_grad, update, sgd, adamw, clip_norm, cosine, warmup, shuffle, batches, Train), changing a model’s field, and a for loop over anything but a range. An exported program is for inference: training happens before it.
  • return inside an if on a tensor is an error (E1706), as Control flow explains.

The parameters of the function being exported are the program’s inputs. Each must have a tensor type with a shape; an int, float or bool parameter with a default is fixed at the default, and anything else is E1705.

Dynamic shapes

Each input’s example tensor and the dynamic_shapes argument of torch.export.export come from the parameters’ types:

  • A literal size, or one made of dims, such as 3*D, is static.
  • Each generic axis, such as b in x: [b n D], is a torch.export.Dim, the same one wherever the axis appears, so [n D] and [n n] Mask share n.
  • A size that is linear in one generic axis is a derived dim: 16gh and (gh 16) are 16 * gh, and n+1 is n + 1. Since every axis has at least one element, each Dim’s least value makes every size derived from it at least 1.
  • Any other size, such as (h w), is Dim.AUTO, and torch works out what it can.

Each generic axis gets an example size different from every other size in the signature, so that torch can’t mistake one axis for another. Examples have the parameters’ dtypes and values in their domains, and are made on the program’s device: the file’s devices, or --devices, or the CPU. The .pt2 file doesn’t keep them.

Torch assumes that a dynamic size isn’t 1 wherever an operation would treat 1 specially, as reshape does, and the program would then refuse inputs where it is. shapescript export traces with torch’s backed_size_oblivious setting first, which reasons about 1 as about any other size, and without it if that fails, so generic axes can be 1 where torch allows it.

Assertions in exported programs

While torch.export traces, dynamic sizes are symbols, not numbers, and no values are known, so checks can’t run in Python. Instead:

  • Each shape check of an as is a torch._check per axis. torch.export takes it as a fact, so rows as [196 D] makes the rest of the function see 196 rows, and the program checks it when it runs, raising RuntimeError if it fails. --checks decides which are kept, as for build.
  • Each unknown size (?k) is a data-dependent size in the program, with what the checker knows of it, that it’s at least 0, passed to torch._check.
  • Each value check, of an as or an argument, is an aten._assert_async in the program, which raises RuntimeError with the check’s message when it fails. --value-checks decides which are kept, and debug keeps them unless the Python that exports runs with -O.
  • Argument checks become facts too; the program also checks its inputs’ shapes against the dynamic shapes itself.

So nonempty_rows(tokens[0]) as [196 D] exports, and the program raises an error when an image has an empty patch.

Control flow in exported programs

An if in an exported function becomes a function for each branch, of the values from before the if that the branches may assign, giving those back with the values both branches bind. If the condition is a bool known while the function is traced, such as one made of dims, Python runs one branch and the program has no trace of the other. Otherwise, for a [] Mask or a bool that depends on a dynamic size, the if is a torch.cond, and both branches are in the program. A branch of a torch.cond gives values rather than returning from the function, so return inside an if on a tensor is E1706; a return inside an if on a bool makes it a Python if, which needs the bool to be known while tracing.

A for i in a..b becomes a function of i and the values from before the loop that it may assign, giving those back. If both ends are numbers while tracing, as with 0..3 or 0..D, the loop is unrolled. Otherwise, as with 0..n for a generic n, it’s a torch.while_loop. A loop with a return inside is a Python loop, which needs its ends known while tracing.

In a torch.cond or a torch.while_loop, an int, float or bool that the block may assign, and a loop’s i, are 0-d tensors, since a program can’t keep a number that changes; a counter in a loop over 0..n works.

shapescript export

shapescript export model.ss --fn causal_attention -o attn.pt2
shapescript export model.ss --fn causal_attention -o attn.pt2 --aoti
shapescript export model.ss --fn forward --hf model_hf [--vllm]
shapescript export model.ss --fn forward --gguf model.gguf [--arch llama]

The last two export a model rather than a function, for Hugging Face and for serving engines. Both take --weights FILE, --seed N, --dtype float32|bfloat16|float16, --tokenizer DIR and --max-len N, and --hf takes --no-verify.

--fn names the function, and defaults to the file’s only export fn. -o defaults to the input with a .pt2 extension. --aoti writes an AOTInductor package instead of the torch.export program, with torch._inductor.aoti_compile_and_package, which needs a C++ compiler; torch._inductor.aoti_load_package loads it. --checks, --value-checks and --devices are as for build.

Tracing runs in Python, which needs torch and einops but not ShapeScript: the compiler generates a module that holds the functions, the runtime helpers they use and the export step, and runs it with --python, or $SHAPESCRIPT_PYTHON, or python3 (python on Windows). The shapescript command that a pip install provides sets SHAPESCRIPT_PYTHON to its own Python. --emit-python FILE writes the module instead; python FILE -o attn.pt2 [--aoti] runs it. A function that doesn’t exist, or a missing --fn when the file doesn’t have exactly one export fn, is an error of the command, with exit code 2; an error from torch while exporting has exit code 1.

The program’s inputs are the tensor parameters, in order. It runs with torch alone:

import torch
attention = torch.export.load("attn.pt2").module()
y = attention(x, w_qkv, w_o)       # any b and n

Exporting from Python

shapescript.export(path, fn=None, output=None, *, aoti=False, checks="always", value_checks="debug", devices=None) does the same in the calling Python and returns the torch.export.ExportedProgram. With output, it writes the .pt2 file, or the AOTInductor package with aoti=True. A file with errors raises shapescript.ShapeScriptError, and a missing function ValueError.

Hugging Face models

shapescript export asr.ss --fn transcribe --hf asr_hf --weights ckpt.pt

shapescript export --hf DIR turns a model into a Hugging Face folder, which transformers loads with trust_remote_code=True and which can be pushed to the Hub. Any model can be, not only a language model: transformers runs any PyTorch module. The function --fn (by default forward) must take the model first (E3306), and is checked as an exported function is, except that the model may be passed around and its fields read; its other parameters are tensors with shapes, or have defaults (E1705). The folder’s ShapeScriptModel holds the model’s fields as its parameters, under their own names, and its forward is the function, compiled: it takes the function’s other parameters, by position or by name, and gives what the function gives.

from transformers import AutoModel
asr = AutoModel.from_pretrained("asr_hf", trust_remote_code=True)
logits = asr(audio)                # transcribe(m, audio), for any b and t

The folder holds config.json, model.safetensors and modeling_shapescript.py, which needs torch, einops and transformers, not ShapeScript. After writing it, the export loads it with transformers, if that’s installed, and checks that it gives what the function gives, on example arguments made from its types, as for torch.export. The options are those of the next section.

If the function is a language model’s that serving engines can run, the folder is one that vLLM serves, as the next section says. If it isn’t, the export says why in a note, and config.json says so too (shapescript.vllm is false, and shapescript.unservable gives the reason); --vllm makes the serving check’s errors errors.

Serving engines

shapescript export --hf turns a language model into a Hugging Face folder that vLLM serves, and --gguf into a GGUF file for llama.cpp. A serving engine doesn’t run a model the way training does: it packs the tokens of many requests along one sequence axis, one request after another, keeps the keys and values of past tokens in a cache of its own, and computes logits only for the tokens it samples. So it runs three parts of the model itself, the token embedding, attention and the output layer, and the rest of the model as written. Usually each model is written again by hand for each engine. Here the rest of the model is the source, compiled, and the checker proves that the engine may run it that way.

The forward function, --fn, by default forward, must have this form:

fn forward(m: Llama, ids: [b n] Index[V], pos: [b n] Index[P]) -> [b n V] Logit over V:
    x = embedding(ids, m.embed)
    for l in 0..L:
        x = block(m, x, pos, l)
    return x |> rms_norm(over=D, weight=m.norm, eps=1e-5) |> linear(m.output)

These are the rules. With --hf, a function that breaks one is exported for transformers only, with a note; with --vllm or --gguf, it’s the error each rule names.

  • Its parameters are the model, the token ids, [b n] Index[V] with V a dim, and optionally their positions, [b n] of an integer dtype. Other parameters need defaults. Anything else is E3301.
  • It ends with the output layer: one return of linear(h, m.W) or h |> linear(m.W), with m.W a param of shape [V D] and no bias, which the engine runs (E3302). m.W may be the embedding table, which ties the two.
  • The token ids are used once, by embedding(ids, m.E) with m.E a param of shape [V D], which the engine runs; the embedding table and the output layer are read nowhere else (E3303).
  • Everything but attention treats each position on its own. Along n, only scaled_dot_product_attention with attn_mask=causal[n n] written in the call may work across positions; the engine puts its own attention in its place, which keeps the requests apart and reads its cache. A sum, a softmax or a contraction across n mixes requests, even where a mask would make padding exact, so attention written by hand, attention with another mask, and positions computed from n are E3304. The check follows n into every function the forward function passes it to, as padding does. Positions come in as a parameter.
  • Attention takes [b H n Dh]: queries, keys and values of one shape, with H and Dh dims, the same for every attention, and at least one attention (E3305). Keys and values with fewer heads are repeated to H before attention, as repeat(b Hkv n Dh -> b (Hkv R) n Dh) does.

Everything else that an exported function can’t do, a serving export can’t either, except that the model may be passed around and its fields read.

The engine runs the attentions in the order the forward function runs them, each with its own part of the cache, so the forward function must run the same attentions on every call; for l in 0..L over a dim does. The checks of devices, dtypes, models and argument values are left out, since engines give the model its inputs on their own device and often in a lower precision; shape checks stay.

Hugging Face folders and vLLM

shapescript export llama.ss --hf llama_hf --weights ckpt.pt
vllm serve llama_hf --trust-remote-code --model-impl transformers --skip-tokenizer-init

The folder holds config.json, generation_config.json, model.safetensors and modeling_shapescript.py: the compiled forward function, the runtime helpers it uses, and the classes ShapeScriptConfig, ShapeScriptModel (the model up to the output layer) and ShapeScriptForCausalLM, which transformers loads with trust_remote_code=True. It needs torch, einops and transformers, not ShapeScript. The embedding table is model.embed_tokens.weight, the output layer lm_head.weight unless it’s tied, and every other field model.<field>, with its shape as declared, stacked layers included. num_hidden_layers is the number of attentions the forward function runs, counted by running it once.

Each attention calls the attention function the model is set to use, from transformers’ table: by default eager, which is the source’s attention, so transformers gives the source’s logits; vLLM’s Transformers backend sets its own, which keeps each layer’s keys and values in its paged cache, and runs everything else through torch.compile and CUDA graphs. vLLM’s own embedding and output layer replace embed_tokens and lm_head. transformers keeps no cache of past keys and values for the model (use_cache is false), so generate there computes each step from the start; vLLM keeps its own. Keys and values with fewer heads than queries are given to the engine repeated, so its cache holds H heads.

An engine reads each attention’s scale once, when it builds the model: ShapeScriptAttention.scaling is the scale the source gives, recorded in config.json as attention_scales when the source gives one, or torch’s default, $1/\sqrt{D_h}$; a scale that differs on a later call is an error. In a model served in bfloat16 or float16, what the source computes in f32 on purpose stays f32, as rotary angles from pos to f32 usually are, and attention and the output layer are given the model’s dtype, as the engine runs them.

After writing the folder, the export loads it with transformers, if that’s installed, and checks that it gives the logits the source computes, to 1e-4 of their size in float32 (5e-2 in the other dtypes); --no-verify (verify=False) skips this.

Without --weights, the weights are new ones made from --seed (by default 0), as M.init(seed=…) makes them; --weights takes a checkpoint that save wrote, whose first model of the forward function’s type is used. --dtype is float32 (the default), bfloat16 or float16. --tokenizer DIR copies a Hugging Face tokenizer into the folder; without one, the engine takes and gives token ids, as --skip-tokenizer-init tells vLLM. --max-len sets max_position_embeddings, by default P of the positions’ Index[P], or 4096. Tensor and pipeline parallelism aren’t supported yet.

GGUF files and llama.cpp

shapescript export llama.ss --gguf llama.gguf --weights ckpt.pt
llama-server -m llama.gguf

A GGUF file holds weights and the name of an architecture, and llama.cpp runs it with its own implementation of that architecture, so a model can only be written as one whose computation is the source’s. --gguf finds which, of llama and gpt2 (or --arch), and which field is which of its tensors, by computing both on random weights in f64: the forward function as the source writes it, and the architecture as llama.cpp computes it. Fields are paired with tensors in the order the forward function first reads them, then in the other orders of fields of one shape, and for llama, with each common rotary base, and with the query and key rows reordered as convert_hf_to_gguf.py reorders Hugging Face’s, for rotary embeddings that turn the two halves of a head rather than neighbouring features. The first pairing whose logits agree to 1e-8 of their size is written. If none does, nothing is written, and the error says why for each architecture: a different norm, a different number of tensors, or the closest pairing’s difference, as for GELU computed exactly where llama.cpp’s gpt2 uses tanh. Every norm’s eps must be one number written in the calls, which the file holds once.

  • llama: rms_norm, rotary position embeddings on neighbouring features, causal attention with as many key and value heads as the key projection has rows for, and a SwiGLU feed-forward layer. examples/llama.ss is one.
  • gpt2: layer_norm, learned position embeddings, a fused query, key and value projection with biases, and GELU with tanh. examples/gpt2.ss is one.

Matrices are written in --dtype, and 1-D tensors in f32. --tokenizer DIR writes a byte-level BPE tokenizer, such as GPT-2’s, from tokenizer.json; without one, the file’s vocabulary names each token <|N|>, which llama.cpp needs, and prompts are token ids. llama.cpp’s own kernels compute in lower precision in places, such as half-precision attention on a GPU and a GELU table on the CPU, so its log-probabilities differ from the source’s by up to about 1e-2 on a GPU and 1e-3 on the CPU, with its cache in f32.

Serving exports from Python

shapescript.export_hf(path, output, *, fn="forward", weights=None, seed=0, dtype="float32", tokenizer=None, max_len=None, verify=True, vllm=False) writes the folder and returns the config.json it wrote; weights may also be a model. shapescript.export_gguf(path, output, *, fn="forward", arch=None, weights=None, …) writes the GGUF file and returns its architecture, or raises ValueError saying why the model isn’t one. A file with errors raises shapescript.ShapeScriptError.