$60 TPU port gets nanochat to GPT-2-grade quality

Karpathy’s nanochat can train a small chatbot on an 8×H100 node in roughly four hours for about $100. A new JAX port, nanochat-jax, set out to keep the original configuration and architecture as close as possible while moving the stack onto a TPU v6e-8. The result: model quality carried over cleanly, but performance did not.

The reproduction focused on recipe 4 from the official nanochat leaderboard, also known as R4 or d24, a model with depth 24 and about 1.4B parameters. Using the same evaluation harness as nanochat, the JAX port reached a CORE score of 0.2695. That is above GPT-2's 0.2565 and slightly above the 0.2512–0.2677 range Karpathy saw across 7 runs of the same recipe.

Training speed was the weak point. On the TPU, MFU came in at about 24% for d24, roughly half of Karpathy’s reported 47–48% on 8×H100 at d20. The reported base-model training loop took 5.29 hours, or 6.02 hours including checkpointing and compilation, versus about 2 hours on the H100 setup.

The full run stopped at SFT rather than continuing to RL. According to the author, the pipeline covers:

• Tokenizer training: 5.3 minutes
• Base model training: 6.02 hours plus 44.5 minutes of evaluation
• SFT: 68.7 minutes plus about 3.5 hours of evaluation

Cost was one of the more interesting results. The TPU run used a single-host 8-chip v6e slice on spot pricing and totaled $60.8 over 12.19 hours, with one preemption and recovery. The same run would have cost about $263 at on-demand rates. For comparison, the post cites nanochat R4 on 8×H100 at roughly $48, and GPT-2 (2019, 1.5B) on TPU v3 x32 at an estimated 168 hours and ~$43,000.

Recommended reading

Google faces publishers' Gemini copyright suit

The software stack stayed close to PyTorch in spirit, but swapped in Flax NNX for neural-network modules, JAX + XLA for computation and compilation, and Pallas for custom kernels such as Splash Attention. Checkpoints remain in PyTorch .pt format for compatibility.

TPU v6e-8 specs and training results

The post also highlights why the TPU behaves differently. Compared with v5p, the v6e doubles bf16 compute to 918 TFLOPs per chip but cuts HBM capacity to 32 GB from 95 GB. Its matrix unit also grew from 128×128 to 256×256, which means tensor dimensions that are not multiples of 256 get padded by XLA, wasting part of the unit.

For the tokenizer, the port trained a 32,768-token vocabulary on 100M characters of ClimbMix. On the training data, it used 1.5% fewer tokens than the GPT-2 tokenizer, matching the same general compression advantage seen in Karpathy’s version.

On the base model, validation bpb landed at 0.7343, versus 0.7185 for R4, though the author notes tokenizer differences make that only a loose reference point. The more important result is that the reproduced model reached the target quality band.

SFT results on chat tasks

After 68.7 minutes of supervised fine-tuning on about 1.07M rows drawn from SmolTalk, MMLU, GSM8K, SpellingBee, and 1,000 identity conversations, the model’s chat performance improved sharply. ChatCORE, averaged over 6 chat tasks, reached 0.3733.

The biggest jump came on generative tasks where the base model had effectively scored 0. After SFT, the post reports:

• SpellingBee: 0.9961
• GSM8K: 0.1008

The author does not compare those numbers directly with Karpathy’s earlier report card because that run used d20 and an older midtraining pipeline that has since been folded into SFT upstream.

The takeaway from the port is narrow but useful: PyTorch-to-JAX parity is possible for quality on this recipe, but on TPU v6e-8 the training stack still leaves a large performance gap versus 8×H100.