GPU는 학습 과정에서 가장 고민되는 부분입니다. 필요할 때 적절한 컴퓨팅 자원이 필요했기 때문에 카카오도 사내에 GPU를 넉넉하게 확보했습니다. 하지만 카카오도 전사적으로 AI에 대한 수요가 늘어나면서 많은 부서가 GPU를 필요로 하게 됐습니다. 카나나는 더 큰 데이터 학습이 필요했고, 점점 부담스러워지는 GPU 확보는 늘 답답함을 안겼습니다.

“기존의 GPU 기반 인프라는 늘 전력 및 예산의 한계를 안고 있었습니다. 이를 해소하는 방법은 다시 더 많은 GPU를 확보하거나, 효율성이 더 높은 새로운 환경으로 전환하는 것이었습니다. TPU는 가장 효과적인 환경이지만 새로운 소프트웨어 개발 환경을 도입해야 한다는 걱정이 있었습니다.”

안재현 AI 엔지니어는 카나나의 LLM 학습량이 늘어나면서 TPU 전환은 자연스럽게 결정됐다고 말했습니다. 단기적으로는 기존 CUDA 기반의 프레임워크를 내려놓고 JAX로 스택을 구성하기 위해 새로운 개발 환경을 공부해야 했지만 길게 보면 비용에 비해 훨씬 높은 컴퓨팅 성능을 얻을 수 있는 필수적인 투자라고 판단한 것입니다.

TPU는 수 백, 수 천 개의 프로세서 코어가 유기적으로 맞물려 텐서플로 프레임워크를 재빠르게 처리하도록 설계되어 있습니다. 무엇보다 필요한 만큼의 코어를 갖춘 시스템이 할당되기 때문에 학습에 필요한 컴퓨팅 환경을 안정적으로 확보할 수 있습니다.

카카오는 기존의 TPU 기반의 모델 개발 경험을 토대로 언어 모델의 학습에서 TPU를 다루는 방법을 잘 알고 있었고, 그 사이에 업그레이드된 TPU v5e를 비롯해 TPU Trillium을 통해 더 강력한 LLM 학습 환경을 기대할 수 있었습니다.

류민호 AI 리서처는 Trillium의 안정성도 높게 평가했습니다. GPU는 하나가 말썽을 일으키면 묶여 있는 시스템이 먹통이 되는 경우도 있는데 TPU의 패키지는 코어 한 두 개가 말썽을 부려도 전체 시스템에 영향을 끼치지 않습니다. 실제로 구글 클라우드의 안정성 만큼이나 TPU도 신뢰도가 높기 때문에 안정성에 대한 걱정, 기술적 고민 없이 AI 연구원들이 프로젝트에만 집중할 수 있었다고 합니다.

카카오는 TPU로 효과적인 학습 환경을 갖추면서 공격적으로 모델의 성능을 높이는 방법을 떠올렸습니다. 기존 모델의 레이어를 더 많이 쌓아 올려 모델 규모를 키우는 방법인 Depth Upscaling(DUS) 방식을 반복적으로 적용했습니다. 류민호 AI 리서처는 덕분에 확보한 시스템 자원보다 훨씬 더 큰 모델을 효과적으로 학습시킬 수 있었다고 말합니다.

“핵심 파이프라인이 검증된 뒤, 강력한 성능을 유지하면서도 추론 효율을 높이기 위해 MoE(Mixture-of-Experts) 모델을 통해 더 진보된 아키텍처 실험을 시작했습니다. 기존 21억 파라미터의 밀집 모델을 134억 파라미터 규모의 MoE 아키텍처로 업사이클링하고 MaxText를 통해 학습했습니다. JAX 스택은 스케줄 결합과 코드 커스터마이징 등 개발의 번거로움을 최소화하면서도 코드와 수학 벤치마크 등에서 큰 성과를 확인해 주었습니다.”

JAX 스택은 고급 모델 개발에서 기대 이상의 성능을 보여주었고, 특정 전문 도메인에 특화도 매끄럽게 이뤄질 수 있다는 가능성을 보여주었습니다. 그 덕분에 카카오는 24.6B 파라미터 규모의 모델을 빠르게 완성해 Gemma-2-27B 등 잘 알려진 모델들과 어깨를 나란히 할 만한 성능을 얻을 수 있었습니다. 여기에 카나나는 한국어와 한국 문화에 대한 부분에서는 월등하게 뛰어난 결과를 보여주기도 했습니다.

“더 만족스러웠던 것은 TPU 환경이 지속적으로 더 좋아진다는 점이었습니다. 지난해 발표된 TPU v5에 비해 2024년 공개된 TPU Trillium은 AI 가속기로서 훨씬 발전된 모습을 보여주었습니다. TPU Trillium은 칩당 성능이 4배 이상 빨라졌고, 칩당 HBM 용량은 16GB에서 32GB로 두 배 늘었고 대역폭도 두 배 넓어졌습니다. 더 높은 성능의 코어를 기존보다 더 많이 연결해서 쓸 수 있기 때문에 학습의 결과를 훨씬 빠르게 받아볼 수 있었습니다.”

TPU의 세대 전환 과정도 수월했습니다. TPU Trillium은 더 나은 성능 뿐 아니라 새로운 요소들을 많이 갖추었는데, 이를 다루는 입장에서는 바뀌는 부분에 대한 기술적 부담을 가질 필요가 없습니다. 모든 프레임워크도 v6e와 적절히 맞물려 업데이트되기 때문에 마치 PC의 CPU나 GPU를 갈아 끼우듯 매끄럽게 v5e 환경을 v6e로 이전할 수 있었습니다. 이는 앞으로도 부담 없이 새로운 세대의 기술을 받아들일 수 있다는 확신으로 이어집니다.

“계속해서 새 GPU가 나오면서 어렵게 확보한 시스템이 낡아가는 것에 대한 부담도 있었습니다. 비용적 손실인 셈입니다. TPU 역시 거의 매년 새 버전이 나오면서 성능이 크게 높아지는데, 기존 환경을 정리하는 부담 없이 곧바로 새 환경으로 이전할 수 있어서 전체 유지 비용 측면에서도 유리했습니다.”

안재현 AI 엔지니어는 TPU를 도입하면서 인공지능과 함께 가는 길에 대한 선택지가 늘었다고 말합니다. 인공지능의 발전에는 GPU뿐 아니라 여러 형태의 하드웨어가 함께 하고 있습니다. TPU는 또 하나의 새로운 방법입니다. 카카오는 이를 통해 카나나의 가능성을 확인하고 인공지능과 카카오의 서비스가 함께 나아갈 방향성을 잡을 수 있었습니다.