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

Every engineer who's worked with "agentic" workflows can tell that this is entirely generated with very little tweaks. It's embarrassing.

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.

Nice work. Regarding your first insight, that's a sneaky one ^^. I experiment a similar issue with fine-tuning a small model with MLX. It was very hard to find out.

For people that are like me : This entire text is AI generated, i feel weird reading it personally, i guess others may not

Great work! Love the lessons learned.

nice