Kakao，作为韩国93%人口的消息平台，面临着现有基础设施的可扩展性问题。工程团队策略性地转移到使用 Google Cloud TPUs 的 JAX 框架，以优化成本和效率。这种举措使得他们开发了专有的 Kanana 模型家族，其中一些模型现在已经开源。团队详细介绍了他们的技术之旅，包括适应 JAX、MaxText 的自定义数据管道和 Mixture-of-Experts（MoE）模型训练。他们对 MaxText 进行了自定义，以动态地混合多个数据源，并修改了令牌处理以提高效率。验证实验确认了 TPU 的性能与基于 GPU 的管道相匹配。进一步的实验涉及将密集模型升级到 MoE 架构在 v5e TPUs 上，展示了 JAX 栈的适合性高级研究。Kakao 观察到 Trillium TPUs 带来了显著的吞吐量增加和成本性能改善。JAX 栈的关键优势包括性能、可扩展性、可定制性和快速特征采纳。JAX 生态系统的模块化设计，包括 MaxText、Flax、Optax 和 Orbax，促进了生产管道和高级研究。