Kakao, qui dessert 93 % de la population de la Corée du Sud sur sa plateforme de messagerie, a rencontré des problèmes de scalabilité avec son infrastructure existante. L'équipe d'ingénierie a stratégiquement migré vers les TPU de Google Cloud en utilisant le framework JAX pour optimiser les coûts et l'efficacité. Ce mouvement a permis le développement de la famille de modèles propriétaires Kanana, dont plusieurs sont maintenant open-sourcés. L'équipe a détaillé leur parcours technique, notamment l'adaptation de JAX, MaxText pour les pipelines de données personnalisés, et la formation de modèles Mixture-of-Experts (MoE). Ils ont personnalisé MaxText pour fusionner dynamiquement plusieurs sources de données et ont modifié le traitement des jetons pour améliorer l'efficacité. Les expériences de validation ont confirmé que les performances des TPU étaient équivalentes à celles de leur pipeline basé sur des GPU. D'autres expériences ont impliqué la transformation d'un modèle dense en une architecture MoE sur des TPU v5e, démontrant la pertinence de la pile JAX pour la recherche avancée. Kakao a observé une augmentation significative du débit et une amélioration de la performance coût avec les TPU Trillium. Les avantages clés de la pile JAX incluent les performances, la scalabilité, la personnalisation et l'adoption rapide de fonctionnalités. La conception modulaire de l'écosystème JAX, notamment MaxText, Flax, Optax et Orbax, facilite à la fois les pipelines de production et la recherche avancée.