A 181M-parameter GPT trained from scratch in pure JAX — no Flax, Equinox, or Keras. Every layer, the training loop, sharding, KV-cache inference, and int8 quantization are written on top of a deliberately minimal two-class abstraction.

Trained on Google TRC TPU. Code: https://github.com/cataluna84/llm-architectures

Two checkpoints are included:

Directory	What it is	Best val loss
`base/params/`	10k-step pretraining on FineWeb10B	3.1271
`sft/params/`	one epoch of instruction tuning on top of base	1.4365

Read the evaluation section before using this. At 181M parameters and 5.2B training tokens, this model is at chance on every multiple-choice benchmark tested. It is a working reference implementation and a study of training dynamics, not a capable assistant.

Architecture


Parameters	181,075,968
Layers	16
Model dim	768
Query heads / KV heads	8 / 4 (GQA)
Head dim	96
MLP hidden dim	3072
Sequence length	2048
Vocab	50,304 (GPT-2 BPE + 9 chat tokens)
Precision	bfloat16
Attention pattern	`L` — every layer global

Design choices worth naming:

Grouped-query attention (8 Q / 4 KV) with RoPE, QK-norm, and logits soft-capping.
ReLU² MLP and parameter-free RMSNorm (no learnable scale).
Attention weights are deliberately 3-D — wq/wk/wv: (d_emb, heads, head_dim), wo: (heads, head_dim, d_emb) — so head sharding is trivial. Anything assuming 2-D matrices (notably Muon's dimension numbers) must be told which axes are batch/reduction/output.
Params carry logical axis names; a ShardingRules table maps them to physical mesh axes, so changing parallelism means editing the rules, not the layers.
During SFT, RoPE frequencies are computed from packed segment positions, not absolute positions.

Tokenizer

GPT-2 BPE extended with nine special tokens:

<|pad|>
<|user_start|>   <|user_end|>
<|assistant_start|>   <|assistant_end|>
<|system_start|>   <|system_end|>
<|tool_start|>   <|tool_end|>

Chat format (what the SFT checkpoint expects):

<|endoftext|><|user_start|>{prompt}<|user_end|>
<|assistant_start|>

Generation stops at <|assistant_end|>.

Pretraining — `base/params/`

FineWeb10B (kjj0/fineweb10B-gpt2), 524,288 tokens/step × 10,000 steps = 5.24B tokens, on a v5e-32 (8 hosts × 4 chips, DDP).

Metric	Value
Best val loss	3.1271 @ step 9919
MFU	25.06%
Throughput	1,137,868 tok/s
Step time p50 / p90 / p99	0.4605 / 0.4619 / 0.4630 s
Wall clock	97.5 min
Total FLOPs	5.696 × 10¹⁸
HBM peak	2.17 / 15.75 GiB
Loss spikes / divergence	0 / none

Optimizer recipe

Muon on hidden matrices, AdamW on embeddings and the unembedding:

peak LR 0.02 | momentum warmup 0.85 → 0.95 over 300 steps | momentum max 0.95
ns_steps 5 | mu_dtype float32 | grad clip 0.5
embedding LR 0.3 | unembedding LR 0.002 | Adam betas 0.8/0.95
cautious weight decay 0.2 | plain weight decay 0.0
WSD schedule, warmdown fraction 0.65 | LR warmup min(300, 1% of steps)

This came from a 43-run sweep on a v5e-64 against a 2σ adoption bar (3-seed noise floor σ = 0.0016; baseline 3.5961). What actually moved:

Knob	Default → adopted	Δ/σ
cautious weight decay	0.01 → 0.2	−14.1σ
embedding LR	0.2 → 0.3	−4.3σ
unembedding LR	0.004 → 0.002	−4.3σ
grad clip	1.0 → 0.5	−3.3σ
all four composed	—	−18σ (3.5671, 3 seeds)

Cautious weight decay was the largest mover and monotone over 0.01 → 0.1 → 0.2 with 0.4 regressing, so 0.2 is an interior optimum. The composed config beats the best single knob by ~3.9σ, so the gains compose rather than overlap.

Defaults stood on gradient accumulation, Adam betas, momentum max, mu_dtype, ns_steps, and LR warmup — extending warmup at a fixed horizon actively hurt.

Schedule, stated carefully: at a matched 2,500-step horizon, WSD-0.65 (3.3230) beats cosine (3.3568) by −0.0338, ≈ −21σ — the largest single effect measured. Sweep tables that compare 2,500-step arms against a 1,000-step baseline show ~−169σ; that number is dominated by training 2.5× longer and should not be read as a schedule effect.

SFT — `sft/params/`

Warm-started from base/params/, one epoch over a packed mixture of smoltalk + MMLU + GSM8K with completion-only loss, on a v5e-64 (16 hosts × 4 chips).

Metric	Value
Best val loss	1.4365 @ step 800 (from 1.596)
Steps	844 — one epoch, data-bound
Tokens	442,499,072
MFU / throughput	26.5% / 2,405,515 tok/s
Wall clock	6.3 min

Data: 567,195 packed rows / 447.6M tokens (smoltalk 459,880 rows, MMLU 99,842, GSM8K 7,473). Sequences are BestFit-packed with segment_ids and completion masks. Optimizer is AdamW (not Muon), init_lr_frac 0.2, final_lr_frac 0.0.

Every SFT hyperparameter is a nanochat default and unswept. Treat the result as a first data point, not a tuned configuration.

Evaluation — read this before drawing conclusions

200 examples/task, greedy decoding, identical harness and settings on both sides.

Task	Base	SFT	Δ	Chance
MMLU	0.230	0.210	−0.020	0.25
ARC-Easy	0.245	0.200	−0.045	~0.25
ARC-Challenge	0.285	0.260	−0.025	~0.25
GSM8K	0.000	0.000	0.000	~0

Both models are at chance, and every delta is inside noise. Standard error at n=200 near chance is ±0.031, so −0.020 to −0.045 are all under ~1.5σ. Reading "SFT made the model worse" into these numbers would be over-interpretation of four coin flips.

This is the expected outcome at this scale. A 181M-parameter model trained on 5.2B tokens (~29 tokens/param) has no latent knowledge for instruction tuning to surface, and SFT teaches format and turn-taking, which accuracy benchmarks structurally do not measure.

What SFT demonstrably did change is the loss (val 1.596 → 1.4365) and the generation behaviour:

Prompt <|endoftext|><|user_start|>What are the benefits of regular exercise? Your response should contain at least 3 sentences. Include keywords such as "health", "reduce", and "improve".<|user_end|> <|assistant_start|>

Completion Regular exercise offers numerous health benefits, particularly improved cardiovascular health and a longer lifespan. Research has shown that regular physical activity can improve cognitive function, enhance mood, and contribute to overall well-being. This could be particularly beneficial for individuals with chronic diseases or conditions that affect daily activities. Additionally, regular exercise can have a positive impact on mental health, reducing symptoms of depression and anxiety.<|assistant_end|>

Three-plus sentences as instructed, all three requested keywords present, and a clean stop on <|assistant_end|> rather than running on. That is the evidence the benchmark table cannot provide.

HumanEval was deliberately excluded from the harness: its task module imports a nonexistent package and executes untrusted model-generated code.

Usage

Checkpoints are Orbax directories, not transformers weights — load them with the repo's code rather than AutoModel.

bash

git clone https://github.com/cataluna84/llm-architectures
cd llm-architectures
uv sync

huggingface-cli download cataluna84/nanogpt-jax-181m --local-dir ./ckpts

# instruction-tuned, chat-formatted, stops on <|assistant_end|>
NANOGPT_MODEL_TYPE=SFT \
NANOGPT_LOAD_PARAMS_CKPT_PATH=./ckpts/sft/params \
python nanogpt/inference.py

# base model, raw continuation
NANOGPT_MODEL_TYPE=pretrained \
NANOGPT_LOAD_PARAMS_CKPT_PATH=./ckpts/base/params \
python nanogpt/inference.py

Runs on CPU — 181M params needs no accelerator for sampling. Inference uses left-padded prompts with right-aligned generation through a KV cache.

Known numerics caveat: exact greedy equivalence between the KV-cache path and the no-cache path holds only on a compact active-KV slice, not the full masked buffer — bf16 flash-attention tiling noise breaks ties.

Limitations

At chance on knowledge benchmarks. Not useful for factual QA, reasoning, or arithmetic. GSM8K is 0.000.
Trained on 5.24B tokens of FineWeb, a filtered web crawl. It reproduces the biases and factual errors of that data and has had no safety tuning, RLHF, or content filtering of any kind.
SFT used a small public mixture for one epoch with unswept hyperparameters.
English only. 2048-token context.
Research and educational artifact. Do not deploy it anywhere its output is trusted.

Citation

Built on the pure-JAX nanoGPT design; see the repository for full attribution and the LICENSE.

bibtex

@software{nanogpt_jax_181m,
  author = {cataluna84},
  title  = {nanoGPT-JAX 181M: pure-JAX GPT pretraining and SFT on TPU},
  url    = {https://github.com/cataluna84/llm-architectures},
  year   = {2026}
}

cataluna84/nanogpt-jax-181m

nanoGPT-JAX 181M — base + SFT