카카오는 메시징 플랫폼을 통해 한국 인구의 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 생태계의 모듈식 설계는 프로덕션 파이프라인과 고급 연구 모두를 용이하게 합니다.