カカオは、メッセージングプラットフォームで韓国の人口の93%にサービスを提供していますが、既存のインフラストラクチャでスケーラビリティの問題に直面していました。エンジニアリングチームは、コストと効率を最適化するために、JAXフレームワークを使用したGoogle Cloud TPUへの移行を戦略的に行いました。この移行により、同社独自のKananaモデルファミリーの開発が可能になり、いくつかのモデルがオープンソース化されました。チームは、JAXの適応、カスタムデータパイプラインのためのMaxText、およびMixture-of-Experts（MoE）モデルのトレーニングを含む技術的な道のりを詳細に説明しました。彼らは、複数のデータソースを動的にブレンドするためにMaxTextをカスタマイズし、効率のためにトークン処理を変更しました。検証実験により、TPUのパフォーマンスがGPUベースのパイプラインと同等であることが確認されました。さらなる実験では、v5e TPUで密なモデルをMoEアーキテクチャにアップサイクルし、JAXスタックの高度な研究への適合性を示しました。カカオは、Trillium TPUでスループットの大幅な増加とコストパフォーマンスの向上を観測しました。JAXスタックの主な利点には、パフォーマンス、スケーラビリティ、カスタマイズ性、および迅速な機能採用が含まれます。MaxText、Flax、Optax、Orbaxを含むJAXエコシステムのモジュラー設計は、本番パイプラインと高度な研究の両方を容易にします。