카카오, Kanana-2 모델 공개: MoE, Muon, FP8 기술로 LLM 훈련 효율 극대화

카나나-2는 전문가 혼합(MoE, Mixture of Experts) 아키텍처를 채택하여, 전체 파라미터 크기는 32B(320억 개)에 달하는 거대 모델의 지능을 보유하면서도, 실제 추론 단계에서는 3B(30억 개) 파라미터만 활성화하여 연산 효율을 극대화했다. MoE는 모델의 크기를 늘리면서도 계산 비용을 절감할 수 있는 효과적인 방법이다. 하지만, MoE는 각 전문가(Expert)를 학습하고 라우팅(Routing)하는 과정에서 추가적인 복잡성이 발생하며, 데이터 불균형(Data Imbalance) 문제에 취약하다는 단점이 있다. 카카오는 MoE 학습에 필요한 여러 커널들을 직접 개발하여 성능 손실 없이 학습 속도를 높이면서 메모리 사용량을 획기적으로 줄였다.

Kanana-2는 기존 AdamW 대신 Muon Optimizer를 사용했다. Muon은 Gradient를 Orthogonalize(직교화)하여 Parameter Update Matrix로 사용하며, Newton-Schulz와 같은 Iterative Algorithm을 사용한다. 특히, Polar Express를 사용하여 더 정확한 Orthogonalization을 가능케 했다. Polar Express는 작은 Scale에서는 Newton-Schulz 대비 유의미한 이득을 보이지 않았지만, Large-scale Training의 후반부로 갈수록 Weight Update의 Noise를 줄이는 것이 중요하다고 판단하여 채택했다. MuonClip을 통해 Attention Logit 발산(Attention Logit Explosion)을 방지하고 훈련 안정성을 확보했다.

카카오는 FP8 Tensor Core를 활용한 훈련을 위해 DeepSeek V3에서 제안된 Fine-grained Scaling을 채택했다. Fine-grained Quantization은 Input Activation은 1x128 Tile, Weight는 128x128 Block 단위로 양자화를 진행한다. DeepSeek V3의 Fine-grained Quantization을 사용하기 위해, GEMM과 Non-GEMM 커널들을 모듈화된 방식으로 조합하고, 필요한 경우 커널을 직접 구현했다. 특히, MoE Quantization 커널을 개발하여 Weight과 Activation 모두에 대한 Quantization을 Expert별로 순차적으로 실행되는 문제를 해결했다. 또한, Forward에 필요한 FP8 Quantized Weight만을 저장하고 Backward 시에는 Block Quantize된 Weight들을 Transpose(전치)하는 커널을 개발하여 GPU 메모리 사용량(GPU Memory Usage)을 절감했다.

수천 억 파라미터 규모의 모델 훈련에서 HP Sweeping은 현실적으로 불가능하므로, 카카오는 2-step HP Transfer를 고안했다. 먼저 Proxy Scale에서 HP의 Scaling Law를 탐색하되, 모델 크기(M)는 고려하지 않고 데이터 양(D) 축으로의 Scaling만 고려한다. D 축으로의 Transfer를 통해 Target D, Proxy M Scale에서의 Optimal LR을 결정한다. 이를 기반으로 MuP를 통해 최종적으로 Target D, Target M Scale에서의 Optimal LR을 결정한다. Batch Size에 대한 Transfer는, Batch Size가 M과 D에 대해 어떻게 Scaling되는가에 대한 연구들의 통일된 결론을 내리기 어려워 고려하지 않았다. 이 전략을 통해 훈련 비용(Training Cost)을 절감하고, 효율적인 HP 탐색을 가능하게 했다.

Large-scale Training에서의 최적화는 GPU를 사용하는 연산에 국한되지 않는다. Multi-GPU Training에서 GPU별 Garbage Collection(GC)의 타이밍이 다르면, 다른 GPU들이 GC GPU를 기다리게 되는 현상이 발생한다. Python의 Automatic GC를 끄고 Manual하게 GC Interval을 설정함으로써 Multi-GPU GC가 Synchronous하도록 개선한 결과, 7% 정도의 성능 향상을 확인했다. 또한, Distributed Asynchronous Checkpointing을 도입하여 체크포인팅에 소요되는 시간을 줄였다. 체크포인팅이 차지하는 스토리지 문제를 해결하기 위해 prefect를 사용하여 체크포인팅 이후의 워크플로우를 자동화했다. 이러한 노력들을 통해 훈련 효율성(Training Efficiency)을 극대화했다.

카카오는 모델 아키텍처 변경에 따른 추론/배포 시의 여러 하드웨어와 프레임워크까지 고려해야 하는 점을 감안하여, DeepSeek V3 아키텍처를 사용했다. 작은 스케일에서의 Data Ablation 결과가 큰 스케일로 Transfer되지 않는 현상과 같이, 실험 결과의 Scalability와 Generalizability를 고려하여 의사결정을 내렸다. Scalable하고 더 간단한 선택지를 우선적으로 고려했으며, Polar Express와 MuP가 그 좋은 예시이다. 데이터, 연구에서 엔지니어링에 이르기까지 모든 디테일에 대한 검증을 진행할 수 없다는 점을 인정하고, 여러 아이템 중 우선순위를 정하는 것이 중요했다. 이러한 전략을 통해 자원 효율성(Resource Efficiency)을 극대화하고, 더 나은 모델 개발을 위한 기반을 마련했다.