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.
- konradha
Every engineer who's worked with "agentic" workflows can tell that this is entirely generated with very little tweaks. It's embarrassing.
- tucan9389
Author here.
Thanks for comments and feedback. I wrote this post in Korean (my native language) and translated it by LLMs. TBH, I spent more time reviewing the post rather than writing. Sorry that I missed the uncomfortable tone and AI-ish. be more careful next time.
Anyway, even though I tried to use coding agent, it was hard to avoid some dead ends and spent more money than I expected. As a newbie in pre-training, hope this baseline codebase helps people who want to reproduce nanochat in JAX + TPU...
If anyone has any useful experiences when you do similar things, please share in the comments.
- dansuh
Great work! Love the lessons learned.