SkewAdam: MoE 학습의 옵티마이저 상태를 어디에 둘 것인가

도입

Mixture-of-Experts(MoE)는 활성화되는 일부 전문가만 사용해 큰 파라미터 규모를 다룰 수 있다는 점에서 매력적이다. 하지만 학습 메모리 관점에서는 여전히 부담이 크다. 가중치와 activation뿐 아니라, 옵티마이저가 저장하는 상태가 전체 메모리 예산에서 매우 큰 비중을 차지할 수 있기 때문이다.

논문 “Where Should Optimizer State Live?”는 이 문제를 직접 겨냥한다. 보고된 6.78B 파라미터 MoE 언어모델에서 bfloat16 가중치는 12.6GB인 반면, AdamW가 유지하는 1차·2차 모멘트 상태는 50.6GB에 달한다. SkewAdam은 MoE 내부의 dense backbone, experts, router가 크기와 그래디언트 통계가 다르므로 동일한 옵티마이저 상태를 줄 필요가 없다는 관찰에서 출발한다.

핵심 내용

• 파라미터 역할별로 상태를 다르게 배치한다. SkewAdam은 약 5%의 파라미터를 차지하는 dense backbone에는 float32 momentum과 factorized second moment를 둔다. 약 95%를 차지하는 experts에는 factorized second moment만 사용하고, 0.01% 미만인 router에는 exact second moment를 유지한다.
• 메모리 절감 폭이 크다. 옵티마이저 상태는 AdamW의 50.6GB에서 1.29GB로 줄어든다. 이는 AdamW 대비 2.6% 수준이다. 전체 피크 학습 메모리도 81.4GB에서 31.3GB로 낮아져 40GB 가속기 예산 안에 들어간다.
• 검증 성능도 뒤처지지 않는다. 동일 초기화와 82M tokens 조건의 비교에서 SkewAdam은 validation perplexity 108.4를 기록했다. 이는 AdamW 126.8, Muon 120.2, Lion 393.7보다 좋은 수치로 제시됐다.
• 정확도와 메모리 효과를 구분한다. 논문은 계층형 할당이 주로 메모리 절감의 원천이며, 성능 측면에서는 momentum을 유지하는 것이 중요하다고 설명한다. 비슷한 factorized estimator를 쓰지만 momentum을 제거한 Adafactor는 큰 격차를 보였다.

의미와 영향

SkewAdam의 핵심은 단순히 옵티마이저 상태를 압축하는 데 있지 않다. MoE 구조의 불균형을 인정하고, 각 구성요소에 맞는 상태를 배치한다는 점이 중요하다. 전문가 파라미터는 압도적으로 많고, 라우터는 매우 작지만 부하 균형에 민감하며, 백본은 모델의 공통 기반을 담당한다. 이런 요소들을 모두 같은 방식으로 처리하는 것은 메모리 측면에서 비효율적일 수 있다.

이 접근은 MoE 학습 시스템 설계에 실용적인 시사점을 준다. 더 큰 GPU를 기다리기보다, 어떤 파라미터에 어떤 옵티마이저 정보를 저장해야 하는지 다시 설계하면 수십 GB 수준의 차이를 만들 수 있다. 특히 제한된 메모리 환경에서 다십억 파라미터 MoE를 학습하려는 팀에 유용한 방향이다.

다만 현재 결과는 특정 모델 규모와 학습 조건에서 나온 것이다. 더 큰 스케일, 더 긴 학습, 다른 MoE 아키텍처에서도 같은 결론이 유지되는지는 추가 검증이 필요하다. 그럼에도 이 논문은 MoE 학습에서 옵티마이저 상태의 “양”뿐 아니라 “위치”가 중요하다는 점을 분명히 보여준다.

출처: Hugging Face Daily Papers