# ONNX export, quantization and the browser runtime
## Export
`train/export.sh CHECKPOINT OUT` produces a directory laid out for
transformers.js:
```
OUT/config.json generation_config.json tokenizer.json tokenizer_config.json
special_tokens_map.json
OUT/onnx/model_q4.onnx 4-bit body, 8-bit embedding, fp32 compute (published)
OUT/SHA256SUMS
OUT/eval/model_q4_full.onnx the same weights with logits for every position
(quant_eval.py only; never published)
```
It runs `optimum-cli export onnx --task text-generation-with-past` on CPU in
a dedicated export environment (`$EXP`, or the tools on `PATH` if `EXP` is
unset), then `quantize.py`. This is a **separate venv** from training because
`optimum-onnx` 0.1 pins transformers 4.57, while training uses transformers
5.x. The exporter needs `accelerate` installed to de-duplicate GPT-2's tied
embedding weights.
## Quantization (`train/quantize.py`)
One file, for ONNX Runtime Web's WebAssembly backend (transformers.js dtype
`q4`, file `model_q4.onnx`):
1. `gemm_to_matmul()` splits GPT-2's `Conv1D` Gemms into `MatMul` + `Add`, so
every body weight is a plain `[K, N]` constant.
2. `dq4()` stores each of the 48 body weights **4-bit**: asymmetric, blocks of
32 along K, an fp32 scale and a uint4 zero point per block, as raw packed
`UINT4` tensors. One blocked `DequantizeLinear` (opset 21) per weight
restores fp32 under the original name.
3. `embedding_int8()` stores the tied token embedding / LM head **once**, as
int8 with one scale per token, in the head's `[768, 50257]` layout. The head
reads it directly; the token lookup is `Gather(axis=1)` + `Transpose`.
4. `last_token_logits()` slices the LM head's input to the last position.
transformers.js only reads the last logits, and full-sequence logits cost a
`[seq × 50257]` matmul and a large output on every prompt.
**The runtime must set `session.disable_quant_qdq = 1`** (`web/brain.ts` passes
it, with `session.disable_prepacking = 1`, through transformers.js
`session_options.extra.session`). With it, ONNX Runtime constant-folds the
`DequantizeLinear` nodes at load and runs plain fp32 kernels. Without it, its
QDQ fusion turns them into `MatMulNBits`, the 4-bit matmul, which in the
WebAssembly build is several times slower (the benchmark below did not finish
loading in 2 minutes with fusion on).
**Measurements** (headless Chromium on a 44-core x86 server, 4 threads, stock
GPT-2 small, `test/bench.html`: 32 greedy tokens after a 58-token prompt;
memory from `performance.measureUserAgentSpecificMemory()`, which includes the
WebAssembly heap at its high-water mark):
| File | Download | Decode | Page memory |
|---|---|---|---|
| fp32 | 498 MB | 34 tokens/s | 1,753 MB |
| 8-bit (`quantize_dynamic`: MatMulInteger, v1's WASM format) | 125 MB | 25–38 tokens/s | 831 MB |
| 4-bit `MatMulNBits` body (ORT's own 4-bit kernel) | 213 MB | 7.2 tokens/s | not measured |
| 4-bit unpacked with BitShift/Cast/Mul ops | 95 MB | 37.7 tokens/s | 3,118 MB (folding keeps every intermediate) |
| 4-bit `DequantizeLinear`, fp16 scales, separate head copy | 91 MB | 34.6 tokens/s | 1,483 MB |
| … with fp32 scales, no prepacking | 97 MB | 34.8 tokens/s | 1,129 MB |
| **… and one shared embedding/head table (shipped)** | **97 MB** | **35.0 tokens/s** | **804 MB** |
Phones were not measured.
**Why the embedding stays 8-bit.** Fake-quantizing the base model on pet data
(loss in nats): fp32 4.229; body 4-bit asymmetric block 32 4.255; tied
embedding 4-bit (block 32 or 16) 5.6–7.0, broken; embedding int8 per row
4.293. `quantize_dynamic`'s 8-bit head also hurt, because it quantizes the
head's input activations per tensor and GPT-2 has large activation outliers;
for the same reason `MatMulNBits` with `accuracy_level=4` (int8 activations)
must not be used here.
**ORT quirk.** Without `disable_quant_qdq`, a per-axis int8
`DequantizeLinear` on the embedding made ONNX Runtime 1.30's graph optimizer
abort with a C++ assertion (`_M_is_engaged()`). With the option set it loads.
On the base model, the exported q4 scores 4.280 against 4.195 for the fp32
ONNX export on the same rows (+2.0%). Numbers for the shipped model are in
`training.md`. Run `check_onnx.py EXPORT_DIR PROMPT [CHECKPOINT]` after any
change to `quantize.py` to catch degenerate output, and `quant_eval.py` for
held-out loss. Both set `disable_quant_qdq` like the browser.
## Optional native hosted runtime
The Rust host loads the same exported q4 graph with native ONNX Runtime and
a native tokenizer. It preserves the QDQ/prepacking settings and empty
per-call KV, and takes the context fit, sampling settings and line/lookup stop
rules from the shared rules in `core/`, as the browser worker does. One mutable session
belongs to each serialized worker. Local mode remains WebAssembly only.
Remote mode does not download a browser model; returning local loads the
pinned compatible brain from cache or downloads it. The server cannot silently
upgrade the selected brain.
See [server setup and library inputs](../server/README.md). Native CPU
correctness and browser transfers are checked separately from speed probes.
The native CUDA path uses I/O binding to keep KV tensors on the GPU between
tokens and returns only logits for CPU sampling, with TF32 disabled. Caches
are per generation and never carry across pets. Real model/HTTP checks and
limited warm-burst measurements are described in [testing](testing.md);
they do not establish sustained capacity or public deployment readiness.
Published model/runtime files remain immutable.
## Browser runtime
transformers.js 4.3.1 with ONNX Runtime Web 1.31.0-dev.20260914-8d85527a0,
vendored by `train/vendor.sh` into `runtime/transformers-4.3.1/`:
- `transformers.min.js`, the **self-contained** browser bundle. The
`transformers.web.js` build imports bare `onnxruntime-web` and needs a
bundler; do not use it.
- the `ort-wasm-simd-threaded*.{mjs,wasm}` builds (plain, `jsep`, `jspi`,
`asyncify`). A browser fetches only the one it needs. The app no longer uses
WebGPU, but the runtime directory is published and immutable, so it is kept
as is.
The ORT version must match the one the transformers.js bundle was built
against; `vendor.sh` checks the version string and refuses otherwise. Upgrading
means a new runtime directory name, never an in-place update.
## `web/brain.ts`
The TypeScript worker is built as standalone `brain.js`; its typed messages
are defined in `web/types.ts`. The main-thread client and loading fallbacks
are in `web/app/brain.ts`. Bundling does not change the vendored runtime or
model session options.
The worker imports the bundle from `CFG.runtime`, sets
`env.allowRemoteModels = false`, `localModelPath = CFG.models` (so
`modelId` resolves to `/models//`), `useBrowserCache = true`,
and `wasmPaths = /`. It then loads `AutoTokenizer` and
`AutoModelForCausalLM` with `device: "wasm"` and `dtype: "q4"`
(`onnx/model_q4.onnx`), runs a one-token warmup and posts `ready`. Download
progress is aggregated across files and drives the egg's cracks.
Before loading the runtime the worker loads the shared rules module (beside
`brain.js`). Each `generate` call carries the request the shared turn built:
a context budget (`ctx`; 1024 for react turns, 512 for idle turns), the token
budget, temperature, top-k and top-p. The worker fits the prompt with the
shared `fit` (counting with the tokenizer), stops with the shared `stop` rule,
and reports `prefillMs` (until the first new token), `tps`
(decoding speed after it) and `prompt`, the exact text it ran after trimming
old lines to fit (shown in the debug panel, `docs/testing.md`).
Generation currently prefills the whole prompt on every call, including each
lookup continuation. There is no KV reuse between calls. Reusing a cache
within one lookup turn is a possible later optimization; reusing one across
turns is harder because the clock, needs and log change near the front of the
header and old chat lines are trimmed. Neither optimization has been
implemented or benchmarked; measure on a phone before prioritizing it.
Files are stored in Cache Storage under `transformers-cache`, keyed by their
absolute URLs (`//onnx/model_q4.onnx` and the four JSON files).
The page also calls `navigator.storage.persist()` without waiting for its
optional permission prompt before starting a download or upgrade. The model
files then come from Cache Storage even if they are gone from the server; the runtime's
`.wasm` comes from the service worker's copy, the HTTP cache or the network.
When the brain is ready the worker also reports the runtime files it fetched,
so the page can ask the service worker to keep them for offline use.
### Which brain runs, and upgrades
On a visit with a hatched pet, `web/app/brain.ts` looks for cached brains among
`modelId` and the page config's `older` list and loads the newest one it has.
An egg always downloads `modelId`.
If that is an older brain, the pet runs on it and the page offers the upgrade:
a dialog once per visit (what's new from `news`, features the new brain adds,
the download size, and a mobile-data warning when `navigator.connection`
reports cellular or Save-Data) plus a banner above the device and tabs and a button in
settings that stay until the visitor upgrades. Upgrading fetches the new
model's files into `transformers-cache` with progress while the old brain keeps
running, waits until the old brain is between turns, then swaps workers; the
new brain loads from the cache without a second download (checked in the
preview: one request for the new `.onnx`). Nothing about the pet changes. If
the download fails the old brain keeps running and the banner offers to try
again.
Once a brain is ready, the page deletes cached files of every other brain under
the models URL (keeping `modelId` too while an older brain runs, so a
half-finished upgrade can resume). "Delete brain" in settings removes them all.
### Fallback
A 90-second watchdog, re-armed by each progress message, retries a stalled
load once with `numThreads = 1` (no `SharedArrayBuffer` needed), then reports
"the model runtime did not start". A load error is reported directly. Not
exercised against a real stall. Reported load errors, worker crashes and
synchronous failures constructing the worker or sending its initial load
message all cancel the watchdog and release any created worker, so a manual
retry starts a fresh worker with fresh progress. The egg shows a retry button;
settings also retain the download button. Retrying keeps saved pet/egg state
and completed cached files. Stream and ONNX protobuf-parse errors evict only
the failing brain's cached ONNX, so retry can replace an incomplete/corrupt
download. Unrelated
runtime failures retain it. A failed stream starts that file over; downloads
do not resume by range.
The service worker leaves `.onnx` requests to the browser, including at custom
model URLs. Its cache-first fetch proxy used to cut off a slow download in
Linux Firefox with `Error in input stream` after about 30 seconds. The model
still uses `transformers-cache`; this bypass does not disable offline brains.
See `testing.md` for the slow-download regression recipe.
### Why no WebGPU
v1 shipped a 4-bit fp16 WebGPU model next to the WASM one, and in a desktop
browser with a GPU it was slower than WASM. That was reported by the owner,
not measured by our scripts, so the cause below is an explanation by
reasoning, not a profile. Generation runs one token at a time with batch 1:
each step dispatches every kernel of the graph (hundreds of small GPU
dispatches for 24 layers), copies the logits back for sampling in JavaScript,
and synchronises with the GPU before the next token. For a model of this size
that fixed per-step overhead is larger than the arithmetic the GPU saves, and
ONNX Runtime's WebGPU backend may also compile new shaders when prompt lengths
change. A CPU with SIMD and a few threads does the same small matrix-vector
work without that round trip. With GPT-2 small the CPU path is fast enough,
so v2 removed WebGPU: one model file, one engine, no adapter probing,
no fp16 export, and no engine setting.
### Threads and cross-origin isolation
WASM threading needs `SharedArrayBuffer`, which needs cross-origin isolation.
The web server should send `Cross-Origin-Opener-Policy: same-origin` and
`Cross-Origin-Embedder-Policy: require-corp` for the app's base path (see
`deployment.md`). Everything the page loads is same-origin (the font is
vendored), so nothing breaks, and `crossOriginIsolated` is true. Adding any
third-party script, font, image or iframe to the page will be blocked by COEP;
host it under `/` instead.
Observed once on the dev server with v1, after flipping isolation on for a
browser profile that had cached the runtime without those headers:
`ort-wasm-simd-threaded.asyncify.mjs` failed with
`net::ERR_BLOCKED_BY_RESPONSE` until the HTTP cache was cleared. It was not
root-caused. Enable the headers before first publishing.
## Speed
`tokenling-v2` (v3 and v4 have the same architecture and file size; not re-measured), measured with `test/bench.html` in headless Chromium on a
44-core x86 server, model served from the same machine, cross-origin isolated:
| Threads | Load (from server) | Decode, greedy / sampled | Prefill, 58 tokens | Page memory |
|---|---|---|---|---|
| 4 (ORT default) | 1.6 s | 34.4 / 32.6 tokens/s | 153 ms | 805 MB |
| 1 | 1.6 s | 14.6 / 14.5 tokens/s | 546 ms | 803 MB |
In the app (same machine, 4 threads) a react turn with a few hundred tokens of
history took 240–290 ms of prefill and decoded at 29–34 tokens/s; most pet
lines are 5–20 tokens. On the live site a fresh profile went from "get an
egg" to a ready brain in 3.5 s over localhost networking. Without cross-origin
isolation WebAssembly runs on one thread (the second row). **Phones were not
measured;** expect several times slower, and use `bench.html` on a device to
find out.