Исследователи представляют SkewAdam — оптимизатор, разработанный для моделей mixture-of-experts (MoE), который распределяет различные типы состояния среди разных групп параметров, чтобы радикально сократить использование памяти. Назначая импульс float32 и факторизованные вторые моменты для основной сети, факторизованные вторые моменты для экспертов и точные вторые моменты для маршрутизатора, SkewAdam уменьшает состояние оптимизатора с 50.6 ГБ до 1.29 ГБ для модели с 6.78B параметров.

  • Пиковая память обучения снижается с 81.4 ГБ до 31.3 ГБ, что укладывается в бюджет акселератора в 40 ГБ.
  • SkewAdam достигает валидационной перплексии 108.4 на 82M токенов, превосходя AdamW (126.8), Muon (120.2) и Lion (393.7).
  • Экономия памяти достигается без потери точности; прирост производительности обусловлен сохранением импульса, а не самим многоуровневым распределением.

Результаты показывают, что то, где распределяется состояние оптимизатора, так же важно, как и его объем, что позволяет эффективно обучаться на оборудовании с ограниченной памятью.