nanochat-jax — checkpoints
Model checkpoints for nanochat-jax, a JAX/Flax port of Andrej Karpathy's nanochat that reproduces the reference recipes on spot TPU (e.g. v6e-8) instead of 8×H100, matching the reference model quality.
These are the original nanochat-format checkpoints (.pt + meta_*.json + tokenizer),
loaded by cloning the nanochat-jax repo and pointing NANOCHAT_JAX_BASE_DIR at a release
folder — not transformers.from_pretrained.
Layout
Each release is a self-contained, directly-loadable base directory:
<release>/<depth>/ ← point NANOCHAT_JAX_BASE_DIR here
tokenizer/
base_checkpoints/<tag>/ model_<step>.pt, meta_<step>.json, optim_<step>_rank0.pt
chatsft_checkpoints/<tag>/ model_<step>.pt, meta_<step>.json, optim_<step>_rank0.pt
Releases
| Release | Depth | Base CORE | ChatCORE (SFT) | Hardware | Notes |
|---|---|---|---|---|---|
0098-speedrun/d24 ⭐ recommended |
24 (~1.4B) | 0.2695 | 0.3733 | TPU v6e-8 (spot) | Full speedrun (tokenizer→base→SFT), $60.8 spot. The run featured in the write-up. |
(More releases — e.g. the canonical 0.274 base — may be added as sibling folders.)
How to load
git clone https://github.com/tucan9389/nanochat-jax.git && cd nanochat-jax
# install per the repo README (uv + jax[tpu]/cpu, torch cpu)
# download ONE release (scope with --include so you don't pull every release)
hf download tucan9389/nanochat-jax --include "0098-speedrun/d24/*" --local-dir ./nc
BASE=$PWD/nc/0098-speedrun/d24
# NB: nanochat-jax reads the tokenizer/eval bundle from the env var and checkpoints from
# --base-dir, so point BOTH at the same release folder.
# chat with the SFT model (CPU works for inference)
JAX_PLATFORMS=cpu NANOCHAT_JAX_BASE_DIR=$BASE PYTHONPATH=$PWD \
python scripts/chat_cli.py -i sft -g d24_speedrun_r4_sft -t 0.0 -p "How do I make coffee?"
# or evaluate the base model
NANOCHAT_JAX_BASE_DIR=$BASE \
python -m scripts.base_eval --source base --model-tag d24_speedrun_r4 --eval core --bf16
Optimizer states (optim_*.pt, ~7.7 GB each) are included so you can resume or branch
training from a checkpoint. If you only want inference/eval, skip them:
--include "0098-speedrun/d24/*" --exclude "*optim*".
⚠️ These are pickle-based
torch.savefiles (loaded withweights_only=False). Only load checkpoints you trust.
License & attribution
MIT, inheriting from nanochat (© 2025 Andrej Karpathy). Community reproduction, not affiliated with or endorsed by Karpathy. Trained on TPU Research Cloud (TRC) credits. Training data (ClimbMix, SmolTalk, etc.) carries its own terms.
Links
- Port: https://github.com/tucan9389/nanochat-jax
- Upstream nanochat: https://github.com/karpathy/nanochat
- Write-up: (blog link — add when published)