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 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).

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.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.

@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}
}