tiny-titan-gpt (JAX / Flax NNX)

A 51.4M-parameter GPT built and trained from scratch — the embedding table, attention, training loop, mixed precision, all by hand — then pretrained on ~3.5B tokens and fine-tuned into a small science-QA chat model. Everything is in JAX / Flax NNX.

▶ Interactive overview / deck — a visual walkthrough of the architecture, training, and results.

The pipeline

Pretrain (3.5B tokens)  →  SFT (chat + science QA)  →  Evaluate (SciQ accuracy)

1. Pretrain — next-token prediction on a streamed mixture (70% FineWeb-Edu, 20% Cosmopedia, 10% FineMath). Produces a base LM that continues text (final val loss ≈ 3.17).
2. SFT — fine-tune the base into a chat model on a blend of science QA (SciQ + OpenBookQA + ARC) and a little chat (smol-smoltalk), in a chat template, computing the loss on assistant tokens only.
3. Evaluate — SciQ test accuracy by multiple-choice likelihood ranking: 47.3% on the 1,000-question test split vs a 25% random baseline (~1.9× chance).

The tiny-titan.ipynb notebook is the Colab driver that runs all three stages top-to-bottom (it calls the code in src/; nothing is redefined in the notebook).

The model

A modern decoder-only transformer (Gemma-style): RMSNorm, RoPE, SwiGLU, QK-norm, pre-norm blocks, final norm, and a tied output projection.

Parameters	51.4M
Layers / heads	8 / 8
Model dim / FFN	512 / 1408 (SwiGLU)
Context	1024
Vocab	50,257 (GPT-2 BPE)
Precision	bf16 mixed

The design decisions are recorded as ADRs in adr/.

Setup

Needs Python 3.11 and uv.

uv sync

This creates .venv and installs everything from pyproject.toml; afterwards uv run … uses that environment automatically.

Run

Chat with the model (terminal)

After fine-tuning, talk to the chat model — it applies the chat template and stops each reply at the model's turn-end token. Each message is atomic: no history carries between turns, so a reply depends only on the current question.

uv run python -m src.inference.main --chat -w weights/tiny-titan-instruct-best.orbax --greedy -r 1.2 -n 64

KV cache (on by default). Generation reuses a key/value cache so each new token doesn't re-process the whole context — much faster. Add --time to print tokens/sec, and --no-cache to compare against re-running the full model each step:

# fast path (default), timed
uv run python -m src.inference.main --chat -w weights/tiny-titan-instruct-best.orbax --greedy -r 1.2 -n 64 --time

# slow path: full re-run per step (sliding window, no maxlen cap)
uv run python -m src.inference.main --chat -w weights/tiny-titan-instruct-best.orbax --greedy -r 1.2 -n 64 --time --no-cache

Note: --time's tokens/sec is only meaningful over many generated tokens — short replies are dominated by prompt prefill, so the rate looks low.

Recommended settings. It's a science-QA model, so favor its best guess over variety:

• Factual QA (recommended): --greedy — deterministic, takes the single most-likely answer. This is what the benchmark rewards and what reads best for short answers.
• A little variety: drop --greedy and use -t 0.3 -p 0.9 (low temperature
  • nucleus sampling). Higher temperatures make a 51M model pick wrong-but-plausible answers, so keep it low for facts.
• -r/--repetition-penalty — defaults to 1.2 (good for long generation); for short QA answers -r 1.0 (off) avoids over-penalizing and is recommended.
• -n/--max-new-tokens — answers are short, so -n 64 is plenty.

Prompts to try — real questions from SciQ (its home turf), each with the expected answer so you can see how it's doing:

What is the opposite of melting? (freezing)

When a meteoroid reaches Earth, what is the remaining object called? (meteorite)

What is the common word for potential difference in a circuit? (voltage)

Dialysis is a treatment for failure of what organs? (kidneys)

What galaxy is our solar system a part of? (milky way)

Open-ended chat ("how are you?", "tell me a story") is weak — at 51M the model learned to answer, not to converse. Quiz it; don't chat with it.

Or use the base model as a raw text-continuation engine (each prompt independent):

uv run python -m src.inference.main -w weights/tiny-titan-base/tiny-titan.orbax

Full options: -w/--weights, -t/--temperature, -p/--top-p, -r/--repetition-penalty, -n/--max-new-tokens, --seed, --greedy, --chat, --cache/--no-cache, --time. Type exit, quit, or Ctrl-D to leave.

Train (Colab)

Open tiny-titan.ipynb in Colab on a GPU runtime and run it top-to-bottom: setup → pretrain → SFT → evaluate. Checkpoints go to Google Drive and resume automatically, so a dead session is fine — just re-run the cell.

Evaluate

uv run python -m src.eval.sciq --weights weights/tiny-titan-instruct.orbax

Reports SciQ test accuracy (length-normalized + raw-sum) against the 25% random baseline, plus a few sample answers.

Layout

tiny-titan.ipynb          # Colab driver: pretrain -> SFT -> eval
src/
  model/                  # TinyTitan architecture + build_model()
  training/
    build_dataset.py      # stream + pack the pretrain corpus -> train/val .bin
    build_sft_dataset.py  # build SFT data -> token .bin + parallel loss-mask .bin
    data.py               # TokenLoader (base) + MaskedTokenLoader (SFT)
    train_base.py         # base pretraining loop (checkpoint/resume)
    train_sft.py          # SFT loop: init from base, masked loss, best-val ckpt
  inference/
    generate.py           # autoregressive loop, sampling, generate_chat
    streamer.py           # byte-safe incremental decode
    loader.py             # restore model from an Orbax checkpoint
    main.py               # Click CLI: continuation + --chat
  eval/
    sciq.py               # SciQ multiple-choice likelihood-ranking scorer
adr/                      # architecture decision records (why each piece exists)

Docs

• adr/ — architecture & training decisions, one short report each.
• src/inference/README.md — how generation works (decoding, temperature, top-p).
• src/eval/README.md — how evaluation works (likelihood ranking, metrics).

Honest scope

At ~51M parameters this is a toy-scale model. The win is the full from-scratch pipeline and a chat model that answers simple science questions above chance — not fluent open-ended conversation, which needs far more scale. SFT teaches format and surfaces what pretraining learned; it can't add knowledge the base never saw.