cataluna84/nanogpt-jax-181m
nanoGPT-JAX 181M — base + SFT
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:
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
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
ShardingRulestable 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).
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:
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).
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.
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.
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.pyRuns 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.
@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}
}