[논문 리뷰] BAGUA: Scaling up Distributed Learning with System Relaxations

Shaoduo Gan, Jiawei Jiang외 12명|arXiv (Cornell University)|2021. 01. 01.

Stochastic Gradient Optimization Techniques참고 문헌 79인용 수 3

한 줄 요약

BAGUA는 모듈형이며 MPI 스타일의 통신 라이브러리로, 분산 딥 러닝에서 고도화된 시스템 이완 기법—예를 들어 양자화, 탈중앙화, 비동기 학습—을 효율적으로 구현할 수 있도록 한다. BAGUA는 유연한 최적화 프레임워크를 통해 과거 최고 수준의 알고리즘을 네이티브로 지원함으로써, 겹침, 융합, 계층적 통신 기능을 통해 다양한 워크로드에서 끝에서 끝까지의 학습 시간 동안 PyTorch-DDP, Horovod, BytePS를 최대 2배 빠르게 성능을 냈다.

ABSTRACT

Recent years have witnessed a growing list of systems for distributed data-parallel training. Existing systems largely fit into two paradigms, i.e., parameter server and MPI-style collective operations. On the algorithmic side, researchers have proposed a wide range of techniques to lower the communication via system relaxations: quantization, decentralization, and communication delay. However, most, if not all, existing systems only rely on standard synchronous and asynchronous stochastic gradient (SG) based optimization, therefore, cannot take advantage of all possible optimizations that the machine learning community has been developing recently. Given this emerging gap between the current landscapes of systems and theory, we build BAGUA, a MPI-style communication library, providing a collection of primitives, that is both flexible and modular to support state-of-the-art system relaxation techniques of distributed training. Powered by this design, BAGUA has a great ability to implement and extend various state-of-the-art distributed learning algorithms. In a production cluster with up to 16 machines (128 GPUs), BAGUA can outperform PyTorch-DDP, Horovod and BytePS in the end-to-end training time by a significant margin (up to 2 times) across a diverse range of tasks. Moreover, we conduct a rigorous tradeoff exploration showing that different algorithms and system relaxations achieve the best performance over different network conditions.

연구 동기 및 목표

• 분산 학습 알고리즘의 이론적 진전과 여전히 표준 동기/비동기 SGD에 의존하는 기존 시스템 사이의 격차를 메우기 위해.
• 양자화, 탈중앙화, 통신 지연과 같은 다양한 시스템 이완 기법을 네이티브로 지원하는 유연하고 모듈형 통신 라이브러리를 설계하기 위해.
• 통합 최적화 프레임워크를 통해 최신 분산 학습 알고리즘의 효율적이고 확장 가능한 구현을 가능하게 하기 위해.
• 다양한 네트워크 조건과 워크로드에서 다양한 알고리즘과 시스템 이완 기법 간의 상호 교환 관계를 경험적으로 평가하기 위해.

제안 방법

• 다양한 통신 패턴을 지원하기 위해 저수준 집합 연산을 추상화하는 모듈형이며 MPI 스타일의 통신 라이브러리 설계.
• 세 가지 핵심 시스템 최적화 통합: 계산-통신 겹침(O), 텐서 융합 및 평탄화(F), 계층적 GPU 통신(H).
• 다양한 시스템 이완 기법 지원: 저정밀도 기울기(QSGD, 1-bit Adam), 탈중앙화 학습(Decen), 비동기 실행(Async).
• 사용자가 시스템 수준 최적화에 직접 매핑되는 프리미티브를 조합하고 확장할 수 있는 통합 프레임워크 제공.
• 최적화 스택을 통한 자동 성능 튜닝 기능 제공으로, 모델 및 네트워크 특성에 따라 동적으로 적응 가능.

실험 결과

연구 질문

• RQ1 알고리즘 전용 로직을 하드코딩하지 않고도 다양한 시스템 이완 기법을 네이티브로 지원할 수 있는 통신 라이브러리를 설계할 수 있는가?
• RQ2 양자화, 탈중앙화, 비동기성과 같은 다양한 시스템 이완 기법이 다양한 워크로드와 네트워크 조건에서 끝에서 끝까지의 학습 성능에 어떤 영향을 미치는가?
• RQ3 겹침, 융합, 계층적 통신과 같은 핵심 시스템 최적화 기법들이 전체 학습 효율성에 미치는 상대적 영향은 무엇인가?
• RQ4 특정 모델과 네트워크 환경에서 가장 높은 성능을 내는 알고리즘 설정은 무엇인가?

주요 결과

• BAGUA는 VGG16, BERT, Transformer, LSTM+AlexNet 워크로드 전반에서 PyTorch-DDP, Horovod, BytePS보다 최대 2배 빠른 끝에서 끝까지의 학습 시간을 달성했다.
• 저대역폭 네트워크에서는 QSGD 및 1-bit Adam과 같은 압축 알고리즘으로 인해 통신 오버헤드가 크게 감소하고 성능 향상이 뚜렷했다.
• 고지연 네트워크에서는 동기화 블로킹을 줄일 수 있는 탈중앙화 알고리즘(Decen-32bits/8bits)이 중심화된 방식을 능가했다.
• 제거 분석 결과, 겹침, 융합, 계층적 통신의 세 가지 최적화 기법이 모두 필수적임을 확인했으며, 이들의 영향은 워크로드에 따라 다름: H는 통신 집약적인 작업에서 유리하고, F는 작은 텐서가 많은 모델에 유리하며, O는 계산 집약적인 환경에서 가장 효과적임을 입증함.
• 비동기 학습(Async)은 느린 노드가 존재하는 환경에서 에포크 시간을 30–50% 감소시켜 이성적임을 입증함.
• 경험적 가이드라인 수립: SGD 기반 최적화기에는 QSGD를, Adam에는 1-bit Adam을, 통신 대비 계산 비율이 낮을 경우 비동기 방법을 사용할 것.