DiLoCo: Распределённое обучение языковых моделей с низким уровнем коммуникации

DiLoCo: Distributed Low-Communication Training of Language Models

November 14, 2023

Авторы: Arthur Douillard, Qixuan Feng, Andrei A. Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc'Aurelio Ranzato, Arthur Szlam, Jiajun Shen

cs.AI

Аннотация

Крупные языковые модели (LLM) стали важным компонентом во многих приложениях машинного обучения. Однако стандартные подходы к обучению LLM требуют использования большого числа тесно связанных акселераторов, где устройства обмениваются градиентами и другими промежуточными состояниями на каждом шаге оптимизации. Хотя создание и поддержка единого вычислительного кластера, содержащего множество акселераторов, является сложной задачей, может быть проще найти несколько вычислительных кластеров, каждый из которых содержит меньшее количество устройств. В данной работе мы предлагаем распределенный алгоритм оптимизации, Distributed Low-Communication (DiLoCo), который позволяет обучать языковые модели на изолированных группах устройств с плохой связью. Этот подход представляет собой вариант федеративного усреднения, где количество внутренних шагов велико, внутренний оптимизатор — AdamW, а внешний оптимизатор — метод Нестерова. На широко используемом наборе данных C4 мы показываем, что DiLoCo на 8 рабочих узлах работает так же эффективно, как полностью синхронная оптимизация, при этом обмениваясь данными в 500 раз меньше. DiLoCo демонстрирует высокую устойчивость к распределению данных на каждом рабочем узле. Он также устойчив к недоступности ресурсов с течением времени и, наоборот, может плавно использовать ресурсы, которые становятся доступными в процессе обучения.

English

Large language models (LLM) have become a critical component in many applications of machine learning. However, standard approaches to training LLM require a large number of tightly interconnected accelerators, with devices exchanging gradients and other intermediate states at each optimization step. While it is difficult to build and maintain a single computing cluster hosting many accelerators, it might be easier to find several computing clusters each hosting a smaller number of devices. In this work, we propose a distributed optimization algorithm, Distributed Low-Communication (DiLoCo), that enables training of language models on islands of devices that are poorly connected. The approach is a variant of federated averaging, where the number of inner steps is large, the inner optimizer is AdamW, and the outer optimizer is Nesterov momentum. On the widely used C4 dataset, we show that DiLoCo on 8 workers performs as well as fully synchronous optimization while communicating 500 times less. DiLoCo exhibits great robustness to the data distribution of each worker. It is also robust to resources becoming unavailable over time, and vice versa, it can seamlessly leverage resources that become available during training.