DiLoCo: 언어 모델의 분산형 저통신 훈련

초록

대규모 언어 모델(LLM)은 머신러닝의 많은 응용 분야에서 핵심 구성 요소로 자리 잡았습니다. 그러나 LLM을 훈련시키기 위한 표준 접근 방식은 많은 수의 긴밀하게 연결된 가속기를 필요로 하며, 각 최적화 단계에서 기기들이 그래디언트와 기타 중간 상태를 교환해야 합니다. 많은 수의 가속기를 호스팅하는 단일 컴퓨팅 클러스터를 구축하고 유지하는 것은 어려운 일이지만, 더 적은 수의 기기를 호스팅하는 여러 컴퓨팅 클러스터를 찾는 것은 더 쉬울 수 있습니다. 본 연구에서는 연결 상태가 좋지 않은 기기 그룹(섬)에서도 언어 모델을 훈련시킬 수 있는 분산 최적화 알고리즘인 DiLoCo(Distributed Low-Communication)를 제안합니다. 이 접근 방식은 페더레이티드 평균화(federated averaging)의 변형으로, 내부 단계 수가 많고 내부 최적화기는 AdamW, 외부 최적화기는 네스테로프 모멘텀(Nesterov momentum)을 사용합니다. 널리 사용되는 C4 데이터셋에서, 8개의 워커를 사용한 DiLoCo는 완전 동기화 최적화와 동등한 성능을 보이면서도 통신량을 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.

DiLoCo: 언어 모델의 분산형 저통신 훈련

DiLoCo: Distributed Low-Communication Training of Language Models

초록

Support