nanochat-jax

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

(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.save files (loaded with weights_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)