연구자들은 SkewAdam을 소개했는데, 이는 전문가 혼합(MoE) 모델을 위해 설계된 옵티마이저로, 서로 다른 매개변수 집단에 서로 다른 상태 유형을 할당하여 메모리 사용을 극적으로 줄입니다. SkewAdam은 백본에는 float32 모멘텀과 분할된 2차 모멘트를, 전문가에는 분할된 2차 모멘트를, 라우터에는 정확한 2차 모멘트를 할당함으로써 6.78B 매개변수 모델에 대해 옵티마이저 상태를 50.6 GB에서 1.29 GB로 줄입니다.

  • 최대 학습 메모리는 81.4 GB에서 31.3 GB로 감소하여 40 GB 가속기 예산 내에 적합합니다.
  • SkewAdam은 82M 토큰에 걸쳐 검증 퍼플렉시티 108.4를 달성하여 AdamW(126.8), Muon(120.2), Lion(393.7)보다 우수한 성능을 보입니다.
  • 메모리 절감은 정확도 손실 없이 이루어지며, 성능 향상은 다단계 할당 자체보다는 모멘텀 유지에서 비롯됩니다.

이 결과는 옵티마이저 상태의 할당 위치가 그 양만큼이나 중요함을 시사하며, 제한된 메모리를 가진 하드웨어에서 효율적인 학습을 가능하게 합니다.