Какао, обслуживающая 93% населения Южной Кореи на своей платформе обмена сообщениями, столкнулась с проблемами масштабируемости с существующей инфраструктурой. Инженерная команда стратегически перешла на Google Cloud TPUs с использованием фреймворка JAX для оптимизации стоимости и эффективности. Это позволило разработать свою проприетарную семейство моделей Kanana, несколько из которых теперь открыты для общего доступа. Команда подробно описала свой технический путь, включая адаптацию JAX, MaxText для настройки пользовательских конвейеров данных, и обучение моделей Mixture-of-Experts (MoE). Они настроили MaxText для динамического объединения нескольких источников данных и модифицировали обработку токенов для эффективности. Эксперименты по валидации подтвердили, что производительность TPU соответствует их конвейеру на основе GPU. Дальнейшие эксперименты включали в себя преобразование плотной модели в архитектуру MoE на v5e TPUs, демонстрируя пригодность стека JAX для передовых исследований. Какао наблюдала значительное увеличение производительности и улучшение стоимости-эффективности с помощью Trillium TPUs. Ключевыми преимуществами стека JAX являются производительность, масштабируемость, настраиваемость и быстрое внедрение новых функций. Модульный дизайн экосистемы JAX, включая MaxText, Flax, Optax и Orbax, облегчает как производственные конвейеры, так и передовые исследования.