Asynchrone Local-SGD Training voor Taalmodellering

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

Samenvatting

Lokale stochastische gradiëntdaling (Local-SGD), ook wel federated averaging genoemd, is een benadering voor gedistribueerde optimalisatie waarbij elk apparaat meer dan één SGD-update uitvoert per communicatie. Dit werk presenteert een empirische studie van {\it asynchrone} Local-SGD voor het trainen van taalmodelen; dat wil zeggen, elke worker werkt de globale parameters bij zodra deze zijn SGD-stappen heeft voltooid. We voeren een uitgebreid onderzoek uit door te analyseren hoe hardware-heterogeniteit van workers, modelgrootte, aantal workers en de optimizer de leerprestaties kunnen beïnvloeden. We constateren dat bij naïeve implementaties asynchrone Local-SGD meer iteraties nodig heeft om te convergeren dan zijn synchrone tegenhanger, ondanks het vaker bijwerken van de (globale) modelparameters. We identificeren momentumversnelling op de globale parameters wanneer worker-gradiënten verouderd zijn als een belangrijke uitdaging. We stellen een nieuwe methode voor die gebruikmaakt van een vertraagde Nesterov-momentumupdate en de lokale trainingsstappen van de workers aanpast op basis van hun rekensnelheid. Deze aanpak, geëvalueerd met modellen tot 150M parameters op de C4-dataset, evenaart de prestaties van synchrone Local-SGD in termen van perplexiteit per update-stap en overtreft deze aanzienlijk in termen van werkelijke rekentijd.

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.