Porting nanochat to a TPU Achieves Quality Parity but Faces Performance Gap

Porting nanochat to a TPU: what carries over from PyTorch, and what breaks #1

Karpathy's nanochat normally runs on an 8×H100 GPU node, and several ports of it to JAX already exist. Among them, my aim was to keep the config and architecture as close to nanochat as possible ( parity) while catching up on both model quality and training performance — the quality (its CORE score) reproduced cleanly, and the performance came only partway. If nanochat is new to you: it's Karpathy's full-stack LLM project — tokenizer training, pretraining, SFT, and RL in a single repository — where about four hours and roughly $100 on one 8×H100 node gets you your own chatbot (nanochat is often called the "$100 speedrun"; those numbers are for d20). For reference, stopping at the GPT-2-grade base model (d24) takes about two hours and roughly $48 on the same node. This post is a record of that port: what carried over unchanged from PyTorch, and what broke on the TPU.

1. Reproduction results

nanochat-jax currently provides a speedrun.sh script that covers base model training and SFT (upstream nanochat goes all the way to RL; this reproduction stops at SFT). The claim above — that the quality reproduced — is based on the CORE score. CORE is the average of accuracies over 22 evaluation tasks, each rescaled so that random guessing scores 0 and a perfect score is 1. Within the same evaluation harness, it lets you compare models against each other, which is why nanochat uses it to judge what counts as "GPT-2 grade". What we reproduced is recipe 4 from the official nanochat LEADERBOARD (R4 from here on; it's also called d24, since depth 24 ≈ 1.4B parameters). The band (0.2512–0.2677) is the score distribution from Karpathy running the same R4 recipe 7 times; this run's 0.2695 lands just above it.

On performance, there's still a gap. MFU is about 24% (d24) — half of Karpathy's measured H100 numbers (47–48% at d20).

Below, we run each script of speedrun.sh on a TPU v6e-8 and check whether the quality Karpathy reported actually comes out.

Model	Chip	CORE (base)	Train time	Train cost	Total cost
GPT-2 (2019, 1.5B)	TPU v3 x32 (est.)	0.2565	168 h	~$43,000 (est.)	-
nanochat R4 (d24)	H100 x8	0.2571	~2 h	~$48 (est.)	-
nanochat-jax (ours)	TPU v6e-8 (spot)	0.2695	5.29 h (6.02 w)	$30 ($130 od)	$60.8 (~$263 od)

CORE scores and the band are from the nanochat LEADERBOARD. w = wall clock including checkpointing and compilation; od = on-demand list price.1

If you spot anything wrong or unclear in this post, or have a question, please let me know at tucan.dev@gmail.com or in the comments — feedback is always welcome :)

---

2. TPU basics

This run used 8 v6e chips (a single-host slice). Per-chip specs across generations:

	v5p	v6e (Trillium)	7x (Ironwood)
HBM capacity	95 GB	32 GB	192 GB
HBM bandwidth	2,765 GB/s	1,638 GB/s	7,380 GB/s
bf16 compute	459 TFLOPs	918 TFLOPs	2,307 TFLOPs
Low-precision compute	Int8 918 TOPs	Int8 1,836 TOPs	FP8 4,614 TFLOPs
MXU (matrix-multiply unit)	128×128	256×256	256×256
Chips per host	4	8	4
1-host topology	2×2×1	2×4	2×2×1
Max pod	8,960 chips	256 chips	9,216 chips

Sources: Cloud TPU System Architecture, v6e, tpu7x (Ironwood). bf16, Int8, and HBM numbers are from each generation's official spec sheet. Low-precision compute is roughly 2× bf16; v5p and v6e hardware-accelerate only up to Int8, and native FP8 support starts with Ironwood.

The v6e's MXU grew to 256×256, from the 128×128 of every generation up to v5p — if a tensor dimension isn't a multiple of 256, XLA pads it with zeros and part of the unit is wasted (more on this in insight 5 of section 4). Meanwhile, HBM is 32GB per chip — a third of v5p's 95GB — while compute is 2×: a compute-heavy, memory-lean design.

Our v6e-8 slice adds up to 7.34 bf16 PFLOPS and costs $4–5/hour on spot (us-central1). Spot is heavily discounted against on-demand, in exchange for GCP being able to reclaim (preempt) it at any time. This run cost $60.8 on spot over 12.19 hours total ($263 at on-demand rates), with one preemption and recovery along the way. You're billed for as long as the node exists (training or not) — forget to delete it, and at $5/h for an 8-chip slice, a day is ≈ $120.

Figure 1. The JAX AI Stack — hardware (CPU/GPU/ TPU) → the XLA compiler → the JAX core → the library layer on top ( Flax and others). (Source: jaxstack.ai, © JAX team)

Of this stack, nanochat-jax uses three software layers — XLA, JAX, and Flax — plus Pallas, which isn't in the figure.

Role	nanochat-jax	PyTorch equivalent
Neural network modules	Flax NNX	torch.nn
Computation + compilation	JAX + XLA (jit·grad·vmap)	eager + autograd + torch.compile
Custom kernel	Pallas (Splash Attention)	Triton
Optimizer	JAX (own implementation, no torch dependency)	torch.optim
Data loading	NumPy (no torch dependency)	torch.utils.data
Checkpoint	PyTorch .pt (for nanochat compatibility)	PyTorch .pt

---

3. Speedrun verification — from the tokenizer to the report card

The pipeline runs tokenizer (5.3m) → base (6.02h + eval 44.5m) → SFT (68.7m + eval ~3.5h), and we run the scripts that speedrun.sh executes, one stage at a time (to run everything in one shot, see Appendix C). For a deep understanding of each stage, see Karpathy's walkthrough post (his is d20, ours is d24); this post focuses on quickly checking measured values against the reference numbers. The midtraining from the walkthrough era (an intermediate stage that pre-taught conversation format and tool use) no longer exists as a separate stage upstream (its data was folded into the SFT mixture), so we also go straight from base to SFT. Setup starts in Appendix A.

Figure 2. The nanochat speedrun pipeline — this reproduction runs tokenizer, base, and SFT (the solid boxes) and skips RL (dashed).

Step 1. Tokenizer (5.3m)

python -m nanochat_jax.dataset -n 170 # ClimbMix train 170 + val 1 shards
python -m scripts.tok_train --max-chars 100000000 --doc-cap 10000 --vocab-size 32768
python -m scripts.tok_eval

We train a vocab-327682 tokenizer on 100M characters of ClimbMix, then compare against GPT-2 how many tokens the same text takes. On the training data (the train row), ours uses 1.5% fewer tokens than the GPT-2 tokenizer. Fewer tokens means more text learned for the same budget — a compression advantage — and Karpathy's tokenizer shows the same pattern, a signal that the reproduction is on track (for Korean, the training data contains no Hangul, so ours spends about 2× the tokens of GPT-4; the full per-domain table is in Appendix D). The number to check: train +1.5%.

Step 2. Base model (6.02h + 44.5m)

# --recipe through --use-real-data: Karpathy's R4 config / --attn-impl through --splash-*: v6e-8 TPU-specific
python -m scripts.base_train \
 --recipe=324e69c --depth=24 --seq-len=2048 --vocab-size=32768 \
 --target-param-data-ratio=9.5 --total-batch-size=1048576 \
 --device-batch-size=2 --grad-accum-steps=32 --grad-accum-impl=fused \
 --warmup-steps=0 --warmdown-ratio=0.5 --final-lr-frac=0.0 \
 --weight-decay=0.2 --matrix-lr=0.02 --embedding-lr=0.3 \
 --unembedding-lr=0.004 --scalar-lr=0.5 \
 --bf16 --cast-embeddings-bf16 --use-real-data \
 --attn-impl=splash --splash-block-q=512 --splash-block-kv=512 --splash-block-kv-compute=256 \
 --matmul-precision=default --lm-head-precision=highest --ve-grad-impl=onehot \
 --checkpoint-every=200 --keep-last-checkpoints=2 \
 --model-tag=d24_speedrun_r4 --no-final-eval

The top half of these arguments is the R4 recipe as-is (730M scaling parameters × 9.5 ≈ 6.9B tokens)3; the only ones you may need to touch are the --splash-* kernel block sizes.4

[step 100] loss=4.517054 ... mfu=24.5%
[step 1000] loss=2.786205 ... mfu=24.7%

Against Karpathy R4's 2-hour base train loop (8×H100), ours takes 5.29h (6.02h wall clock including checkpointing and compilation), at about 24% MFU5.6 (The two lines above are an excerpt; I've uploaded the full training log of this run — the header also prints flops/token and peak TFLOPS.)

RESULTS_DIR=$HOME/.cache/nanochat-jax/results/d24_speedrun_r4 && mkdir -p $RESULTS_DIR

python -m scripts.base_eval --source base --model-tag d24_speedrun_r4 \
 --eval bpb,core --eval-steps 20 --device-batch-size 1 --max-per-task -1 \
 --out $RESULTS_DIR/base_eval.json --partial-out $RESULTS_DIR/base_eval_core_partial.json --bf16

CORE (the 22-task metric from section 1) comes out to 0.2695 — above GPT-2's 0.2565, and in the same class as the R4 band (0.2512–0.2677). The val bpb of 0.7343 is a reference number only, since tokenizer differences mix into it (R4: 0.7185).

Figure 3. The dashed lines in the middle and right panels are the R4 baselines — by the end of training, both metrics reach the baselines, landing at R4 level (the curves are from an earlier reproduction run).

Step 3. SFT (68.7m + eval ~3.5h)

# SFT training (68.7m)
python -m scripts.chat_sft --model-tag d24_speedrun_r4 --output-model-tag d24_speedrun_r4_sft \
 --num-iterations -1 --device-batch-size 1 --max-seq-len 2048 --total-batch-size 524288 \
 --sft-scope full --eval-every 200 --eval-steps 4 \
 --chatcore-every -1 --chatcore-max-cat -1 --chatcore-max-sample 24 \
 --load-optimizer 1 --log-every 10 --bf16
# scoring — the 3 multiple-choice tasks in one go, the 3 generative tasks one by one
python -m scripts.chat_eval -i sft -g d24_speedrun_r4_sft --max-new-tokens 512 --batch-size 8 \
 --dtype bfloat16 --out $RESULTS_DIR/chat_eval_sft.json --task-name "ARC-Easy|ARC-Challenge|MMLU"
for task in SpellingBee HumanEval GSM8K; do
 python -m scripts.chat_eval -i sft -g d24_speedrun_r4_sft --task-name "$task" \
 --max-new-tokens 512 --dtype bfloat16 --num-samples 1 --temperature 0 --top-k 50 \
 --jit-gen 1 --out $RESULTS_DIR/chat_eval_$(echo "$task" | tr 'A-Z' 'a-z').json
done

Train for 68.7 minutes on a mixture (~1.07M rows) of SmolTalk, MMLU, GSM8K, SpellingBee7, and 1,000 identity conversations8, and the base model that could only repeat the question back becomes a model that answers it. ChatCORE (the centered average over 6 chat tasks) is 0.3733 — the comparison line here is the base model itself. On the same 6 tasks, base scored essentially 0 on the generative ones, so most of the generative score is ability SFT added: SpellingBee 0.9961, GSM8K 0.1008 (math is RL's job). I don't compare 1:1 against the report card in Karpathy's walkthrough — that one is d20 and ran on the older midtraining pipeline.

Figure 4. Right panel: SFT (blue) is ahead of base (gray) on every task — you can see the generative tasks, which scored 0 on base, now score above zero.

Beyond the scores, checking with actual prompts:

JAX_PLATFORMS=cpu NANOCHAT_JAX_BASE_DIR= PYTHONPATH= \
python scripts/chat_cli.py -i sft -t 0.0 -p "How do I make coffee?"

Feed the same prompt to base and to SFT, and the difference is easy to see.

Prompt: How do I make coffee?
 BASE : How do I make coffee? How do I make coffee? ... (repeats the question)
 SFT : Making coffee is a simple process... you'll need a coffee maker, ...

Prompt: Spell the word 'banana' letter by letter.
 BASE : The word 'banana' is made up of two letters. ... (fails to spell)
 SFT : banana:b,a,n,a,n,a

Report card

python -m nanochat_jax.report generate

| Metric | BASE | SFT | RL |
|-----------------|----------|----------|----|
| CORE metric | 0.269486 | - | - |
| ARC-Challenge | - | 0.509386 | - |
| ARC-Easy | - | 0.622896 | - |
| GSM8K | - | 0.100834 | - |
| HumanEval | - | 0.140244 | - |
| MMLU | - | 0.369534 | - |
| SpellingBee | - | 0.996094 | - |
| train bpb | 0.731504 | - | - |
| val bpb | 0.734295 | - | - |
| ChatCORE metric | - | 0.373265 | - |

End to end: 12h14m wall clock (12.19h billed), $60.8 on spot (~$263 on-demand), one preemption with a 17-minute recovery (RL is out of scope, so its column stays empty — kept to match nanochat's format). The CORE reproduction is done. The next two sections cover the missteps on the way there, and what they cost.

---

4. Five things I learned from the TPU/JAX port

The most important thing the port taught me: if you can avoid it, don't. I decided to do it anyway, so here are the problems I ran into and some insights of my own, trimmed down to five. Each item is tagged with which side it hurt: quality (CORE) or performance (MFU). The first is the bug that fooled us the longest.

Insight 1. [Quality] Some bugs give wrong results without raising an error — data corruption from two lines of NumPy

When quality collapsed, the first thing I suspected was the chip, and the second was the bf16 embeddings. But what had quietly gone wrong was not the model — it was how the training loop fed the data: no error, no NaN, no shape mismatch, the kind of bug that shows up only as lower quality. Below is the code in base_train.py that collects the grad-accum microbatches.

idx_chunks, target_chunks = [], []
for _ in range(grad_accum_steps):
 inputs_np, targets_np = next(train_loader)
 idx_chunks.append(np.asarray(inputs_np)) # ← collects the views as-is (the bug)
 target_chunks.append(np.asarray(targets_np))
idx_all = jnp.asarray(np.stack(idx_chunks, axis=0)) # stacked later
targets_all = jnp.asarray(np.stack(target_chunks, axis=0))

Follow train_loader into dataloader.py, and you find row_buffer (line 112), which reuses a single buffer for speed and hands out views (zero-copy slices) of it. But the path above collects the microbatches in a list and calls np.stack later. Since np.asarray doesn't copy something that is already an array,9 every collected microbatch ends up pointing at the last contents of the reused buffer — the batch for one step effectively becomes N copies of the last microbatch. That is data corruption: the token count stays the same, only the contents are wrong.

The original PyTorch code draws one microbatch, immediately runs forward/backward to accumulate it into the gradients, then draws the next — so the data is consumed before the buffer is reused, and the trap doesn't exist there. We, on the other hand, collect all the microbatches and hand them to a single compiled step (lax.scan) — and while they were being collected, the buffer got overwritten (why we needed this new path in the first place is explained by insight 2 below). The fix is one line: copy at capture time.

 idx_chunks.append(np.array(inputs_np, copy=True)) # copy at capture time
 target_chunks.append(np.array(targets_np, copy=True))

The v6e runs from the period with this bug were stuck in a CORE band of 0.08–0.13; the run right after the fix scored 0.274 (that 0.274 is the published canonical run, separate from the speedrun's 0.2695 — both are R4-level). Also, the fix is already in main, so the commands you'd run from the speedrun section above are safe. (This data-feeding code actually had one more "sibling bug".10)

If a dataloader can hand out views of a reused buffer, then on any path that stacks them later, the safe first move is to check whether they get copied.

Insight 2. [Performance] It's better to merge the training step into a single jit

compute_grads = jax.jit(compute_grads) # compile unit 1
accumulate = jax.jit(accumulate) # compile unit 2
apply_update = jax.jit(apply_update) # compile unit 3

acc = zero_grads
for mb in microbatches: # DON'T — three programs chained by a Python loop
 acc = accumulate(acc, compute_grads(state, mb))
state = apply_update(state, acc)

@jax.jit # DO — the whole step is one program
def train_step(state, microbatches):
 def micro_step(acc, mb): # same logic as above — only the boundary moves inside
 return accumulate(acc, compute_grads(state, mb)), None
 acc = lax.scan(micro_step, zero_grads, microbatches)[0]
 return apply_update(state, acc)

Port the PyTorch code as-is and you get the DON'T above. On GPU that structure is fine — asynchronous CUDA and torch.compile hide the host-side dispatch cost. On TPU, an A/B of the two structures moves MFU from 16% to 22%.11

The cause is the boundaries. XLA optimizes only within one jit program, so splitting the step into three programs leaves the device idle between programs (in our first implementation, idle time reached 26.7% of the step), and the compiler can't overlap operations across them either. The MaxText docs call wrapping the whole step in a single jit the "primary mechanism" behind their performance. Meanwhile, the first (manual) method in the official optax grad-accum example shows this structure with no warning, so it's easy to copy if you miss the scan alternative further down the page.

Rule: one compiled program per step. When performance is low, look at device idle time before you look at kernels (it shows in the XProf timeline) — if the idle time stands out, suspect the structure (the jit boundaries) first.

Insight 3. [Performance] The Splash attention kernel's defaults need tuning

# jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py
# TODO(apaszke,sharadmv): Select better parameters based on a heuristic.
block_q: int = 128 # TRAP — all 8 tile fields default to 128
block_kv: int = 128
block_q_dkv: int = 128 # ...5 of them are backward-only

python scripts/base_train.py ... \
 --splash-block-q=512 --splash-block-kv=512 --splash-block-kv-compute=256 # DO — the values we use at seq 2048

Splash is the TPU flash-attention kernel that ships inside JAX. Eight tile-size fields decide its performance — open the source and the defaults are all 128, with the authors' TODO still sitting right above them.

In an A/B that changed only the tiles, MFU went from 22.4% to 41.2%,11 and the loss matched to six decimal places — exactly what the docstring says: tiles barely touch the numerics and mostly move performance. Most of the runtime sat in the backward pass. Why flash-attention's backward costs more than its forward is covered in the FlashAttention-2 paper; the trade-offs of Splash tiles are in this write-up. MaxText also uses its own 512 instead of the JAX defaults.

The best values depend on the sequence length — in our seq-4096 experiments the optimum was 1024, and the published run (seq 2048) uses 512. The sweep from tpuchat, an independent port, converged to the same 1024 for its shapes — and 2048 tiles were actually slower. Bigger is not automatically better.

Since the loss isn't sensitive to the tiles (see the docstring), a tile sweep even before profiling is a reasonable first move.

Insight 4. [Quality] By default, TPUs compute fp32 matmuls in bf16

Our training command carries a flag called --lm-head-precision highest. For a while I believed it was protecting quality; in reality, it does nothing. How that happened is tied to a TPU-specific behavior.

Request an fp32 matmul on a TPU, and even though the dtype stays fp32, under the default settings the multiplication runs in bf16. This is not hidden behavior — it's a documented default — and it's a trap well known enough that a Google researcher wrote in an issue that his own team had been bitten by it for roughly the fourth time.

a = np.random.default_rng(0).standard_normal((1024, 1024)).astype(np.float32)
ref = a.astype(np.float64) @ a.T.astype(np.float64)
out = jax.jit(jnp.dot)(a, a.T) # fp32 inputs, default precision
assert np.abs(np.asarray(out, np.float64) - ref).max() / np.abs(ref).max() \
 --zone= --accelerator-type=v6e-8 --version=v6e-ubuntu-2404 --spot

Once READY appears, you can connect. It normally shows up within a few minutes, and you'll know about a failure within a minute.

From READY on, you're being billed ~$5/hour. Connect right away. (a) is an interactive shell; (b) and (c) are for one-shot commands and background runs.

A new project may have a TPU quota of 0. Check and request v6e quota for your target region on the console's Quotas page first — as Appendix B shows, quota was the #1 cause of failed allocations.

# (a) interactive shell — for quick commands
gcloud compute tpus tpu-vm ssh --zone=

# (b) one-shot command — put multi-line scripts in the REMOTE variable via a quoted heredoc
REMOTE=$(cat 
EOF
)
gcloud compute tpus tpu-vm ssh --zone= --ssh-flag="-oConnectTimeout=30" --command="$REMOTE"

# (c) long-running command — detach from the session with setsid nohup and run in the background
REMOTE=$(cat ' > ~/logs/command0.log 2>&1 &
echo LAUNCHED
EOF
)
gcloud compute tpus tpu-vm ssh --zone= --ssh-flag="-oConnectTimeout=30" --command="$REMOTE"

Once you're on the machine, clone the nanochat-jax repository and install the dependencies.

sudo rm -f /tmp/libtpu_lockfile # required before the first jax process
git clone https://github.com/tucan9389/nanochat-jax.git
curl -LsSf https://astral.sh/uv/install.sh | sh
uv venv --python 3.12 ~/venv && source ~/venv/bin/activate
uv pip install "torch~=2.11.0" --index-url https://download.pytorch.org/whl/cpu
cd ~/nanochat-jax && uv pip install -e ".[tpu,dev]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
python -c "import jax; print(jax.devices())"

If you see eight TPU v6 lite devices, setup is done.

Appendix B) History and statistics of trying to allocate TPU spot instances

TPU spot attempt statistics by region, April–July 2026. This is an excerpt of the main chip–region combinations, so the rows add up to less than the full statistics (2,885 creation attempts, 1,632 failures).

Chip	Region	Attempts→successes	Failure causes	Slice price (spot, May/Jul)	Actual usage
v6e-8	us-east4	763→ 0	quota 0 ×738	$3.39/$2.90·h	—
v6e-8	us-central1	575→ 208 (36%)	capacity ×321, quota ×34	$4.12/$4.99·h	150.1h $659, 92 preemptions
v6e-8	asia-southeast1	222→ 0	capacity ×176, zone unsupported ×46	$2.14·h (cheapest)	—
v6e-8	us-west1	110→0	capacity ×105	$4.12·h	—
v6e-8	asia-northeast1	88→ 85 (97%)	capacity ×1	$8.74·h (2.1×)	20.0h $175, 3 preemptions
v6e-8	us-east5	29→0	capacity ×19, quota ×10	$10.21·h	—
v6e-8	europe-west4 and others	28→0	mixed	$11.10·h	—
v6e-32	all 6 regions	107→ 0	all quota 0	($8.55+·h)	tried on May 3 and May 5, then gave up
v5p-32	us-east5	58→10 (17%)	quota ×40, capacity ×8	$25.54·h	27.0h $691, 1 preemption
v5p-32	us-central1	23→ 23 (100%)	—	$13.24·h	16.3h $216, 2 preemptions
v5p-32	europe-west4	4→2	zone unsupported ×2	$20.50·h	7.4h $153
v5p-32	us-east1/us-east4	15→0	quota / zone unsupported	—	—
v5p-8	us-east5	27→ 27 (100%)	—	$6.38·h	19.2h $123, 3 preemptions
v5p-8	europe-west4	4→3		$5.12·h	33.8h $173 (31.4h, the longest run)
v5p-16	us-east5	13→4	quota ×5, capacity ×3	$12.77·h	4.7h $60
v6e-1	europe-west4	17→ 17 (100%)	—	$1.39·h	6.5h $9, 6 preemptions
v6e-1	us-east5	22→6	capacity ×11, quota ×5	$1.28·h	0.7h $1
v6e-16	asia-northeast1/ew4	8→6	quota ×2	$17.5–22.2·h	1.8h $36

Appendix C) Cheatsheet of GCP commands

Placeholders: = TPU name, = zone (e.g. us-central1-b), = GCS bucket, = bucket region (same region as the TPU zone). In a non-interactive shell on a Mac, run export PATH=/Users//google-cloud-sdk/bin:$PATH first.

In your local machine (Mac)

# allocate TPU — sweep zones until you land on one with capacity (success rates and prices: Appendix B)
for z in us-central1-b us-central1-a us-central1-c us-east4-a asia-southeast1-b; do
 gcloud alpha compute tpus tpu-vm create \
 --zone=$z --accelerator-type=v6e-8 --version=v6e-ubuntu-2404 --spot --quiet && break
done

# check the status of machine
gcloud compute tpus tpu-vm describe --zone= --format='value(state)' # READY / PREEMPTED
gcloud compute tpus tpu-vm list --zone=

# run a command on the TPU machine
gcloud compute tpus tpu-vm ssh --zone= # interactive shell
gcloud compute tpus tpu-vm ssh --zone= --command="" # one-shot command (for multi-line scripts, see heredoc patterns (b)(c) in Appendix A)
gcloud compute tpus tpu-vm scp :~/logs/31_base_train.log ./logs/ --zone= # retrieve logs

# deallocate the TPU machine (after deleting, sweep every zone you tried for zombie instances — create is async)
gcloud compute tpus tpu-vm delete --zone= --quiet
for z in us-central1-b us-central1-a us-central1-c us-east4-a asia-southeast1-b; do
 gcloud compute tpus tpu-vm list --zone=$z 2>&1 | grep -v "Listed 0"
done

# anything else?
gcloud storage buckets create gs:// --location= # same region as the TPU zone
gcloud storage rsync -r gs:///base_checkpoints /path/to/ssd/base_checkpoints # permanent mirror of the results
gcloud auth login; gcloud config set project # first time only

In your TPU instance (VM)

# core command within speedrun.sh
sudo rm -f /tmp/libtpu_lockfile # required before the first jax process
bash runs/speedrun.sh # the production pipeline in one shot
bash runs/runcpu.sh # d2/CPU wiring check (only checks it runs to the end)
# the stages speedrun.sh runs in order (same as Steps 1-3 in the main text):
python -m nanochat_jax.dataset -n 8 # B-1 for the tokenizer
python -m nanochat_jax.dataset -n 170 & # B-2 for training (background)
python -m scripts.tok_train --max-chars 100000000 --doc-cap 10000 --vocab-size 32768 # B-3
python -m scripts.tok_eval # B-4
python -m scripts.base_train # B-5
python -m scripts.base_eval --source base --model-tag d24_speedrun_r4 --eval bpb,core ... # B-6
python -m scripts.chat_sft # B-7
python -m scripts.chat_eval -i sft -g d24_speedrun_r4_sft ... # B-8
python -m nanochat_jax.report generate # B-9

# monitoring base training — the GCS checks and preemption recovery below only work if you mirror checkpoints to (periodic gcloud storage rsync recommended)
tail -5 ~/logs/31_base_train.log | grep -E 'step|bpb' # last step lines
ps -ef | grep '[b]ase_train' # process alive ([b] avoids grep matching itself)
df -h / | tail -1 # disk
gcloud storage ls -l gs:///base_checkpoints/d24_speedrun_r4/ | tail -8 # latest ckpt in GCS

# monitoring SFT training
tail -5 ~/logs/60_sft_chain.log | grep -E 'step|val_bpb'
ps -ef | grep '[c]hat_sft'

# anything else? — preemption recovery (restore from GCS, then continue)
# training: python -m scripts.base_train ... --resume-from-step 
# eval: python -m scripts.chat_eval ... --start-index 

Appendix D) Full tok_eval compression tables

diff % = (baseline tokens − our tokens) ÷ baseline tokens. Positive means we fit the same bytes into fewer tokens (a compression advantage). ratio = bytes/token.

vs GPT-2

| Text | Bytes | GPT-2 tok | ratio | Ours tok | ratio | diff % |
|---------|---------|-----------|-------|----------|-------|---------|
| news | 1803 | 392 | 4.60 | 406 | 4.44 | -3.6% |
| korean | 893 | 745 | 1.20 | 729 | 1.22 | +2.1% |
| code | 1259 | 576 | 2.19 | 410 | 3.07 | +28.8% |
| math | 690 | 348 | 1.98 | 338 | 2.04 | +2.9% |
| science | 1098 | 253 | 4.34 | 238 | 4.61 | +5.9% |
| train | 2948778 | 631304 | 4.67 | 621881 | 4.74 | +1.5% |
| val | 3024593 | 653067 | 4.63 | 645862 | 4.68 | +1.1% |

vs GPT-4

| Text | Bytes | GPT-4 tok | ratio | Ours tok | ratio | diff % |
|---------|---------|-----------|-------|----------|-------|---------|
| news | 1803 | 387 | 4.66 | 406 | 4.44 | -4.9% |
| korean | 893 | 364 | 2.45 | 729 | 1.22 | -100.3% |
| code | 1259 | 309 | 4.07 | 410 | 3.07 | -32.7% |
| math | 690 | 295 | 2.34 | 338 | 2.04 | -14.6% |
| science | 1098 | 245 | 4.48 | 238 | 4.61 | +2.9% |
| train | 2948778 | 611619 | 4.82 | 621881 | 4.74 | -1.7% |
| val | 3024593 | 631183 | 4.79 | 645862 | 4.68 | -2.3% |

---

Footnotes

1. GPT-2's ~$43k is Karpathy's retrospective estimate (TPU v3 ×32 × 168h). The R4 row's ~2h and ~$48 are my estimate: a 2-hour base train loop times ~$24/hour for 8×H100. The ~$73 total is from Karpathy's 2026 tweet (GPT-2-grade d24, 8×H100, 3.04h) — that run is the default d24 (data-ratio 12), so it's the same d24 as our R4 (ratio 9.5) but a separate recipe. Our precise CORE value is 0.269486 (see the report card); the text uses 0.2695. ↩

2. The common belief that "rounding the vocab up to a hardware-friendly multiple makes it faster" comes from a GPU-era tweet by Karpathy ( 50257→50304, rounded up to a multiple of 64, ~25% faster via kernel occupancy). On TPU, XLA automatically pads matrix dimensions to multiples of the MXU tile (128 on v5p, 256 on v6e) ( Cloud TPU performance guide), and 32768 is already such a multiple, so nothing is wasted on padding. On how vocab size relates to model scale, the vocab scaling law ( Tao et al. 2024: larger models deserve larger vocabularies, but sublinearly) is a good read. ↩

3. The ~6.9B tokens come from the tokens:params ratio nanochat uses. It follows the compute-optimal ratio concept established by Chinchilla (the params–tokens combination that gives the lowest loss for a given compute budget), but the value itself is nanochat's own measurement, not Chinchilla's 20: an early scaling fit of ~8 ( miniseries #420) was later re-measured to a compute-optimal ~10.5, which made the script default 12 ( #481), and R4 uses 9.5, slightly undertrained relative to that. Karpathy's guess for why the value comes out below 20 is the influence of the Muon optimizer, among other things (#420). The ratio multiplies the non-embedding (scaling) parameters (d24 ≈ 730M): 730M × 9.5 ≈ 6.9B tokens. Of the ClimbMix shards, training uses about 150; 171 is the download count with some margin (170 train + 1 validation). ↩ ↩

4. Block sizes are for performance and have almost no effect on numerics (Pallas BlockSizes docstring). The hard constraints mostly live in the Pallas splash kernel (JAX 0.10.0, the version we use), not in our code — block_kv must be an integer multiple of block_kv_compute, and block_kv_compute must be a multiple of 128 (=NUM_LANES) ( kernel L967–970); block_q and block_kv must also be multiples of 128 ( L667 and L673). Blocks must be ≤ seq_len ( L1446–1449) and must fit in on-chip VMEM (exceeding it gives a compile/allocation error — splash itself doesn't check this). The only constraint our wrapper enforces directly is block_kv_compute % 128 ( attention.py#L82-L85). MaxText's default is also 512 (our starting value), and larger models raise only block_q ( gemma2-9b on v6e uses sa_block_q=2048; block_kv stays at 512). Tune with XProf — the splash kernel tags its block sizes into the trace as XProf metadata. ↩

5. In his homelab d24 reproduction ( #710), matt-langston reports his DGX Spark at 29.6% MFU and puts the general H100 band at 50–65% (no depth given; the same thread also explains the difference between GPU utilization and MFU). Karpathy's own measured anchors are ~47–48% for the d20 speedrun and 47.40% for d34. nor, who ported modded-nanogpt to the same v6e-8, also landed at 15→23%, and the GPT-3 training MFU in the PaLM paper's table is 21.3%. So our 24% is not far outside this distribution. That said, these numbers differ in seq_len, model architecture, and how FLOPs are counted, so they can't be compared head-on — MFU was originally a metric for comparing system efficiency on the same model ( PaLM §4.1). ↩ ↩

6. For scaling and rooflines on TPU, Google DeepMind's scaling-book is a textbook-level resource. ↩

7. The general way to add narrow abilities like this via midtraining/SFT is laid out in Karpathy's "counting r in strawberry (and how to add abilities generally)" ( nanochat #164). ↩

8. Karpathy, "Guide: infusing identity to your nanochat" ( nanochat #139) — mixing 1,000 synthetic identity conversations into the midtraining/SFT mixture. ↩

9. np.asarray not copying an input that is already an ndarray is NumPy's documented behavior — a slice is a view sharing the same buffer, and changes through a view show up in the original. NumPy inside data pipelines is a well-known source of silent bugs: Tanel Pärnamaa downloaded and audited 100K+ GitHub repositories that import PyTorch, and reported that among the repos combining a custom dataset, NumPy's random number generator, and multi-process loading, over 95% had a related bug (PyTorch's official tutorials, OpenAI, and NVIDIA code included; post). Not exactly the same bug as ours, but the same family. ↩

10. An earlier version of the port had a "sibling bug", too — the grad-accum accumulation loop was missing entirely, so the model trained on only 1/N of the configured tokens while the LR schedule decayed as if full batches were being consumed (silent under-training — the early loss looks normal, and only the final quality suffers). Upstream nanochat derivesgrad_accum_steps from total_batch_size and asserts the two divide evenly, which blocks this kind of mismatch at the source. Grad-accum is just an easy place for things to silently go wrong: HuggingFace and Unsloth also shipped grad-accum bugs (different mechanism — loss normalization) for a while in 2024 before fixing them ( Unsloth and HuggingFace), and it's a known point of confusion in na