카카오의 AI 모델 ‘카나나’, 구글 클라우드 TPU로 남들과 다른 길을 걷게 된 이야기

카카오가 ‘카나나(Kanana)’라는 이름으로 통합 AI브랜드를 새롭게 선보였습니다. 생성형 AI를 비롯해 카카오가 직접 개발하고 운영하는 인공지능 서비스와 모델 등을 카나나라는 이름으로 정의하는 밑그림입니다. 카카오는 우리나라에서 가장 많은 사람들이 쓰는 서비스를 보유하고 있기 때문에 인공지능이 맞물려서 만들어낼 수 있는 가치가 많습니다.

인공지능을 통한 새로운 카카오의 경험을 보여주려면 그에 걸맞는 고성능 모델이 필요합니다. 카카오는 GPU 뿐 아니라 구글 클라우드의 TPU를 함께 이용해서 많은 문제들을 효과적으로 잘 풀어내고 있습니다. 최신 TPU v6e aka. TPU Trillium까지 적용해 거대 모델을 빠르게 구축한 것도 눈여겨 볼 부분입니다.

물론 TPU의 적용은 강점 만큼이나 생각할 부분도 많습니다. 카나나 모델을 만든 카카오의 라이언(류민호), 스티브(안재현) 두 AI 엔지니어와 나눈 TPU 속 진짜 이야기를 들어볼까요?

#얼마 전 if(kakaoAI)2024에서 통합 AI 모델 ‘카나나(Kanana)’가 주목을 받았습니다. 카카오는 오래 전부터 인공지능에 진심이었고, 아예 인공지능 중심의 공동체를 운영하면서 언어 모델이나 이미지 생성 AI에 공을 들여온 것 같아요.

카나나(Kanana)는 ‘가장 나다운 AI’ 라는 의미를 내포하고 있습니다. 인공지능 서비스는 여러 형태로 우리 일상에 다가오고 있잖아요. 인공지능이 다른 기술보다 더 빠르게 익숙해지는 이유는 사람을 이해하려고 한다는 점에 있어요.

기술이 발전할수록 어쩔 수 없이 어려워지는 부분들이 있습니다. 그 모든 것을 다 따라 배우는 건 부담스러운 일이고요. 언어와 소통이라는 컴퓨팅의 중요한 부분이 생성형 AI를 통해 해소되면서 할 수 있는 일들이 크게 늘어날 거라고 봐요. 서비스 곁에서 이용자들이 하려는 것들을 빠르게 거들어준다면 지금까지 카카오가 펼쳐 온 기술들을 더 넓게 활용할 수 있는 것은 물론이고 앞으로 더 적극적으로 서비스를 고도화할 수 있을 겁니다. 카나나가 사람과 세상을 이해하는 만큼 카카오의 서비스가 더 풍성해지는 것이지요.

#카나나가 이해해야 하는 대화의 범위가 상당히 클 것 같습니다. 범용 서비스에서 이용자들의 의도를 정확히 이해하려면 이를 처리하는 데에 상당한 컴퓨팅 파워가 필요할 텐데, 구글 클라우드와 TPU로 그 문제를 풀어낸 계기는 뭔가요?

구글 클라우드의 TPU(Tensor Processing Unit)는 대규모 언어 모델(LLM)의 학습을 맡았습니다. 대규모 언어 모델의 학습에는 기존 언어 모델에 비하여 대규모의 컴퓨팅 자원이 필요합니다. 카나나 언어 모델은 초기에는 GPU를 통해 개발됐어요. 카나나는 국내에서 카카오 서비스의 경험을 높이면서 글로벌로 영역을 넓히는 것이 목표였기 때문에 영어와 한글을 모두 잘 하도록 만드는 것이 아주 중요했죠. 그렇게 탄생한 모델이 Kanana-Essence-8B 모델이었습니다. 해당 모델(Kanana-Essence-8B)의 자세한 개발과정과 상세한 성능은 카카오 테크 블로그에서 확인할 수 있습니다.

하지만 여전히 모델의 성능을 더 높일 수 있는 방법을 찾아야 했습니다. 그래서 먼저 기존 모델의 레이어를 더 많이 쌓아 올려서 모델 규모를 키우는 방법인 Depth Upscaling(DUS) 방식을 선택했습니다. 특히 DUS 방법을 한 번만 적용하는 것이 아닌 반복적으로 여러 번 적용함으로써 기존 모델보다 훨씬 더 큰 모델을 상대적으로 적은 시스템 자원으로 학습시키는 것이 성능을 끌어올리는 핵심 기술이었습니다.

아래 표는 DUS 방식을 적용해 24.6B 사이즈로 테스트를 진행한 Kanana 모델의 성능입니다. MMLU는 영어 이해, KMMLU는 한국어 이해, HAE-RAE는 한국 문화, GSM8K는 수학 그리고 HumanEval 및 MBPP는 코딩에 대한 벤치마크입니다. Gemma-2-27B 모델과 비교해 Kanana-24.6B (DUS) 모델은 영어 성능을 비슷하게 끌어올리면서 한국어 성능은 훨씬 더 우수한 성능을 보였습니다. 이러한 시도를 거쳐 DUS 방식의 효용성을 검증할 수 있었고, 향후 학습이 완료된 카나나 모델에서는 수학과 코딩 실력을 더 끌어올릴 예정입니다.

추가적으로, 큰 모델을 더 적은 자원을 통해 효율적으로 운영하는 것도 놓칠 수 없었습니다. GPU는 늘 갈증을 남겼지요. 카카오 자체적으로도 여러 인공지능 연구를 위해 GPU를 확보해서 운영하고 있는데 필요한 성능을 넉넉하게 확보하는 것이 쉽지 않았어요. GPU는 늘 부족했고, 병렬 처리로 학습 모델의 규모를 키우는 데에는 GPU의 확장성(scalability)이 걸림돌이라는 생각을 지우기 어려웠습니다. 근본적으로 다른 대안이 필요했습니다. 구글 클라우드의 TPU는 그 답답함을 적절하게 해소해 주었다고 판단하고 있습니다.

자체개발 모델 외에도 각 서비스의 기능에 따라 다양한 AI 모델을 활용할 수 있도록 검증된 오픈 모델에 카카오의 데이터를 학습시켜 CPT 모델을 개발하기도 했습니다. 처음부터 70B 정도의 대규모 모델을 학습시키기 위해서는 엄청난 자원이 필요한데요. 오픈 모델과 TPU v6e를 통해 높은 수준의 한국어 대응이 가능한 모델을 빠르게 확보할 수 있었습니다.

#머신러닝 프레임워크를 바꾸는 건 쉽지 않았을 것 같은데 그 과정은 어땠을까요? GPU에서 하던 일을 그대로 TPU로 넣는 것은 아니잖아요?

네, 사실 GPU와 TPU는 밑바탕이 다르기 때문에 대규모 언어 모델을 설계하는 프레임워크도 완전히 다릅니다. GPU경우에는 GP(General Purpose) GPU이지만, TPU의 경우에는 행렬연산(Matrix multiplication)에 최적화된 머신러닝을 위한 목적으로 설계된 전용 가속기이기 때문입니다. 초기에 카나나 LLM은 자연어 처리에 특화된 Megatron-LM 프레임워크를 기반으로 학습이 이뤄졌어요. Megatron-LM은 오픈소스 프로젝트가 진행중이고, 언어 모델 학습에 대해서는 전 세계의 표준 기술처럼 자리 잡으면서 수많은 경험들이 공유되었어요.

데이터 전처리부터 학습, 추론까지 일련의 과정이 매끄럽게 이뤄져 있다 보니 초기 설계에 대한 부담은 상당히 적었습니다. 당연히 처음 고민은 TPU가 이 Megatron-LM을 대신해 줄 수 있느냐에 있었어요. 구글 클라우드와 그 부분에서 머리를 맞대었고, TPU에서 운영되는 MaxText 기반 프레임워크가 Megatron-LM의 자리를 채울 수 있다는 확신을을 갖을 수 있게 되었습니다.

실제로 모델 생성을 위한 모든 프로세스를 Google 이 제공하는 MaxText로 옮기는 것은 꽤 어려운 일이었지만 결국 원하는 결과를 만들어 낼 수 있었습니다. 기존의 모델과 학습 데이터는 쉽게 변환이 이뤄졌고, 데이터 처리와 학습 알고리즘까지 Megatron-LM과 거의 비슷했습니다.

물론 처음에는 그 다른 부분에 대해서 공부해야 하는 부분이 있습니다. 업무가 바쁘고 가야 할 길이 많은데 자료를 찾고 실제 데이터를 넣어가며 답을 찾아가는 과정이 초반에는 답답하기도 했지만 구글 클라우드의 TPU 전담팀의 기술 지원으로 답을 찾아낼 수 있었어요. 그리고 어느 수준을 넘어가면서 TPU는 전환 과정의 노력에 대한 보상(성능, 가격 등)을 확실히 주더라구요.

#TPU는 낯선 환경인데 흔히 엔지니어들 사이에서는 새롭고 낯선 기술은 모두 장애 요인이라는 우스갯소리가 있어요. TPU도 어떻게 보면 GPU로 다 잘 해 낼 수 있는데 전환에 대한 결정과 설득이 쉽지 않았을 것 같아요.

맞아요. 실제로 연구원들도 낯선 시스템, 프레임워크에 대한 불안함이 있었습니다. ‘TPU로 우리가 하던 일들을 완성할 수 있을까?’라는 막막함인 거죠. 결론부터 말하자면 지금 카나나는 TPU를 바탕으로 아주 빠르게 고도화되는 중이고, TPU가 목표를 달성하는 길을 더 쉽게 만들어 주었다고 말할 수 있어요.

아직 언어 모델을 비롯한 생성형 AI는 컴퓨팅 파워에 영향을 많이 받는데, 일정 수준 이상으로 다룰 수 있게 되면서는 GPU와 TPU를 적절히 판단하고 골라서 쓸 수 있는 기반을 마련했다는 점도 큽니다. 사실 TPU에 대한 부담은 실제 다뤄가면서 많이 사라졌습니다. 오픈소스 프레임워크는 파이토치가 많기 때문에 텐서플로, JAX 기반으로 바꾸는 과정에서 별도의 개발이 따라야 하는 부분도 있기는 합니다. 하지만 필요한 프레임워크들은 대부분 이미 준비가 되어 있었고, 개발 과정이나 기간도 큰 차이는 없었습니다. 결국은 익숙함의 차이일 뿐이라고 봅니다.

#여전히 인공지능이라고 하면 GPU를 먼저 떠올리고 기본 환경도 GPU를 중심으로 짜여 있어요. 단적으로 TPU가 GPU를 대신할 만큼의 강점은 뭔가요?

처음에는 모든 LLM 모델을 TPU v5e로 이전했습니다. 앞서 GPU의 자원 확보에 대한 이야기를 했는데, TPU는 미리 계약한 만큼의 용량(Capacity)를 확실히 보장받을 수 있습니다. 언제라도 필요한 만큼 자원을 끌어다 쓸 수 있다는 얘기죠.

저희가 학습 과정에서 많은 컴퓨팅 자원이 필요했던 것은 결국 모델의 규모를 늘리면서 정확도를 높이기 위해서였어요. 카나나의 LLM은 동일한 레이어를 계속 쌓아서 프리 트레이닝을 하면서 적은 컴퓨팅 자원으로 모델 크기를 키우는 방식인데, TPU는 이에 필요한 성능을 자연스럽게 뽑아내 주었습니다.

당장의 학습 과정에서 효과가 GPU 대비하여 명확하게 드러나기 시작했고, TPU에 익숙해질수록 TPU의 자체의 잠재력을 끌어내기 수월해졌습니다. 특히, TPU는 Jupiter 네트워킹을 통해 수 만 개 칩(Chip)을 유기적인 하나의 덩어리(Physical Pod)로 묶을 수 있어 스케일 아웃에 대한 유연성이 대단했습니다. 사실 GPU는 작은 모델을 다룰 때는 효과가 좋습니다. 하지만 GPU는 GPU가 장착 되어 있는 물리적인 서버단위로 시스템이 묶이는 한계가 있기 때문에 그 한계치인 8개 GPU를 넘어서려면 다른 시스템과 연결해야 합니다.

인피니밴드로 묶는 것도 점점 성능이 좋아지긴 하지만 물리적으로 하나의 시스템 안에 있는 것과 다른 시스템을 묶는 것은 여전히 다릅니다. 그걸 하나하나 튜닝하는 것이 또 많은 경험과 노력이 필요해요. 그리고 혹시라도 GPU가 하나라도 먹통이 되면 전체 시스템으로 그 영향이 옮겨가기도 합니다.

TPU는 최소 256개 부터 수 만 개에 이르는 모든 칩 유닛이 단일 패브릭(Pod)으로 연결되어 있기 때문에 데이터 입출력 과정에서 생기는 병목이 획기적으로 적습니다. 그래서 더 높은 성능이 필요하면 TPU 칩의 수를 더 늘리면 됩니다. 물론 칩의 개수에 완벽하게 100% 비례해서 성능이 높아지는 것은 아니지만 그에 근접하는 확장성을 보여주었어요. 특히 필요한 시스템 자원을 정확히 예측할 수 있다는 점에서는 GPU와 비교할 수 없을 만큼 좋았던 부분입니다.

#TPU의 강점들이 클라우드의 관리형 서비스와 많은 부분에서 비슷해 보이는 것 같네요. 새로 나온 GPU를 제때 구하지 못해서 많은 기업들이 고민을 하고 있잖아요.

맞아요. GPU는 자원 확보 때문에 인하우스에 시스템을 갖추는 경우가 많고, 카카오도 크게 투자해서 학습과 추론 전 과정에 GPU를 활용하고 있는데, 그 유지보수도 만만치 않습니다. 그리고 이 시스템이 상당히 비싼데, 인공지능 시장이 커지면서 GPU의 업데이트가 잦아지고 있어요.

이걸 계속 새 제품으로 바꾸는 것도 부담이 적지 않아요. 그렇다고 기존 GPU를 억지로 끌고 가는 것도 적절한 일이 아니고요. 하지만 TPU는 구글의 딥마인드에서 연구되고, 구글 클라우드에서 운영되기 때문에 관리부터 유지보수에 대한 부담이 없죠.

올해 구글 I/O에서 TPU의 v6e aka. Trillium 업데이트가 예고된 것도 인상적이었습니다. TPU Trillium은 HBM의 용량이나 대역폭이 두 배 늘어나서 256개 칩이 메모리 병목 없이 데이터를 처리합니다. 그래서 자연스럽게 모델이 커져도 성능을 꾸준히 유지할 수 있다는 강점도 뒤따릅니다. 전반적으로 v6e는 v5e에 비해 더 낮은 비용으로 더 많은 일을 해낼 수 있다는 부분이 보였습니다. 구글도 Gemini 2.0을 TPU v6e로 학습했다는 점도 의미가 있고 Trillium이라는 멋진 이름도 마음에 들었어요.

TPU Trillium은 얼마 전에 공개되긴 했지만 저희는 미리 써볼 수 있는 기회를 얻었습니다. 기술적으로 TPU v5e를 TPU Trillium으로 전환하는 것을 늦출 이유가 없죠. 서둘러서 전환을 결정했습니다. 사실 TPU Trillium으로 업그레이드하는 과정은 아주 수월했습니다. 프레임워크와 기본적인 기능들을 그대로 활용할 수 있기 때문에 마치 PC에서 CPU를 바꿔 꽂는 것처럼 모든 과정이 매끄러웠어요.

실제로 적용해 보니 분산 처리를 비롯해 더 큰 규모의 모델을 다룰 수 있을 것이라는 점들이 눈에 보이기 시작했습니다. TPU Trillium은 256개의 유닛이 하나처럼 매끄럽게 병렬 연산을 처리해 냈고 별다른 처리 없이도 곧바로 기존 v5e에 비해 3배 가량 더 높은 성능을 보여줬습니다. 더욱이 TPU Trillium 1개의 Pod(Full Pod, 256 chips)의 성능 이상이 필요하다면, TPU Multi-Slice 를 통해 2개 이상의 Pod 로 확장할 수 있습니다. 더 높은 파라미터를 갖는 대규모 언어 모델의 학습 성능에 맞출 수 있는 것이죠.

생각해 보면 인프라 관리에 대한 고민 없이 TPU의 업데이트 주기에 맞춰 더 나은 모델 학습 환경이 저절로 따라오는 건데요, 이런 수준으로 업데이트가 이어진다고 하면 앞으로의 서비스를 더 여유있게 그려볼 수 있겠다는 기대가 생기더라고요.

#인공지능 서비스의 모델 크기가 커지면서 TPU도 인공지능 서비스를 꾸리는 데에 효과적인 방법으로 자리를 잡아가는 것 같습니다. TPU에 대한 기술적 기대나 관심이 많은데 낯설다는 점 때문에 접근하지 못하는 경우가 많은데 돌아보면 카카오는 일찍부터 TPU와 친하게 지내 온 게 지금까지 이어 온 것 같네요.

TPU는 똑똑하지만 GPU만큼 많은 부분에 두루 쓸 수 있는 환경을 제공하는 건 아닙니다. 그리고 TPU가 GPU를 완전히 대체할 수 있는 것도 아니죠. 이걸 단점으로 볼 게 아니라 TPU가 다른 특성을 갖고 있다고 받아들이는 게 필요하다고 생각해요. TPU는 머신 러닝을 위한 행렬연산(Matrix multiplication)에 최적화되어 있고, TPU가 GPU대비하여 가격 및 성능 관점에서 효율적인 부분이 있습니다. 결국, TPU로 더욱 더 잘 할 수 있는 영역을 잘 파악하는 게 중요한 거죠.

저희 카카오는 지난 TPU v3 시절부터 구글과 함께 경험하면서 TPU의 이러한 특성을 이해하려는 노력을 해 왔었습니다. 돌아보면 카카오브레인(현, 카카오) 시절부터 TPU를 이용해 KoGPT 언어 모델을 개발했고, 이미지 생성 AI인 칼로(Karlo)를 통해 TPU와 멀티모달 데이터셋의 궁합을 고민하기도 했어요. 그 경험들을 통해 TPU의 특성을 이해하기 시작했고, 성능과 개발 환경이 자리를 잡아가는 Trillum(TPU v6e)을 통해서 본격적으로 효과를 보여주는 것 같아요.

그래서 지금은 GPU와 TPU에 적절한 업무를 할당해서 인공지능 모델들을 다룰 수 있게 됐다고 자부합니다. 서비스를 설계할 때 고를 수 있는 선택지가 늘어난 셈이죠. TPU는 아직 낯설고 모두에게 활짝 열려 있는 것은 아니지만 그 강점이 명확하고 점차 그 활용 영역이 넓어질 것이라고 확신합니다.

GPU에 익숙한 환경이라면 TPU는 첫 발을 내딛는 게 어려운 것은 사실입니다. 구글도 적극적으로 기술 지원을 해 주고 있지만 가장 중요한 것은 TPU로 할 수 있는 일을 찾고 효과적인 결과물을 만들어 내겠다는 의지가 아닐까요. TPU가 잘 할 수 있는 일을 찾아내고 적절한 모델을 구성한다면 TPU는 같은 시간 동안 더 높은 생산성과 효율적인 비용이라는 두 마리 토끼를 잡을 수 있는 방법이라고 생각합니다.

우리가 원하는 답을 얻는 데까지 가는 길은 한 가지가 아닙니다. 인공지능 기술 만큼이나 GPU, 뉴럴 프로세서 등 이를 처리하는 하드웨어도 빠르게 바뀌고 있습니다. TPU는 모델을 처리하는 새로운 방법을 제시해 주었고 카나나에 대한 가능성을 폭넓게 바라볼 수 있게 됐습니다. TPU로 학습된 카나나(Kanana)가 본격적으로 서비스되면서 바꾸어 놓을 카카오 서비스들도 지켜봐주세요.