Porting nanochat to TPU v6e: PyTorch Wins and JAX Breaks

Karpathy's nanochat LLM ported to TPU v6e-8 with JAX achieves identical CORE scores but only 24% MFU vs 48% on H100. The post details what carries over from PyTorch and what breaks on TPU, including Pallas kernel tuning and memory constraints.

4 min readJul 18, 2026

Porting nanochat to TPU v6e: PyTorch Wins and JAX Breaks

nanochat on TPU: The $100 speedrun meets JAX

Karpathy's nanochat — a full-stack LLM project from tokenizer training to RL — normally runs on an 8×H100 GPU node and costs ~$100 for a 4-hour speedrun. Now there's a JAX port targeting TPU v6e-8 slices. The results: CORE score of 0.2695 (above the reference band of 0.2512–0.2677) but MFU of only ~24%, half of the H100's 47–48%. The port is documented in a GitHub discussion by tucan9389.

Reproduction results: quality matches, performance lags

The port reproduces recipe 4 (d24, 1.4B parameters) from nanochat's leaderboard. Base model training took 5.29h compute (6.02h wall clock) on a TPU v6e-8 spot instance, costing $30 ($130 on-demand). SFT added 68.7 minutes. Total wall clock: 12.19h, costing $60.8 spot ($263 on-demand). The CORE score of 0.2695 lands just above Karpathy's 7-run band of 0.2512–0.2677. MFU at 24% is the main gap.

TPU basics: v6e's 256×256 MXU and 32GB HBM

The v6e (Trillium) has a 256×256 matrix multiply unit (MXU), up from 128×128 in previous generations. If a tensor dimension isn't a multiple of 256, XLA pads with zeros, wasting compute. HBM is only 32GB per chip — a third of v5p's 95GB — but compute is 2×. This compute-heavy, memory-lean design forces careful kernel tuning.

Speedrun pipeline: tokenizer, base, SFT

The pipeline runs three stages:

Tokenizer (5.3m): Trains a vocab-32768 tokenizer on 100M characters of ClimbMix. It uses 1.5% fewer tokens than GPT-2 tokenizer on the training data — a compression advantage matching Karpathy's results.

Base model (6.02h): Uses the R4 recipe with --splash-* kernel block sizes tuned for v6e. Key arguments:

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 \
  --attn-impl=splash --splash-block-q=512 --splash-block-kv=512 --splash-block-kv-compute=256 \
  --bf16 --cast-embeddings-bf16 --use-real-data

Training logs show loss dropping from 4.517 to 2.786, with MFU ~24.5%. CORE eval yields 0.2695.

SFT (68.7m): Fine-tunes on a mixture of SmolTalk, MMLU, GSM8K, SpellingBee, and 1,000 identity conversations. ChatCORE reaches 0.3733, with SpellingBee at 0.9961 and GSM8K at 0.1008 (math improvement left for RL).

What carries over from PyTorch

The architecture and config (recipe R4) port unchanged. Data loading uses NumPy, not torch.utils.data. Checkpoints remain PyTorch .pt for compatibility. The optimizer is implemented in pure JAX, no torch dependency.

What breaks on TPU

Attention kernel: PyTorch uses FlashAttention via Triton; TPU requires Pallas SplashAttention. Block sizes must be tuned (--splash-block-q=512, --splash-block-kv=512, --splash-block-kv-compute=256).
Performance: MFU drops to 24% due to MXU padding and memory bandwidth limits.
Preemption: Spot TPUs can be reclaimed; this run had one preemption with 17-minute recovery.
Cost: On-demand $263 vs H100's ~$48, though spot at $60.8 is comparable.

Key technical details

v6e-8 slice: 7.34 bf16 PFLOPS, 32GB HBM per chip, 256×256 MXU.
CORE score: 0.2695 (above GPT-2's 0.2565 and R4 band).
MFU: ~24% vs H100's 47–48%.
Training time: 5.29h compute (base) vs H100's 2h.
Cost: $30 spot base training vs H100's ~$48.

Why it matters

This port demonstrates that JAX + TPU can match PyTorch + GPU quality for LLM training, but performance optimization is still immature. For developers exploring TPU alternatives to NVIDIA GPUs, this is a realistic benchmark: expect ~2× longer training time and ~2× lower MFU, but with potentially lower spot costs. The need for manual kernel tuning (Pallas vs Triton) is a significant friction point.

Next steps

The author plans to improve MFU by optimizing Pallas kernel block sizes and exploring XLA compilation flags. For now, if you want to run nanochat on TPU, expect to invest time in kernel tuning and accept slower training. The JAX port is open-source at the GitHub discussion link.

Editor's Take

I've run nanochat on H100s before, and the TPU port's 24% MFU is disappointing but expected. The v6e's 32GB HBM is a real bottleneck for 1.4B parameter models — you can't even fit the optimizer states without gradient accumulation. I think the real value here is the detailed breakdown of what breaks, especially the Pallas attention kernel tuning. If you're serious about TPU for LLMs, you'll need to invest in custom kernels, not just rely on XLA.

— DevDigest Editorial

Key Takeaways

•Expect ~2× longer training time on TPU v6e vs H100 for similar quality, with MFU around 24%.
•Tune Pallas SplashAttention block sizes (--splash-block-q, --splash-block-kv, --splash-block-kv-compute) to match TPU MXU dimensions (256×256).
•Use spot TPU instances for cost savings but plan for preemption recovery (this run had one 17-minute recovery).

Why It Matters

This is a real-world benchmark of training a full LLM pipeline on TPU v6e using JAX, showing quality parity but significant performance gaps. For developers considering TPU for cost or availability reasons, this provides concrete numbers and identifies specific pain points like MXU padding and memory constraints.

#jax#tpu#karpathy#nanochat#LLM training

Get the weekly digest

Every Sunday - top tech stories, industry breakthroughs, and developer tools delivered to your inbox.

No spam, unsubscribe anytime.