대규모 언어 모델(LLM)의 효율적인 훈련은 여전히 큰 과제이며, 특히 혼합 전문가(MoE) 모델은 방대한 매개변수로 인해 막대한 메모리를 요구합니다. 최근 공개된 SkewAdam이라는 새로운 최적화 기법은 이러한 MoE 모델 훈련의 메모리 문제를 혁신적으로 해결하며, 훈련 비용과 접근성을 크게 낮출 잠재력을 보여주고 있습니다.
SkewAdam은 MoE 모델의 각기 다른 구성 요소(백본, 전문가, 라우터)가 균일하지 않다는 점에 착안하여, 각 부분에 필요한 최적화 상태(optimizer state)를 차등적으로 할당합니다. 예를 들어, 67.8억 매개변수 MoE 모델 훈련 시 AdamW는 50.6GB의 최적화 상태 메모리를 필요로 하지만, SkewAdam은 이를 1.29GB로 97.4% 절감합니다. 이로 인해 피크 훈련 메모리는 81.4GB에서 31.3GB로 줄어들어, 단일 40GB GPU에서도 대규모 MoE 모델 훈련이 가능해집니다. SkewAdam은 백본(embeddings, attention, dense FFN)에는 운동량(momentum)과 인자화된(factored) 2차 모멘트를, 전문가(experts)에는 인자화된 2차 모멘트만, 그리고 라우터(router)에는 전체 정밀도의 2차 모멘트를 할당하는 전략을 사용합니다.
이러한 메모리 효율성뿐만 아니라, SkewAdam은 기존 AdamW, Muon, Lion과 같은 최적화 기법들보다 더 낮은 검증 퍼플렉시티(perplexity)를 달성하며 우수한 성능을 입증했습니다. 이는 대규모 MoE 모델을 더 적은 자원으로 훈련할 수 있게 함으로써, 연구자와 개발자들이 고성능 모델을 구축하는 데 필요한 장벽을 낮추는 데 기여할 것입니다. 특히 개인 개발자나 소규모 팀에게는 값비싼 고사양 GPU 여러 대 없이도 최신 MoE 모델을 실험하고 개발할 수 있는 기회를 제공하여, AI 기술의 민주화에 중요한 역할을 할 것으로 기대됩니다.