Kakao, que atiende al 93% de la población de Corea del Sur en su plataforma de mensajería, enfrentó problemas de escalabilidad con su infraestructura existente. El equipo de ingeniería transitó estratégicamente a las unidades de procesamiento de tensors (TPUs) de Google Cloud utilizando el marco de trabajo JAX para optimizar costos y eficiencia. Este movimiento permitió el desarrollo de su familia de modelos propietarios Kanana, con varios modelos ahora de código abierto. El equipo detalló su viaje técnico, incluyendo la adaptación de JAX, MaxText para pipelines de datos personalizados, y entrenamiento de modelos de Expertos en Mezcla (MoE). Personalizaron MaxText para combinar dinámicamente múltiples fuentes de datos y modificaron el procesamiento de tokens para eficiencia. Los experimentos de validación confirmaron que el rendimiento de las TPU coincidía con su pipeline basada en GPU. Experimentos adicionales involucraron reciclar un modelo denso en una arquitectura MoE en TPUs v5e, demostrando la idoneidad de la pila JAX para investigación avanzada. Kakao observó un aumento significativo en el rendimiento y una mejora en el costo-rendimiento con las TPUs Trillium. Las ventajas clave de la pila JAX incluyen rendimiento, escalabilidad, personalizabilidad y adopción rápida de características. El diseño modular del ecosistema JAX, incluyendo MaxText, Flax, Optax y Orbax, facilita tanto pipelines de producción como investigación avanzada.