nanochat en una TPU: qué se conserva de PyTorch y qué se rompe al migrar a JAX

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

Un desarrollador portó nanochat, el proyecto full-stack de LLM de Karpathy, a JAX para ejecutarlo en una TPU v6e-8. La calidad se reprodujo: el modelo base alcanzó un CORE de 0.2695, superando la banda de referencia de Karpathy en H100. Sin embargo, el rendimiento quedó a la mitad: el MFU fue del 24% frente al 47-48% de las H100. El artículo detalla los cambios necesarios, como el uso de kernels Pallas y ajustes de precisión, y analiza las diferencias arquitectónicas entre TPU y GPU que explican la brecha.