[논문 리뷰] Decentralized Training of Foundation Models in Heterogeneous Environments

Binhang Yuan, Yongjun He외 7명|arXiv (Cornell University)|2022. 06. 02.

Stochastic Gradient Optimization Techniques인용 수 23

한 줄 요약

이 논문은 지리적으로 분산되고 이질적인 GPU에서 GPT-스타일 기초 모델을 학습하기 위한 스케줄러와 시스템 최적화를 제시하여 중앙 집중식 기준선 대비 상당한 속도 향상을 달성한다.

ABSTRACT

Training foundation models, such as GPT-3 and PaLM, can be extremely expensive, often involving tens of thousands of GPUs running continuously for months. These models are typically trained in specialized clusters featuring fast, homogeneous interconnects and using carefully designed software systems that support both data parallelism and model/pipeline parallelism. Such dedicated clusters can be costly and difficult to obtain. Can we instead leverage the much greater amount of decentralized, heterogeneous, and lower-bandwidth interconnected compute? Previous works examining the heterogeneous, decentralized setting focus on relatively small models that can be trained in a purely data parallel manner. State-of-the-art schemes for model parallel foundation model training, such as Megatron, only consider the homogeneous data center setting. In this paper, we present the first study of training large foundation models with model parallelism in a decentralized regime over a heterogeneous network. Our key technical contribution is a scheduling algorithm that allocates different computational "tasklets" in the training of foundation models to a group of decentralized GPU devices connected by a slow heterogeneous network. We provide a formal cost model and further propose an efficient evolutionary algorithm to find the optimal allocation strategy. We conduct extensive experiments that represent different scenarios for learning over geo-distributed devices simulated using real-world network measurements. In the most extreme case, across 8 different cities spanning 3 continents, our approach is 4.8X faster than prior state-of-the-art training systems (Megatron).

연구 동기 및 목표

• 비싼 데이터센터 클러스터에 대한 의존도를 줄이기 위해 대형 기초 모델의 비중앙 집중형 이질적 네트워크 학습을 촉진한다.
• 이질적 환경에서 파이프라인 및 데이터 병렬성을 위한 태스크렛-대-디바이스 할당 문제를 수식화한다.
• 분산 학습에서 통신 비용을 최소화하기 위한 비용 모델과 최적화 알고리즘을 개발한다.
• 처리량 향상을 위해 통신을 계산과 겹치게 하는 시스템 수준의 최적화를 시연한다.
• 실제 네트워크 측정을 시뮬레이션하는 지리적으로 분산된 GPU 클러스터에서 접근법을 평가한다.

제안 방법

• 학습을 디바이스에 할당된 태스크렛(마이크로배치, 계층 단계)의 집합으로 모델링한다.
• 이중 수준의 비용 모델을 사용하여 통신 비용을 데이터 병렬 및 파이프라인 병렬 구성요소로 분해한다.
• 데이터P-비용(DataP-Cost)과 파이프라인P-비용(PipelineP-Cost)을 축소된 그래프와 개방 루프 순회하는 판매원 문제로 형식화한다.
• 특화된 로컬 서치를 갖춘 하이브리드 유전 알고리즘을 사용하여 거의 최적에 가까운 할당을 찾는다.
• CUDA 스트림 간에 계산과 통신을 파이프라인화하는 시스템 수준의 최적화를 도입한다.

실험 결과

연구 질문

• RQ1 이질적 네트워크에서 통신 비용을 최소화하기 위해 분산된 GPU 태스크렛을 어떻게 할당하나요?
• RQ2 이질성 하에서 데이터 병렬 및 파이프라인 병렬 통신 비용을 어떻게 모델링하나요?
• RQ3 진화적 스케줄링 접근법이 무작위나 Kernighan-Lin 전략을 능가하는 할당을 찾을 수 있나요?
• RQ4 지리적으로 분산된 환경에서 스케줄러 및 시스템 최적화가 엔드투엔드 학습 처리량에 미치는 영향은 무엇인가요?
• RQ5 분산 학습이 데이터센터 기준선에 얼마나 근접할 수 있나요?

주요 결과

• 제안된 스케줄러 및 시스템 최적화는 전 세계에 걸친 지오분산 환경에서 최첨단 중앙집중식 시스템보다 엔드투엔드 학습을 3.8–4.8배 빠르게 수행한다.
• 전 세계적인 지오분산 시나리오에서 본 방법은 GPT3-1.3B 학습에 대해 Megatron-1.3B보다 4.8배, Deepspeed보다 3.6배 빠르다.
• 스케줄러가 없으면 성능 향상이 감소하며, 제거 실험에서 최대 2.7배 느려진 결과가 나타난다.
• 데이터 센터 조건에서 비분산 접근은 Megatron/Deepspeed에 비해 1.7–3.5배 느리나, 네트워크가 최대 100배 느려도 여전히 비교적 준수하다.
• 통신과 계산의 중첩을 가능하게 하는 시스템 최적화가 처리량을 더 향상시키고, 스케줄러가 무작위 할당 대비 속도 향상에 크게 기여한다.
• 특화된 로컬 탐색이 비용 모델에 맞춰진 제거 실험은 표준 Kernighan-Lin 기반 탐색보다 성능이 우수하다는 것을 보여준다.