TL;DR
대형 MoE 학습에서 옵티마이저 상태가 VRAM 병목이 되어 모델 학습을 어렵게 만드는 문제를 해결하기 위해 SkewAdam은 파라미터 그룹별로 옵티마이저 상태의 정밀도와 저장 방식을 계층적으로 달리하여 Backbone에는 Momentum과 팩토링된 2차 모멘트를, Experts에는 팩토링된 2차 모멘트만, Router에는 정확한 2차 모멘트를 유지하는 정책을 적용했다. 이 접근으로 옵티마이저 상태 메모리는 AdamW의 50.6GB에서 1.29GB로 97.4% 감소했고 피크 VRAM은 81.4GB에서 31.3GB로 줄어 단일 40GB GPU에서 6.78B MoE 학습이 가능해졌다는 실험 수치가 보고되었다. 논문과 GitHub 코드는 재현 가능한 근거를 제공하며 메모리 절감과 수렴성·라우터 안정성의 균형을 유지한 설계라는 의미를 가진다.
섹션별 상세
이미지 분석
이 그래프는 SkewAdam, Lion, Muon, AdamW 네 옵티마이저의 피크 VRAM(파란색)과 옵티마이저 상태(주황색)를 병렬로 제시해 메모리 사용 구조를 비교한다. SkewAdam의 옵티마이저 상태가 1.29GB로 표기되어 AdamW의 50.55GB와 비교했을 때 97% 수준의 감소가 시각적으로 확인되며 피크 VRAM도 81.4GB에서 31.3GB로 감소한 수치가 레이블로 표시되어 있다. 그래프에는 40GB GPU 한계선이 점선으로 표시되어 SkewAdam만 단일 40GB GPU에 맞춰진다는 결론을 시사한다.
각 옵티마이저별 Peak VRAM과 옵티마이저 상태 메모리를 비교한 막대그래프이며 SkewAdam이 옵티마이저 상태를 크게 줄인 수치를 보여준다.
이 그림은 Steps 축에 따라 Validation Perplexity가 감소하는 경향을 로그 스케일로 표시해 네 옵티마이저의 수렴 속도와 최종 퍼플렉시티 수준을 비교한다. SkewAdam(파란색)은 학습 초반과 중후반 모두에서 Muon과 AdamW에 비해 비슷하거나 더 낮은 퍼플렉시티를 보이며 수렴 성능이 희생되지 않았다는 주장을 시각적으로 뒷받침한다. 해당 플롯은 메모리 절감 조치가 실제 최적화 성능과 안정성에 미치는 영향을 검증하는 근거로 사용될 수 있다.
훈련 단계별 검증 퍼플렉시티를 네 옵티마이저로 비교한 로그 스케일 선그래프로 수렴 경향을 보여준다.
용어 해설
- Mixture-of-Experts
- — Mixture-of-Experts는 여러 개의 전문화된 서브네트워크(Experts)와 입력을 적절한 Expert로 분배하는 Router로 구성된 아키텍처로, 일부 파라미터만 선택적으로 활성화하여 모델의 용량을 키우면서 연산 비용을 제한하는 방식으로 작동한다. 입력 토큰마다 Router가 활성화할 Expert를 선택하고 선택된 Expert들만 계산 자원을 소비하므로 대형 모델에서 계산 효율성과 파라미터 효율성 사이의 트레이드오프를 조정하는 데 중요하다. 이 논문 맥락에서는 Expert의 대규모 파라미터 집합 때문에 옵티마이저 상태 메모리가 전체 메모리 병목이 되는 현상이 핵심 문제로 제기되었다.
- Optimizer State
- — 옵티마이저 상태는 각 파라미터에 대해 유지되는 1차 및 2차 모멘트나 모멘텀 등의 추가 텐서들을 의미하며, 특히 Adam 계열에서는 각 파라미터마다 여러 텐서를 저장해 모델 파라미터 자체보다 더 많은 메모리를 차지할 수 있다. 옵티마이저 상태 크기는 배치 처리, 체크포인트, 분산 학습 설계에 직접적인 영향을 미치므로 대규모 모델 학습에서 VRAM 사용 최적화의 핵심 대상이 된다. 이 연구는 옵티마이저 상태를 파라미터 특성에 따라 계층적으로 줄이는 방식으로 VRAM 병목을 해소하는 것을 목적으로 삼았다.
- Factored 2nd Moment
- — 팩토링된 2차 모멘트는 전체 2차 모멘트 텐서를 완전하게 저장하는 대신 행렬 분해나 축별 축소를 통해 저차원 표현으로 압축해 저장하는 기법으로, 저장 공간을 줄이되 2차 정보의 근사치를 유지한다. 이 방식은 파라미터 수가 매우 큰 Expert 영역에서 메모리 비용을 줄이면서 적절한 학습 안정성을 확보하는 데 사용될 수 있다. 논문에서는 Expert 파라미터에 대해 2차 모멘트를 팩토링해 옵티마이저 상태를 크게 축소하는 것이 핵심 전략으로 제시되었다.
- Router
- — 라우터는 입력을 적절한 Expert로 배분하는 MoE 구성 요소로, 각 입력에 대해 어떤 Expert가 활성화될지를 결정하는 확률적 또는 결정적 매핑을 수행한다. 라우터의 출력을 안정적으로 유지하지 못하면 특정 Expert에 부하가 치우치거나 학습이 불안정해질 수 있으므로 라우터의 수치적 정확성 및 업데이트 안정성은 MoE 학습에서 중요한 고려사항이 된다. 본 연구는 라우터 파라미터에 대해 정확한 2차 모멘트를 유지하는 예외적 정책을 둬 라우터 안정성을 보존했다.
코드 예제
언급된 도구
논문 구현 및 코드 제공
언급된 리소스
AI 요약 · 북마크 · 개인 피드 설정 — 무료
출처 · 인용 안내
인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.