Port en TPU de $60 consigue que nanochat alcance calidad de GPT-2

El nanochat de Karpathy puede entrenar un pequeño chatbot en un nodo 8×H100 en aproximadamente cuatro horas por unos $100. Un nuevo port a JAX, nanochat-jax, se propuso mantener la configuración y la arquitectura originales lo más cerca posible mientras movía la pila a un TPU v6e-8. El resultado: la calidad del modelo se trasladó de forma limpia, pero el rendimiento no.

La reproducción se centró en la receta 4 del tablero oficial de nanochat, también conocida como R4 o d24, un modelo con profundidad 24 y unos 1.4B parámetros. Usando el mismo arnés de evaluación que nanochat, el port JAX alcanzó una puntuación CORE de 0.2695. Eso está por encima del 0.2565 de GPT-2 y un poco por encima del rango 0.2512–0.2677 que Karpathy vio en 7 ejecuciones de la misma receta.

La velocidad de entrenamiento fue el punto débil. En el TPU, el MFU quedó en alrededor del 24% para d24, aproximadamente la mitad del 47–48% informado por Karpathy en 8×H100 para d20. El bucle de entrenamiento del modelo base reportado tardó 5.29 horas, o 6.02 horas incluyendo guardado de puntos de control y compilación, frente a unas 2 horas en la configuración H100.

La ejecución completa se detuvo en SFT en lugar de continuar hasta RL. Según el autor, la canalización cubre:

• Entrenamiento del tokenizador: 5.3 minutos
• Entrenamiento del modelo base: 6.02 horas más 44.5 minutos de evaluación
• SFT: 68.7 minutos más unas 3.5 horas de evaluación

El coste fue uno de los resultados más interesantes. La ejecución en TPU usó una slice de 8 chips v6e de un solo host con precios spot y sumó $60.8 en 12.19 horas, con una preempción y recuperación. La misma ejecución habría costado unos $263 a precios bajo demanda. Para comparar, la publicación cita nanochat R4 en 8×H100 en aproximadamente $48, y GPT-2 (2019, 1.5B) en TPU v3 x32 en unas 168 horas estimadas y ~$43,000.

Recomendado

Google enfrenta demanda de editoriales por derechos de autor contra Gemini

La pila de software se mantuvo cercana en espíritu a PyTorch, pero cambió a Flax NNX para los módulos de redes neuronales, JAX + XLA para el cálculo y la compilación, y Pallas para kernels personalizados como Splash Attention. Los checkpoints siguen en formato .pt de PyTorch por compatibilidad.

Especificaciones del TPU v6e-8 y resultados de entrenamiento

La publicación también destaca por qué el TPU se comporta de forma diferente. Comparado con el v5p, el v6e duplica el cómputo bf16 hasta 918 TFLOPs por chip pero reduce la capacidad HBM a 32 GB desde 95 GB. Su unidad matricial también creció de 128×128 a 256×256, lo que significa que las dimensiones de tensor que no son múltiplos de 256 son rellenadas por XLA, desperdiciando parte de la unidad.

Para el tokenizador, el port entrenó un vocabulario de 32,768 tokens con 100M de caracteres de ClimbMix. En los datos de entrenamiento, usó un 1.5% menos de tokens que el tokenizador de GPT-2, reproduciendo la misma ventaja general de compresión vista en la versión de Karpathy.

En el modelo base, la bpb de validación quedó en 0.7343, frente a 0.7185 para R4, aunque el autor señala que las diferencias de tokenizador hacen que eso sea solo un punto de referencia aproximado. El resultado más importante es que el modelo reproducido alcanzó la banda de calidad objetivo.

Resultados de SFT en tareas de chat

Tras 68.7 minutos de ajuste fino supervisado en unos 1.07M de filas extraídas de SmolTalk, MMLU, GSM8K, SpellingBee y 1,000 conversaciones de identidad, el rendimiento de chat del modelo mejoró notablemente. ChatCORE, promediado en 6 tareas de chat, alcanzó 0.3733.

El mayor salto se dio en tareas generativas donde el modelo base había obtenido efectivamente 0. Tras SFT, la publicación informa:

• SpellingBee: 0.9961
• GSM8K: 0.1008

El autor no compara esos números directamente con la hoja de calificaciones anterior de Karpathy porque esa ejecución usó d20 y un pipeline de entrenamiento intermedio más antiguo que desde entonces se ha integrado en SFT aguas arriba.

La conclusión del port es limitada pero útil: la paridad PyTorch-a-JAX es posible en cuanto a calidad para esta receta, pero en TPU v6e-8 la pila de entrenamiento todavía deja una gran brecha de rendimiento frente a 8×H100.