[논문리뷰] Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training

링크: 논문 PDF로 바로 열기

저자: Nuemaan Malik

---

1. Key Terms & Definitions (핵심 용어 및 정의)

• SkewAdam: 파라미터의 역할(Backbone, Experts, Router)에 따라 최적화 상태(Optimizer state)를 차등 할당하는 메모리 효율적 옵티마이저 프레임워크입니다.
• Tiered State Allocation: 파라미터 그룹별로 서로 다른 통계적 중요도와 밀도를 고려하여 Momentum 버퍼와 Second-moment 추정 방식을 차별적으로 적용하는 전략입니다.
• Factored Second Moments: 고차원 행렬의 2차 모멘트를 행(row)과 열(column) 벡터의 곱으로 분해하여 저장함으로써 Memory 사용량을 획기적으로 줄이는 기법입니다.
• MoE (Mixture-of-Experts): 전체 파라미터 중 일부만을 활성화하여 연산 효율을 높이는 구조로, 대규모 모델에서 Optimizer state가 전체 메모리 점유율의 가장 큰 비중을 차지하는 모델 아키텍처입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 Mixture-of-Experts (MoE) 모델 학습 시 발생하는 방대한 Optimizer state 메모리 점유 문제를 해결하기 위해 고안되었습니다. 기존 옵티마이저(예: AdamW)는 모델의 전체 파라미터를 동일한 방식으로 취급하여, 드물게 활성화되는 Expert 파라미터에도 불필요하게 많은 float32 메모리를 할당합니다. 이러한 균일한 할당 정책은 6.78B-parameter 모델 기준 Optimizer state로만 약 50.6 GB를 소모하게 하여 학습 효율을 저해합니다 [Figure 1]. 저자들은 Backbone, Experts, Router의 통계적 특성이 서로 다르다는 점에 착안하여, 각 계층에 적합한 최적화 자원을 선택적으로 배분하는 새로운 접근 방식의 필요성을 제시합니다.

Figure 1 — 계층별 최적화 상태 할당 구조

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 논문은 SkewAdam이라는 계층적 상태 할당 정책을 제안하며, 이를 통해 Optimizer state를 기존 대비 97.4% 감소시키는 데 성공하였습니다 [Figure 1]. Backbone 계층에는 float32 momentum과 Factored second moment를 할당하고, Expert 계층에는 메모리 점유가 큰 Momentum을 제거하고 Factored second moment만 적용하며, Router에는 Exact second moment를 할당하여 학습 안정성을 유지합니다. 실험 결과, SkewAdam은 6.78B-parameter MoE 모델 학습에서 Peak training memory를 기존 81.4 GB에서 31.3 GB로 대폭 낮추면서도, 검증 데이터 기준 108.4의 Perplexity를 기록하여 AdamW (126.8), Muon (120.2), Lion (393.7) 대비 우수한 성능을 입증하였습니다 [Table 1]. 특히, 제안된 계층화 할당 정책은 동일한 Perplexity를 유지하면서도 Uniform allocation 방식 대비 20배 적은 Optimizer state를 사용한다는 점이 핵심입니다.

4. Conclusion & Impact (결론 및 시사점)

본 연구는 MoE 모델 학습 시 옵티마이저 상태를 계층별로 차등 할당하는 전략이 메모리 효율성을 극대화하면서도 학습 성능을 보존할 수 있음을 입증하였습니다. 이 연구는 대규모 모델 학습에서 단순히 메모리 총량을 줄이는 것보다, 파라미터별로 최적화 자원을 전략적으로 배치하는 설계 원칙이 중요함을 시사합니다. 향후 더 깊은 모델 구조와 다양한 학습 환경에서도 본 정책의 유효성을 검증함으로써, 리소스 제약이 있는 환경에서의 거대 모델 학습 대중화에 기여할 것으로 기대됩니다.

Figure 2 — 학습 수렴도 및 메모리 비교

Figure 3 — 라우팅 안정성 비교

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.