SkewAdam 발표: MoE 학습 시 옵티마이저 메모리 사용량 97% 절감
MoE 파라미터 최적화를 통해 6.7B 모델을 40GB GPU에 올릴 수 있음. 학습 효율화에 필수적인 기술임.
요약
Mixture-of-Experts(MoE) 모델 학습 시 발생하는 방대한 VRAM 점유 문제를 해결하기 위한 새로운 옵티마이저 'SkewAdam'이 공개되었습니다. 기존 AdamW 옵티마이저는 12.6GB 모델 학습 시 옵티마이저 상태값만 50.6GB를 소모하지만, SkewAdam은 계층적 상태 할당 방식을 통해 메모리 사용량을 97%까지 절감합니다. 구체적으로 전체 파라미터의 5%인 백본에는 모멘텀과 인수분해된 2차 모멘트를, 95%를 차지하는 전문가(Experts) 층에는 인수분해된 2차 모멘트만을 적용합니다. 0.01% 미만인 라우터에는 정확한 2차 모멘트를 할당하여 성능을 유지하면서도 효율성을 극대화했습니다. 이 방식을 통해 40GB GPU 환경에서 6.7B 규모의 MoE 모델을 학습할 수 있게 되었습니다.
AI가 원문을 요약한 내용으로, 부정확할 수 있습니다.
원문 제목 SkewAdam: A tiered optimizer that cuts MoE state memory by 97% (fits a 6.7B MoE on a 40GB GPU) [R]
원문 보기 ↗