DiLoCo: Gedistribueerde training van taalmodelen met weinig communicatie

DiLoCo: Distributed Low-Communication Training of Language Models

November 14, 2023

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

cs.AI

Samenvatting

Grote taalmodelen (LLM) zijn een cruciaal onderdeel geworden in veel toepassingen van machine learning. Standaardbenaderingen voor het trainen van LLM vereisen echter een groot aantal nauw met elkaar verbonden accelerators, waarbij apparaten bij elke optimalisatiestap gradiënten en andere tussenliggende staten uitwisselen. Hoewel het moeilijk is om een enkel rekencluster te bouwen en te onderhouden dat veel accelerators host, kan het gemakkelijker zijn om meerdere rekenclusters te vinden die elk een kleiner aantal apparaten hosten. In dit werk stellen we een gedistribueerd optimalisatie-algoritme voor, Distributed Low-Communication (DiLoCo), dat het mogelijk maakt om taalmodelen te trainen op eilanden van slecht verbonden apparaten. De aanpak is een variant van federated averaging, waarbij het aantal binnenstappen groot is, de binnenste optimizer AdamW is, en de buitenste optimizer Nesterov momentum is. Op de veelgebruikte C4-dataset laten we zien dat DiLoCo op 8 workers even goed presteert als volledig synchrone optimalisatie, terwijl het 500 keer minder communiceert. DiLoCo toont een grote robuustheid ten opzichte van de gegevensverdeling van elke worker. Het is ook robuust tegen het onbeschikbaar worden van bronnen over tijd, en omgekeerd kan het naadloos gebruikmaken van bronnen die tijdens het trainen beschikbaar worden.

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.