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.

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 <base>/models/<id>/), useBrowserCache = true, and wasmPaths = <runtime>/. 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.

Each generate call takes a context budget (ctx; 1024 for react turns, 512 for idle turns) 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 (<models>/<id>/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 <base>/ 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.