DeMo : Optimisation du Momentum Découplé

DeMo: Decoupled Momentum Optimization

November 29, 2024

Auteurs: Bowen Peng, Jeffrey Quesnelle, Diederik P. Kingma

cs.AI

Résumé

L'entraînement de grands réseaux neuronaux nécessite généralement de partager les gradients entre les accélérateurs via des interconnexions spécialisées à haute vitesse. En nous inspirant des principes de traitement du signal de décomposition en fréquence et de compaction d'énergie, nous démontrons que la synchronisation des états complets de l'optimiseur et des paramètres du modèle pendant l'entraînement est inutile. En dissociant les mises à jour de momentum et en permettant une divergence contrôlée des états de l'optimiseur entre les accélérateurs, nous obtenons une convergence améliorée par rapport aux optimiseurs de pointe. Nous présentons Découplage de Momentum (DeMo), un optimiseur fusionné et un algorithme parallèle de données qui réduit les exigences de communication entre les accélérateurs de plusieurs ordres de grandeur. Cela permet l'entraînement de grands réseaux neuronaux même avec une bande passante réseau limitée et un matériel hétérogène. Notre méthode est agnostique de la topologie et indépendante de l'architecture, et prend en charge un entraînement distribué synchrone par horloge scalable avec une surcharge de calcul et de mémoire négligeable. Les résultats empiriques montrent que les modèles entraînés avec DeMo correspondent ou dépassent les performances des modèles équivalents entraînés avec AdamW, tout en éliminant le besoin d'interconnexions à haute vitesse lors de la pré-formation de modèles de base à grande échelle. Une implémentation PyTorch de référence en open source est publiée sur GitHub à l'adresse https://github.com/bloc97/DeMo

English

Training large neural networks typically requires sharing gradients between accelerators through specialized high-speed interconnects. Drawing from the signal processing principles of frequency decomposition and energy compaction, we demonstrate that synchronizing full optimizer states and model parameters during training is unnecessary. By decoupling momentum updates and allowing controlled divergence in optimizer states across accelerators, we achieve improved convergence compared to state-of-the-art optimizers. We introduce {De}coupled {Mo}mentum (DeMo), a fused optimizer and data parallel algorithm that reduces inter-accelerator communication requirements by several orders of magnitude. This enables training of large neural networks even with limited network bandwidth and heterogeneous hardware. Our method is topology-agnostic and architecture-independent and supports scalable clock-synchronous distributed training with negligible compute and memory overhead. Empirical results show that models trained with DeMo match or exceed the performance of equivalent models trained with AdamW, while eliminating the need for high-speed interconnects when pre-training large scale foundation models. An open source reference PyTorch implementation is published on GitHub at https://github.com/bloc97/DeMo