Entrenamiento Asíncrono Local-SGD para Modelado de Lenguaje

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

Resumen

El descenso de gradiente estocástico local (Local-SGD), también conocido como promediado federado, es un enfoque de optimización distribuida en el que cada dispositivo realiza más de una actualización de SGD por comunicación. Este trabajo presenta un estudio empírico del Local-SGD {\it asíncrono} para entrenar modelos de lenguaje; es decir, cada trabajador actualiza los parámetros globales tan pronto como termina sus pasos de SGD. Realizamos una investigación exhaustiva examinando cómo la heterogeneidad del hardware de los trabajadores, el tamaño del modelo, el número de trabajadores y el optimizador podrían afectar el rendimiento del aprendizaje. Descubrimos que, con implementaciones ingenuas, el Local-SGD asíncrono requiere más iteraciones para converger que su contraparte síncrona, a pesar de actualizar los parámetros del modelo (global) con mayor frecuencia. Identificamos la aceleración por momento en los parámetros globales cuando los gradientes de los trabajadores están desactualizados como un desafío clave. Proponemos un método novedoso que utiliza una actualización de momento de Nesterov retrasada y ajusta los pasos de entrenamiento local de los trabajadores según su velocidad de cálculo. Este enfoque, evaluado con modelos de hasta 150 millones de parámetros en el conjunto de datos C4, iguala el rendimiento del Local-SGD síncrono en términos de perplejidad por paso de actualización y lo supera significativamente en términos de tiempo de reloj.

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.