Asynchronous Local-SGD Training for Language Modeling

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Liu, Bo, Chhaparia, Rachita, Douillard, Arthur, Kale, Satyen, Rusu, Andrei A., Shen, Jiajun, Szlam, Arthur, Ranzato, Marc'Aurelio
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866912040600207360
author Liu, Bo
Chhaparia, Rachita
Douillard, Arthur
Kale, Satyen
Rusu, Andrei A.
Shen, Jiajun
Szlam, Arthur
Ranzato, Marc'Aurelio
author_facet Liu, Bo
Chhaparia, Rachita
Douillard, Arthur
Kale, Satyen
Rusu, Andrei A.
Shen, Jiajun
Szlam, Arthur
Ranzato, Marc'Aurelio
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2401_09135
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Asynchronous Local-SGD Training for Language Modeling
Liu, Bo
Chhaparia, Rachita
Douillard, Arthur
Kale, Satyen
Rusu, Andrei A.
Shen, Jiajun
Szlam, Arthur
Ranzato, Marc'Aurelio
Machine Learning
Computation and Language
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.
title Asynchronous Local-SGD Training for Language Modeling
topic Machine Learning
Computation and Language
url https://arxiv.org/abs/2401.09135