DiLoCo: Addestramento Distribuito di Modelli Linguistici a Bassa Comunicazione

DiLoCo: Distributed Low-Communication Training of Language Models

November 14, 2023

Autori: Arthur Douillard, Qixuan Feng, Andrei A. Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc'Aurelio Ranzato, Arthur Szlam, Jiajun Shen

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) sono diventati un componente fondamentale in molte applicazioni di apprendimento automatico. Tuttavia, gli approcci standard per l'addestramento degli LLM richiedono un numero elevato di acceleratori strettamente interconnessi, con dispositivi che si scambiano gradienti e altri stati intermedi a ogni passo di ottimizzazione. Sebbene sia difficile costruire e mantenere un singolo cluster di calcolo che ospiti molti acceleratori, potrebbe essere più semplice trovare diversi cluster di calcolo, ciascuno con un numero inferiore di dispositivi. In questo lavoro, proponiamo un algoritmo di ottimizzazione distribuito, Distributed Low-Communication (DiLoCo), che consente l'addestramento di modelli linguistici su isole di dispositivi scarsamente connesse. L'approccio è una variante della media federata, in cui il numero di passi interni è elevato, l'ottimizzatore interno è AdamW e l'ottimizzatore esterno è il momento di Nesterov. Sul dataset ampiamente utilizzato C4, dimostriamo che DiLoCo su 8 worker performa altrettanto bene quanto l'ottimizzazione completamente sincrona, comunicando 500 volte in meno. DiLoCo mostra una grande robustezza rispetto alla distribuzione dei dati di ciascun worker. È inoltre robusto alla disponibilità di risorse che diminuisce nel tempo e, viceversa, può sfruttare senza interruzioni le risorse che diventano disponibili durante l'addestramento.

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.