Kakao, das mit seiner Messaging-Plattform 93 % der südkoreanischen Bevölkerung bedient, sah sich mit Skalierbarkeitsproblemen seiner bestehenden Infrastruktur konfrontiert. Das Ingenieurteam wechselte strategisch zu Google Cloud TPUs unter Verwendung des JAX-Frameworks, um Kosten und Effizienz zu optimieren. Dieser Schritt ermöglichte die Entwicklung ihrer proprietären Kanana-Modellfamilie, von der mehrere Modelle inzwischen Open Source sind. Das Team detaillierte seinen technischen Weg, einschließlich der Anpassung von JAX, MaxText für benutzerdefinierte Datenpipelines und des Trainings von Mixture-of-Experts (MoE)-Modellen. Sie passten MaxText an, um mehrere Datenquellen dynamisch zu kombinieren, und modifizierten die Token-Verarbeitung zur Effizienzsteigerung. Validierungsexperimente bestätigten, dass die TPU-Leistung ihrer GPU-basierten Pipeline entsprach. Weitere Experimente umfassten das Upcycling eines dichten Modells in eine MoE-Architektur auf v5e TPUs, was die Eignung des JAX-Stacks für die fortgeschrittene Forschung demonstriert. Kakao verzeichnete mit Trillium TPUs eine deutliche Durchsatzsteigerung und eine verbesserte Kosten-Leistungs-Bilanz. Zu den Hauptvorteilen des JAX-Stacks gehören Leistung, Skalierbarkeit, Anpassbarkeit und schnelle Feature-Übernahme. Das modulare Design des JAX-Ökosystems, einschließlich MaxText, Flax, Optax und Orbax, erleichtert sowohl Produktionspipelines als auch fortgeschrittene Forschung.