Porting nanochat to a TPU: What Carries Over from PyTorch and What Breaks

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

I ported Karpathy's nanochat to a Google Cloud TPU v6e using JAX and Flax, aiming for architectural parity with the original PyTorch version. While the model quality exceeded expectations with a strong CORE score, training performance lagged significantly behind H100 GPU benchmarks. This post details the speedrun results, hardware trade-offs, and specific code changes required to make the transition work.

"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."