GPT-X3-150M — ONNX (fp32 / fp16 / int8 / int4)

ONNX exports and quantizations of AxiomicLabs/GPT-X3-150M, Axiomic Labs' hybrid GDN2 / gated-GQA base model. All credit for the model, its training and its data goes to Axiomic Labs; this repo only changes the file format and numeric precision.

The models run with plain onnxruntime + numpy + tokenizers — no PyTorch, no custom kernels.

Files

*Perplexity on a small mixed evaluation set (narrative, science, Python code, civic prose; ~250 tokens). The int8 perplexity sitting slightly below fp32 is noise at this sample size, not an improvement. "KL vs fp32" is the mean per-token KL divergence of the next-token distribution from the fp32 export, and "Top-1 agreement" is how often the most likely next token matches fp32.

Which one should I use?

• Want output closest to the original: model_fp16.onnx (GPU) or model_int8.onnx (CPU).
• Want small and fast on CPU: model_int8emb.onnx.
• Every MB counts: model_int4mix.onnx, accepting ~7% higher perplexity.

Quick start

pip install onnxruntime numpy tokenizers huggingface_hub

from pathlib import Path
from huggingface_hub import snapshot_download

repo = Path(snapshot_download("ahmadw/GPT-X3-150M-ONNX",
                              allow_patterns=["*.json", "*.py", "onnx/model_int8emb.onnx"]))

cd <downloaded repo>
python onnx_generate.py "Once upon a time," --onnx onnx/model_int8emb.onnx --max-new-tokens 64
python onnx_generate.py "def is_prime(n):" --onnx onnx/model_int8emb.onnx --temperature 0.7

Or from Python:

from tokenizers import Tokenizer
from onnx_generate import GPTX3OnnxModel, generate

model = GPTX3OnnxModel("onnx/model_int8emb.onnx")
tokenizer = Tokenizer.from_file("tokenizer.json")
print(generate(model, tokenizer, "Once upon a time,", max_new_tokens=64))

GPT-X3-150M is a base model: it continues text, it is not instruction-tuned.

Graph signature

The GDN2 recurrence is inherently a per-token loop, so instead of baking a fixed sequence length into the graph, each model is a single-token decode step with the caches passed explicitly:

Outputs are logits [batch, 65536] followed by new_state.* in the same order. Start with zero-filled states (KV caches with past = 0) and feed prompt tokens one at a time; GPTX3OnnxModel.prefill() / step() in onnx_generate.py do exactly this. All inputs and outputs are float32 in every variant, so the variants are drop-in interchangeable.

How the variants were made

fp32 export (export_onnx.py): torch.onnx.export(dynamo=True) of a decode-step wrapper that reuses the original module weights and reimplements the GDN2 step and gated GQA step with explicit state; batch size and KV length are dynamic. It was validated position by position against the PyTorch model (max |logit diff| 2.9e-4, identical argmax) and produces the same greedy text as transformers generate().

Quantizations (quantize_onnx.py):

• fp16 — onnxruntime.transformers.float16, keeping float32 I/O.
• int8 — onnxruntime dynamic quantization, per-channel symmetric int8 weights for all MatMuls. The tied embedding / LM-head table stays fp32, which also leaves the 65 536-wide LM head in fp32.
• int8emb — int8, plus the tied embedding / LM-head table stored once in 8-bit with one scale per 32-value block, consumed by a MatMulNBits (8-bit) LM head and by an in-graph dequantizing lookup. Block-wise scales are essential here: two hidden channels of the embedding (356 and 350) are ~10× larger than the rest, and per-row int8 scaling crushed every other channel (KL 0.26 in testing). With block scales, KL drops to 0.012.
• int4mix — int8emb, plus 4-bit symmetric block-32 (MatMulNBits, RTN) weights for the attention and GDN2 projections; the MLP stays int8.

Why not full int4? Tested and rejected. 4-bit quantization of every MatMul raised perplexity 12–27% (KL 0.07–0.16). The MLP turned out to be the most int4-sensitive block (KL 0.055 alone), and HQQ only partially recovered it (KL 0.043) at a large CPU speed cost. 150M parameters at hidden size 576 leave little redundancy for 4-bit weights.

Evaluation

eval_onnx_quality.py reproduces the table above: per-token perplexity, KL and top-1 agreement against the fp32 export on a fixed text set, plus 48-token greedy generations from four prompts, reporting how many tokens match fp32 before diverging.

Divergences were inspected: the fp16/int8 split at token 17 is a near-tie in fp32 itself (0.068 vs 0.065 probability), and the diverged continuations of the int8 builds remain fluent and on topic. The evaluation set is small; treat these numbers as a sanity check and measure on your own domain before deploying.

Reproducing

# Inside a clone of AxiomicLabs/GPT-X3-150M (needs torch, transformers, onnx, onnxscript, onnxruntime)
HF_MODULES_CACHE=/tmp/hf_modules python export_onnx.py   # -> onnx/model.onnx (+ verification)
python quantize_onnx.py                                  # -> fp16, int8, int8emb, int4mix
python eval_onnx_quality.py                              # quality table

Built with torch 2.10, onnx 1.20, onnxscript 0.6.2 and onnxruntime 1.28.

Limitations

• Prompts are prefilled token by token (no parallel prefill graph), so time-to-first-token grows linearly with prompt length. This is cheap at 150M parameters but noticeable for very long prompts.
• model_int8emb.onnx and model_int4mix.onnx use the com.microsoft MatMulNBits operator and therefore require ONNX Runtime; other ONNX runtimes may not support them.
• Tested on CPU (Apple Silicon) only.
• All limitations and biases of the base model apply; see the original model card.

License

Apache-2.0, same as the base model.