SkewAdam: MoE 학습에서 옵티마이저 상태 메모리를 계층적으로 줄여 6.78B 모델을 단일 40GB GPU에 탑재한 옵티마이저

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

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 코드는 재현 가능한 근거를 제공하며 메모리 절감과 수렴성·라우터 안정성의 균형을 유지한 설계라는 의미를 가진다.

섹션별 상세

01

대형 MoE 학습에서 옵티마이저 상태가 VRAM 병목이 되는 문제가 핵심이었다. 글은 옵티마이저 상태가 모델 파라미터보다 메모리를 훨씬 더 소비하는 관찰에서 출발하며 입력으로는 전체 모델 파라미터와 기존 옵티마이저의 텐서 집합을 받아 문제를 정의했다. 제시된 근거로 AdamW의 옵티마이저 상태가 50.6GB에 달한다는 수치가 제공되었으며 이는 단일 GPU 메모리 한계를 초과해 분산 설계나 오프로드를 강제한다고 결론지었다.

02

SkewAdam의 작동 원리는 파라미터 그룹별로 옵티마이저 상태의 정밀도와 보관 방식을 차등화하는 계층적 할당 정책이다. 입력으로는 파라미터 집합을 받고 처리 단계에서 Backbone 파라미터에는 Momentum과 팩토링된 2차 모멘트를 결합해 저장하며 Experts에는 팩토링된 2차 모멘트만 유지하고 Router에는 정확한 2차 모멘트를 유지하는 방식으로 출력은 각 그룹별로 축소된 옵티마이저 상태이다. 구현 근거로 Backbone이 모델의 5%를 차지하고 Experts가 95%를 차지한다는 가정 하에 팩토링 기법을 적용해 메모리를 극적으로 줄였다는 기술적 선택이 제시되었다.

03

하드웨어 결과는 메모리 감소 수치로 구체화되었다. 옵티마이저 상태가 50.6GB에서 1.29GB로 97.4% 감소했고 피크 VRAM은 81.4GB에서 31.3GB로 떨어졌다는 측정값이 게시되어 단일 40GB GPU에서 6.78B MoE 학습이 가능해졌다고 보고되었다. 이 수치는 분산이나 옵티마이저 오프로드 없이도 더 큰 모델을 단일 장비에서 다룰 수 있다는 실질적 의미를 가지며 하드웨어 설계와 비용 구조에 직접적인 영향을 미친다.

04

성능과 안정성 측면에서는 수렴성과 라우터 안정성이 희생되지 않았다는 주장이 포함되었다. 글은 수렴이나 라우터 불안정성 문제 없이 메모리 절감 효과를 달성했다고 명시하고 있으며 그 근거로 실험 결과와 함께 arXiv 논문과 GitHub 코드 링크를 제공해 재현 가능성을 확보했다. 이로 인해 실제 적용 시 옵티마이저 상태 절감과 모델 동작 안정성 사이의 트레이드오프를 줄일 가능성이 제기되었다.

이미지 분석

Chart
이 그래프는 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이 옵티마이저 상태를 크게 줄인 수치를 보여준다.

Chart
이 그림은 Steps 축에 따라 Validation Perplexity가 감소하는 경향을 로그 스케일로 표시해 네 옵티마이저의 수렴 속도와 최종 퍼플렉시티 수준을 비교한다. SkewAdam(파란색)은 학습 초반과 중후반 모두에서 Muon과 AdamW에 비해 비슷하거나 더 낮은 퍼플렉시티를 보이며 수렴 성능이 희생되지 않았다는 주장을 시각적으로 뒷받침한다. 해당 플롯은 메모리 절감 조치가 실제 최적화 성능과 안정성에 미치는 영향을 검증하는 근거로 사용될 수 있다.
훈련 단계별 검증 퍼플렉시티를 네 옵티마이저로 비교한 로그 스케일 선그래프로 수렴 경향을 보여준다.

용어 해설

혼합 전문가(MoE)(Mixture-of-Experts): — Mixture-of-Experts는 여러 개의 전문화된 서브네트워크(Experts)와 입력을 적절한 Expert로 분배하는 Router로 구성된 아키텍처로, 일부 파라미터만 선택적으로 활성화하여 모델의 용량을 키우면서 연산 비용을 제한하는 방식으로 작동한다. 입력 토큰마다 Router가 활성화할 Expert를 선택하고 선택된 Expert들만 계산 자원을 소비하므로 대형 모델에서 계산 효율성과 파라미터 효율성 사이의 트레이드오프를 조정하는 데 중요하다. 이 논문 맥락에서는 Expert의 대규모 파라미터 집합 때문에 옵티마이저 상태 메모리가 전체 메모리 병목이 되는 현상이 핵심 문제로 제기되었다.
옵티마이저 상태(Optimizer State): — 옵티마이저 상태는 각 파라미터에 대해 유지되는 1차 및 2차 모멘트나 모멘텀 등의 추가 텐서들을 의미하며, 특히 Adam 계열에서는 각 파라미터마다 여러 텐서를 저장해 모델 파라미터 자체보다 더 많은 메모리를 차지할 수 있다. 옵티마이저 상태 크기는 배치 처리, 체크포인트, 분산 학습 설계에 직접적인 영향을 미치므로 대규모 모델 학습에서 VRAM 사용 최적화의 핵심 대상이 된다. 이 연구는 옵티마이저 상태를 파라미터 특성에 따라 계층적으로 줄이는 방식으로 VRAM 병목을 해소하는 것을 목적으로 삼았다.
팩토링된 2차 모멘트(Factored 2nd Moment): — 팩토링된 2차 모멘트는 전체 2차 모멘트 텐서를 완전하게 저장하는 대신 행렬 분해나 축별 축소를 통해 저차원 표현으로 압축해 저장하는 기법으로, 저장 공간을 줄이되 2차 정보의 근사치를 유지한다. 이 방식은 파라미터 수가 매우 큰 Expert 영역에서 메모리 비용을 줄이면서 적절한 학습 안정성을 확보하는 데 사용될 수 있다. 논문에서는 Expert 파라미터에 대해 2차 모멘트를 팩토링해 옵티마이저 상태를 크게 축소하는 것이 핵심 전략으로 제시되었다.
라우터(Router): — 라우터는 입력을 적절한 Expert로 배분하는 MoE 구성 요소로, 각 입력에 대해 어떤 Expert가 활성화될지를 결정하는 확률적 또는 결정적 매핑을 수행한다. 라우터의 출력을 안정적으로 유지하지 못하면 특정 Expert에 부하가 치우치거나 학습이 불안정해질 수 있으므로 라우터의 수치적 정확성 및 업데이트 안정성은 MoE 학습에서 중요한 고려사항이 된다. 본 연구는 라우터 파라미터에 대해 정확한 2차 모멘트를 유지하는 예외적 정책을 둬 라우터 안정성을 보존했다.

코드 예제

text

언급된 도구

GitHub repository중립링크

논문 구현 및 코드 제공

언급된 리소스

문서arXiv preprint

GitHubSkewAdam GitHub