Treinamento Assíncrono de SGD Local para Modelagem de Linguagem

Asynchronous Local-SGD Training for Language Modeling

January 17, 2024

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

cs.AI

Resumo

O gradiente descendente estocástico local (Local-SGD), também conhecido como média federada, é uma abordagem de otimização distribuída na qual cada dispositivo realiza mais de uma atualização de SGD por comunicação. Este trabalho apresenta um estudo empírico do Local-SGD {\it assíncrono} para o treinamento de modelos de linguagem; ou seja, cada trabalhador atualiza os parâmetros globais assim que conclui suas etapas de SGD. Realizamos uma investigação abrangente examinando como a heterogeneidade de hardware dos trabalhadores, o tamanho do modelo, o número de trabalhadores e o otimizador podem impactar o desempenho de aprendizado. Descobrimos que, com implementações ingênuas, o Local-SGD assíncrono leva mais iterações para convergir do que sua contraparte síncrona, apesar de atualizar os parâmetros (globais) do modelo com mais frequência. Identificamos a aceleração de momento nos parâmetros globais quando os gradientes dos trabalhadores estão desatualizados como um desafio fundamental. Propomos um método inovador que utiliza uma atualização de momento de Nesterov atrasada e ajusta as etapas de treinamento local dos trabalhadores com base em sua velocidade de computação. Essa abordagem, avaliada com modelos de até 150 milhões de parâmetros no conjunto de dados C4, iguala o desempenho do Local-SGD síncrono em termos de perplexidade por etapa de atualização e o supera significativamente em termos de tempo de execução real.

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.