Addestramento Asincrono Local-SGD per il Modellamento del Linguaggio

Asynchronous Local-SGD Training for Language Modeling

January 17, 2024

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

cs.AI

Abstract

La discesa del gradiente stocastica locale (Local-SGD), anche nota come media federata, è un approccio all'ottimizzazione distribuita in cui ogni dispositivo esegue più di un aggiornamento SGD per comunicazione. Questo lavoro presenta uno studio empirico della Local-SGD {\it asincrona} per l'addestramento di modelli linguistici; ovvero, ogni worker aggiorna i parametri globali non appena ha completato i suoi passi SGD. Condurremo un'indagine approfondita esaminando come l'eterogeneità hardware dei worker, la dimensione del modello, il numero di worker e l'ottimizzatore possano influenzare le prestazioni di apprendimento. Scopriamo che, con implementazioni naive, la Local-SGD asincrona richiede più iterazioni per convergere rispetto alla sua controparte sincrona, nonostante aggiorni i parametri del modello (globali) più frequentemente. Identifichiamo l'accelerazione del momento sui parametri globali quando i gradienti dei worker sono obsoleti come una sfida chiave. Proponiamo un metodo innovativo che utilizza un aggiornamento del momento di Nesterov ritardato e regola i passi di addestramento locali dei worker in base alla loro velocità di calcolo. Questo approccio, valutato con modelli fino a 150M di parametri sul dataset C4, eguaglia le prestazioni della Local-SGD sincrona in termini di perplessità per passo di aggiornamento e la supera significativamente in termini di tempo di clock.

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.