Entraînement Asynchrone Local-SGD pour la Modélisation du Langage

Asynchronous Local-SGD Training for Language Modeling

January 17, 2024

Auteurs: Bo Liu, Rachita Chhaparia, Arthur Douillard, Satyen Kale, Andrei A. Rusu, Jiajun Shen, Arthur Szlam, Marc'Aurelio Ranzato

cs.AI

Résumé

La descente de gradient stochastique locale (Local-SGD), également appelée agrégation fédérée, est une approche d'optimisation distribuée où chaque appareil effectue plus d'une mise à jour SGD par communication. Ce travail présente une étude empirique de la Local-SGD {\it asynchrone} pour l'entraînement de modèles de langage ; c'est-à-dire que chaque travailleur met à jour les paramètres globaux dès qu'il a terminé ses étapes SGD. Nous menons une investigation approfondie en examinant comment l'hétérogénéité matérielle des travailleurs, la taille du modèle, le nombre de travailleurs et l'optimiseur pourraient influencer les performances d'apprentissage. Nous constatons qu'avec des implémentations naïves, la Local-SGD asynchrone nécessite plus d'itérations pour converger que sa contrepartie synchrone, malgré une mise à jour plus fréquente des paramètres (globaux) du modèle. Nous identifions l'accélération par momentum sur les paramètres globaux lorsque les gradients des travailleurs sont obsolètes comme un défi majeur. Nous proposons une nouvelle méthode qui utilise une mise à jour de momentum de Nesterov retardée et ajuste les étapes d'entraînement locales des travailleurs en fonction de leur vitesse de calcul. Cette approche, évaluée avec des modèles allant jusqu'à 150 millions de paramètres sur le jeu de données C4, correspond aux performances de la Local-SGD synchrone en termes de perplexité par étape de mise à jour, et la dépasse significativement en termes de temps écoulé.

English

Local stochastic gradient descent (Local-SGD), also referred to as federated averaging, is an approach to distributed optimization where each device performs more than one SGD update per communication. This work presents an empirical study of {\it asynchronous} Local-SGD for training language models; that is, each worker updates the global parameters as soon as it has finished its SGD steps. We conduct a comprehensive investigation by examining how worker hardware heterogeneity, model size, number of workers, and optimizer could impact the learning performance. We find that with naive implementations, asynchronous Local-SGD takes more iterations to converge than its synchronous counterpart despite updating the (global) model parameters more frequently. We identify momentum acceleration on the global parameters when worker gradients are stale as a key challenge. We propose a novel method that utilizes a delayed Nesterov momentum update and adjusts the workers' local training steps based on their computation speed. This approach, evaluated with models up to 150M parameters on the C4 dataset, matches the performance of synchronous Local-SGD in terms of perplexity per update step, and significantly surpasses it in terms of wall clock time.