논문: https://arxiv.org/abs/2607.19058 코드 (GitHub): https://github.com/nuemaan/skewadam

안녕하세요 여러분, 저는 Mixture-of-Experts (MoE) 학습 시 발생하는 거대한 VRAM 병목 현상을 해결하기 위해 설계된 새로운 옵티마이저 (optimizer)에 관한 프리프린트 (preprint)를 방금 발표했습니다. MoE를 학습시켜 보셨다면, 옵티마이저 상태 (optimizer state)가 일반적으로 메모리 예산에서 가장 큰 단일 항목이라는 것을 알고 계실 것입니다. 예를 들어, AdamW는 단지 12.6 GB 모델을 업데이트하기 위해 50.6 GB의 상태 메모리를 사용합니다.

저는 계층형 상태 할당 (tiered state allocation) 방식을 사용하여 이 문제를 해결하고자 SkewAdam을 구축했습니다. 모든 파라미터 (parameter)를 동일하게 취급하는 대신, 파라미터의 동작에 따라 정밀도 (precision)를 할당합니다:

• 백본 (Backbone, 파라미터의 5%): 모멘텀 (Momentum) + 인수분해된 2차 모멘트 (Factored 2nd moment)
• 전문가 (Experts, 파라미터의 95%): 인수분해된 2차 모멘트 (Factored 2nd moment)만 사용
• 라우터 (Router, 파라미터의 0.01% 미만): 정확한 2차 모멘트 (Exact 2nd moment)

하드웨어 결과:
옵티마이저 상태 메모리가 50.6 GB에서 1.29 GB로 감소합니다 (97.4% 절감). 피크 학습 메모리 (Peak training memory)는 81.4 GB에서 31.3 GB로 감소합니다. 이를 통해 수렴 (convergence) 성능이나 라우터 안정성을 희생하지 않고도 6.78B MoE를 단일 40GB GPU에 여유롭게 올릴 수 있습니다.