Асинхронное обучение Local-SGD для языкового моделирования

Asynchronous Local-SGD Training for Language Modeling

January 17, 2024

Авторы: Bo Liu, Rachita Chhaparia, Arthur Douillard, Satyen Kale, Andrei A. Rusu, Jiajun Shen, Arthur Szlam, Marc'Aurelio Ranzato

cs.AI

Аннотация

Локальный стохастический градиентный спуск (Local-SGD), также известный как федеративное усреднение, представляет собой подход к распределенной оптимизации, при котором каждое устройство выполняет более одного шага SGD перед обменом данными. В данной работе представлено эмпирическое исследование асинхронного Local-SGD для обучения языковых моделей, где каждый рабочий процесс обновляет глобальные параметры сразу после завершения своих шагов SGD. Мы проводим всестороннее исследование, изучая, как неоднородность оборудования рабочих устройств, размер модели, количество рабочих процессов и оптимизатор могут повлиять на производительность обучения. Мы обнаруживаем, что при наивной реализации асинхронный Local-SGD требует больше итераций для сходимости по сравнению с синхронным подходом, несмотря на более частые обновления глобальных параметров модели. Мы выявляем ключевую проблему, связанную с ускорением импульса на глобальных параметрах, когда градиенты рабочих процессов устаревают. Мы предлагаем новый метод, который использует отложенное обновление импульса Нестерова и корректирует количество локальных шагов обучения рабочих процессов в зависимости от их скорости вычислений. Этот подход, протестированный на моделях с количеством параметров до 150 миллионов на наборе данных C4, демонстрирует сопоставимую с синхронным Local-SGD производительность в терминах перплексии на шаг обновления и значительно превосходит его по времени выполнения.

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.