[논문 리뷰] DiLoCo: Distributed Low-Communication Training of Language Models

Arthur Douillard, Qixuan Feng외 7명|arXiv (Cornell University)|2023. 11. 14.

Topic Modeling인용 수 7

한 줄 요약

DiLoCo는 대규모 언어 모델을 위한 분산 학습 프레임워크를 제안하며, AdamW를 사용해 큰 내부 최적화 단계(500단계)를 수행하고 외부 Nesterov 모멘텀 업데이트를 적용함으로써 통신을 극도로 줄입니다. 이로써 C4 데이터셋에서 완전히 동기화된 학습과 비교해도 성능가능성을 유지하면서 통신량을 500배 줄였고, 데이터 분포, 워커 가용성 변화, 이질적인 장치에 대한 강건성도 입증했습니다.

ABSTRACT

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.

연구 동기 및 목표

• 가속기들이 공유된 위치에 있지 않거나 통신 대역폭이 제한된 환경에서 대규모 언어 모델을 학습하는 데 도전하는 것.
• 엄격하게 동기화되고 고대역폭이 필요한 클러스터를 필요로 하는 완전히 동기화된 분산 학습의 확장성 및 신뢰성 문제를 해결하는 것.
• 서로 다른 네트워크 환경에 분산된 클러스터(‘아일랜드’) 간의 통신이 열악한 상황에서도 학습을 가능하게 하는 것.
• 데이터 분포 변화, 워커 장애, 컴퓨팅 자원의 동적 가용성 변화에 대한 강건성을 향상시키는 것.
• 통신 오버헤드를 최소화하면서도 완전히 동기화된 학습과 동등한 모델 성능과 효율성을 유지하는 것.

제안 방법

• 통신 빈도를 줄이기 위해 큰 내부 단계(H=500)를 사용하는 연합 평균 기반 프레임워크를 채택.
• 각 워커의 로컬 모델 업데이트에 표준 LLM 학습 관행에 맞게 AdamW를 내부 최적화기로 사용.
• 모든 워커로부터의 기울기를 집계한 후, 전역 모델 파라미터를 업데이트하기 위해 Nesterov 모멘텀을 외부 최적화기로 적용.
• 모델 업데이트를 로컬 학습 단계와 드물게 발생하는 전역 동기화 단계로 분리하여 노드 간 통신 빈도를 500단계에 한 번으로 줄임.
• 다양한 하드웨어 유형과 위치에 있는 워커들이 작동할 수 있도록 하되, 주기적인 파라미터 동기화만 수행.
• 각 외부 최적화 단계 후에 업데이트된 전역 모델을 모든 워커에게 재배포하여 반복적인 정밀 조정을 가능하게 함.

실험 결과

연구 질문

• RQ1 드물게 통신하는 분산 학습 방법이 대규모 언어 모델 학습에서 완전히 동기화된 학습과 비교해 성능이 유사한가?
• RQ2 다양한 워커 샤드 간 데이터 분포 변화에 대해 이 방법은 얼마나 강건한가?
• RQ3 학습 중 워커 가용성의 동적 변화(예: 장애, 새로운 자원 추가)를 어떻게 처리할 수 있는가?
• RQ4 8명 이상의 워커로 확장할 경우 이 방법은 효과적으로 작동하는가? 통신 감소의 한계는 무엇인가?
• RQ5 클러스터 내에서 비균일하거나 이질적인 장치를 사용할 경우에도 이 방법이 학습 안정성과 수렴성을 유지할 수 있는가?

주요 결과

• DiLoCo는 C4 데이터셋에서 완전히 동기화된 학습과 비교해 퍼플렉서티 측면에서 성능이 유사하거나 이를 초월함에도 불구하고 통신 데이터량을 500배 줄였습니다.
• 워커 샤드 간 다양한 데이터 분포 변화에 대해 안정적인 학습과 수렴을 유지하는 데 성공했습니다.
• DiLoCo는 워커 가용성의 동적 변화를 유연하게 처리하며, 새로운 워커를 원활하게 통합하고 장애 복구를 수행할 수 있었습니다.
• 저대역폭 연결을 가진 지리적으로 분산된 클러스터 간 학습이 가능해졌으며, 통신 주기를 몇 분에 한 번으로 줄여 기존의 수백 밀리초 단위의 주기에서 벗어났습니다.
• DiLoCo는 완전히 동기화된 모델과 동일한 모델 효율성과 추론 속도를 유지하며, 추론 시 성능 저하가 없었습니다.
• DiLoCo는 8명의 워커를 초과할 경우 수익 감소 현상이 나타나지만, 이질적이고 느슨하게 연결된 환경에서 강력한 벽시계 시간 효율성과 강건성을 입증했습니다.