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

https://github.com/tucan9389/nanochat-jax/discussions/1

by tucan9389 • 8 days ago

54 12 8 days ago

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.

tucan93892• 4 days ago

Thank you. Agreed, we should be more careful when doing this. Was it the JAX training codebase?

adityashankar• 5 days ago

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

embedding-shape2• 5 days ago

Don't tell me you don't enjoy and learn a ton from barely coherent nuggets like this?

The most important thing the port taught me: if you can avoid it, don't. I decided to do it anyway, so here are the problems I ran into and some insights of my own, trimmed down to five. Each item is tagged with which side it hurt: quality (CORE) or performance (MFU). The first is the bug that fooled us the longest.

Chris2048• 5 days ago

TBH llm generated text is usually better than this..

attila-lendvai• 5 days ago

he's just dogfooding us all...

konradha• 5 days ago

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

stephantul• 5 days ago

Even the title is a Claudeism, it makes me sad

tahakotwal12• 5 days ago

nice

tucan93892• 4 days ago

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• 2 hours ago

Great work! Love the lessons learned.

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

Related Stories