Des poids structurés pour réduire le coût de Muon à grande échelle
En termes simples
Le préentraînement d’un modèle de langue demande beaucoup de calcul. Cet article présente Monarch–Muon, un optimiseur qui stocke certains poids du modèle en morceaux plus petits et met à jour ces morceaux à l’aide d’étapes adaptées aux matrices.
Lors d’essais menés en un seul entraînement sur deux types de modèles, la perte finale sur les données de validation était plus faible qu’avec AdamW lorsque les deux méthodes utilisaient la même structure en blocs. Des mesures séparées ont aussi relevé un temps d’entraînement et une consommation de mémoire plus faibles à la plus grande taille testée, sur un seul H100.
La recherche
Le préentraînement des modèles de langue peut-il bénéficier de mises à jour adaptées aux matrices, qui peuvent améliorer l’optimisation, sans payer tout le coût de stockage et de calcul des poids denses ?
L’article teste si une représentation structurée de certains poids peut rendre ces mises à jour moins coûteuses tout en conservant leur avantage d’optimisation.
Monarch–Muon applique des mises à jour adaptées aux matrices à des facteurs de poids structurés, pour conserver les avantages d’optimisation tout en réduisant le stockage et le coût de l’orthogonalisation.
Monarch–Muon met à jour des facteurs de poids structurés avec des étapes adaptées aux matrices. Les chercheurs l’ont comparé à AdamW, à structure de facteurs identique, lors de préentraînements denses et mixture-of-experts sur FineWeb-Edu.
La perte finale sur l’ensemble de validation était plus faible avec Monarch–Muon dans les quatre comparaisons menées en un seul entraînement. Ces résultats suggèrent que les mises à jour adaptées aux matrices peuvent être utiles avec les structures testées.
Les comparaisons de qualité reposent sur un seul entraînement pour chaque taille de modèle. Les mesures système, menées séparément, ne permettent pas d’établir le compromis entre qualité et matériel lors d’un entraînement distribué plus long.
La preuve de convergence suppose des mises à jour polaires exactes, alors que l’optimiseur pratique les approxime au moyen d’un nombre fini d’itérations Newton–Schulz.
Analyse
Détails
- Auteurs de l’article original
- Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov, Dmitrii Vlashchuk, Vladimir Aletov
- Soumis à arXiv
- Licence
- CC BY 4.0