nanochat auf einer TPU: Was von PyTorch überlebt und was bricht
Porting nanochat to a TPU: what carries over from PyTorch, and what breaks
Karpathys nanochat, normalerweise auf einem 8×H100-GPU-Knoten trainiert, wurde auf eine TPU v6e-8 portiert. Der Port reproduziert die Qualität des Basismodells (CORE-Score 0,2695, leicht über dem Referenzband) und erreicht eine MFU von etwa 24 % – etwa die Hälfte der H100-Werte. Der Beitrag dokumentiert die Portierung, vergleicht die Ergebnisse mit den Referenzwerten und erläutert die Unterschiede zwischen PyTorch und JAX/Flax auf TPUs, inklusive der Auswirkungen der neuen 256×256-MXU auf die Effizienz.
Die v6e-MXU wuchs auf 256×256, von den 128×128 jeder Generation bis zur v5p – wenn eine Tensordimension kein Vielfaches von 256 ist, füllt XLA sie mit Nullen auf und ein Teil der Einheit wird verschwendet.