TinyGPT*
A 95M-parameter language model — pretrained from scratch on ~6.5B tokens of OpenWebText with a single RTX 3070 Ti (8 GB VRAM), instruction-tuned on Stanford Alpaca, and ONNX-converted so it runs locally in your browser.
* small enough to fit on a consumer GPU; large enough to write a real sentence.
Run it right now, in this tab.
Click to download the ONNX weights (~536 MB, cached on first load). Inference runs entirely on your device via WebGPU, with WASM fallback. Nothing about your prompts leaves the page.
A note — at 95M parameters, TinyGPT speaks English fluently but isn't reliable on facts. It will write a coherent sentence about anything, and confidently make things up. Treat it as a coherence demo, not a knowledge model.
Click to download the model and run it locally.
A modern decoder-only transformer, shrunk down.
Same building blocks as Llama / Mistral — Grouped-Query Attention, Rotary positions, RMSNorm — with weight tying between the embedding and the LM head to keep the parameter count honest.
- Parameters
- 95M
- Layers
- 12
- Embedding dim
- 768
- FFN dim
- 2048
- Query heads
- 12
- KV heads
- 4 (GQA, ratio 3)
- Hidden dim
- 2048
- Context
- 512 tokens
- Vocab
- 50 304 (tiktoken gpt2)
- Position
- Rotary (RoPE, θ = 10 000)
- Normalization
- RMSNorm (ε = 1e-6)
- Activation
- GELU
- Weight tying
- lm_head ↔ token_embedding
- Pretrain corpus
- OpenWebText (~6.5B tokens)
- Hardware
- 1 × RTX 3070 Ti (8 GB)
- Instruction tune
- Stanford Alpaca (52K)
Fitting a 95M model into 8 GB of VRAM.
6.7B tokens of OpenWebText, on a single RTX 3070 Ti — the same GPU you can buy used for the price of one month of ChatGPT Pro.
The whole point of TinyGPT was the constraint. 95M parameters in float32 is 380 MB just for the weights — before you account for gradients, Adam's momentum + variance state (which is 2× the weights), activations, and the KV cache. Naive training OOMs in seconds.
The tricks that made it fit:
- Mixed-precision (AMP) — forward in float16 / bfloat16, master weights in float32, dynamic loss scaling. Cuts activation memory roughly in half.
- Gradient accumulation — micro-batches that fit, accumulated to a sane effective batch size before each optimizer step.
- Fused linear + cross-entropy via the Liger Triton kernel — skips materializing the (seq × vocab) logits tensor during the loss, which is otherwise the single biggest activation in the graph.
- Gradient checkpointing on the transformer blocks — recompute activations on the backward pass instead of storing them.
Stanford Alpaca, response-only loss.
After pretraining, the base model was fine-tuned on the 52K-instruction Stanford Alpaca dataset using a standard SFT loop with a critical detail: response-only loss masking — the model is only penalized for its predictions after the ### Response: marker, not for parroting the instruction back.
This same template is used in the browser demo — your input gets wrapped before tokenization, and generation stops when the model emits one of the stop tokens listed on the right.
PyTorch checkpoint to ONNX to WebGPU.
Browsers can't load .pth directly — pickle is a code-execution format, no browser has a Python interpreter, and even if they did the model code is needed too. So the checkpoint is exported once, offline, to ONNX with three specific changes that make a Llama-style architecture web-friendly:
- RoPE rewritten in real arithmetic. The original uses
torch.view_as_complex— ONNX has no complex tensor type. Replaced with cos/sin tables passed as inputs, mathematically identical. - KV cache externalized as flat tensors. The PyTorch forward returns a list of
(K, V)tuples; ONNX wants a flat signature. 24 paired tensors (12 layers × 2) go in, 24 come out. - Causal mask as an explicit input, not SDPA's
is_causal=Trueflag. Lets one graph handle both the prefill (multi-token) and decode (single token + KV cache) phases.
The resulting ONNX runs on onnxruntime-web. WebGPU is the primary execution provider; WASM (CPU) handles fallback. Tokenization on the browser side uses js-tiktoken with the same gpt2 encoding the model was trained on.