Les chercheurs présentent SkewAdam, un optimiseur conçu pour les modèles à mélange d'experts (MoE) qui alloue différents types d'état à des populations de paramètres distinctes afin de réduire drastiquement l'utilisation de la mémoire. En attribuant le momentum float32 et les moments seconds factorisés au noyau, les moments seconds factorisés aux experts, et les moments seconds exacts au routeur, SkewAdam réduit l'état de l'optimiseur de 50,6 Go à 1,29 Go pour un modèle de 6,78 milliards de paramètres.

  • La mémoire maximale de formation passe de 81,4 Go à 31,3 Go, s'intégrant dans un budget d'accélérateur de 40 Go.
  • SkewAdam atteint une perplexité de validation de 108,4 sur 82 millions de tokens, surpassant AdamW (126,8), Muon (120,2) et Lion (393,7).
  • Les économies de mémoire sont réalisées sans perte de précision ; le gain de performance provient de la conservation du momentum plutôt que de l'allocation en couches elle-même.

Les résultats suggèrent que l'emplacement où l'état de l'optimiseur est alloué est aussi important que son volume, permettant une formation efficace sur du matériel à mémoire limitée.