Asynchronous Local-SGD Training for Language Modeling
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.
65a8a864fee7ba6063c7e431
Bo Liu
65a8a864fee7ba6063c7e432
Rachita Chhaparia
65a8a864fee7ba6063c7e433
Arthur Douillard
admin_assigned
2024-01-18T10:06:56.510Z
622792366303bf1dc304f49f
/avatars/975c1cc3eb2f97cf8e848162056d5bea.svg
Arthur Douillard
user
ArthurDouillard
65a8a864fee7ba6063c7e434
Satyen Kale
65a8a864fee7ba6063c7e435
Andrei A. Rusu
65a8a864fee7ba6063c7e436
Jiajun Shen
65a8a864fee7ba6063c7e437
Arthur Szlam
65a8a864fee7ba6063c7e438
Marc'Aurelio Ranzato
Asynchronous Local-SGD Training for Language Modeling
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.
Asynchronous Local-SGD with delayed Nesterov momentum and adaptive local training steps improves performance and clock time for large language models training.
Local-SGD federated averaging asynchronous worker hardware heterogeneity model size number of workers optimizer momentum acceleration Nesterov momentum C4 dataset perplexity