DiLoCo: Treinamento Distribuído de Modelos de Linguagem com Baixa Comunicação

DiLoCo: Distributed Low-Communication Training of Language Models

November 14, 2023

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

cs.AI

Resumo

Modelos de linguagem de grande escala (LLM) tornaram-se um componente crítico em muitas aplicações de aprendizado de máquina. No entanto, as abordagens padrão para treinar LLMs exigem um grande número de aceleradores fortemente interconectados, com dispositivos trocando gradientes e outros estados intermediários a cada etapa de otimização. Embora seja difícil construir e manter um único cluster de computação que hospede muitos aceleradores, pode ser mais fácil encontrar vários clusters de computação, cada um hospedando um número menor de dispositivos. Neste trabalho, propomos um algoritmo de otimização distribuído, Distributed Low-Communication (DiLoCo), que permite o treinamento de modelos de linguagem em ilhas de dispositivos mal conectados. A abordagem é uma variante da média federada, onde o número de passos internos é grande, o otimizador interno é o AdamW e o otimizador externo é o momento de Nesterov. No amplamente utilizado conjunto de dados C4, mostramos que o DiLoCo em 8 trabalhadores tem um desempenho tão bom quanto a otimização totalmente síncrona, enquanto se comunica 500 vezes menos. O DiLoCo exibe grande robustez à distribuição de dados de cada trabalhador. Ele também é robusto a recursos que se tornam indisponíveis ao longo do tempo e, inversamente, pode aproveitar de forma contínua recursos que se tornam disponíveis durante o treinamento.

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.