Karpathy的nanochat移植TPU实测

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

我把Karpathy的nanochat项目从PyTorch移植到了JAX，并在TPU v6e-8上进行了完整复现。这次实验旨在保持架构高度一致，同时验证模型质量与训练性能。结果显示，CORE评分成功复现甚至略超预期，达到了0.2695，证明了模型质量的可迁移性。但在训练效率上，TPU的MFU仅为24%，远低于H100的47%，暴露了硬件差异带来的挑战。本文详细记录了从Tokenizer训练到SFT的全流程，分析了XLA编译、Pallas算子等JAX生态组件在TPU上的表现，为开发者提供了宝贵的实战参考。

"如果nanochat对你来说是陌生的：它是Karpathy的全栈LLM项目，将tokenizer训练、预训练、SFT和RL整合在一个仓库中，只需约四小时和100美元，你就能在单台8×H100节点上获得自己的聊天机器人。"