[논문 리뷰] Whale: A Unified Distributed Training Framework.

Ang Wang, Xianyan Jia외 4명|arXiv (Cornell University)|2020. 11. 18.

Advanced Neural Network Applications참고 문헌 39인용 수 4

한 줄 요약

Whale은 새로운 IR 추상화와 메타 기반 비용 모델을 통한 자동 병렬화를 통해 데이터 병렬, 모델 병렬, 파이프라인 병렬, 연산자 분할, 하이브리드 병렬 전략을 모두 지원하는 통합된 분산 학습 프레임워크이다. 64개의 GPU에서 BERT-Large 학습을 Horovod 데이터 병렬 기반 대비 2.32배 빠르게 하고, 대규모 이미지 분류 작업을 14.8배 빠르게 하며, TensorFlow와의 호환성을 유지하면서 코드 수정을 최소화한다.

ABSTRACT

Data parallelism (DP) has been a common practice to speed up the training workloads for a long time. However, with the increase of data size and model size, DP has become less optimal for most distributed training workloads. Moreover, it does not work on models whose parameter size cannot fit into a single GPU's device memory. To enable and further improve the industrial-level giant model training, we present Whale, a unified distributed training framework. It provides comprehensive parallel strategies including data parallelism, model parallelism, operator sharding, pipeline, hybrid strategy, and automatic parallel strategy. To express complex training strategies effectively and efficiently in one framework, Whale IR is designed as the basic unit to explore and implement different distributed strategies. Moreover, Whale enables automatic parallelism upon using a meta-driven cost model. Whale is compatible with TensorFlow and can easily distribute training tasks by adding a few code lines without changing user model code. To the best of our knowledge, Whale is the first work that can support various hybrid distributed strategies within one framework. In our experiment of Bert Large model, Whale pipeline strategy is 2.32 times faster than Horovod data parallelism (HDP) on 64 GPUs. In a large-scale image classification task (100,000 classes), Whale hybrid strategy, which consists of operator sharding and DP, is 14.8 times faster than HDP on 64 GPUs.

연구 동기 및 목표

• 대규모 모델과 데이터셋에 대한 스케일링에서 데이터 병렬의 한계를 해결한다.
• 단일 GPU 메모리 용량을 초과하는 초거대 모델의 효율적 학습을 가능하게 한다.
• 단일 시스템 내에서 다양한 병렬 전략을 지원하는 통합 프레임워크를 제공한다.
• 메타 기반 비용 모델을 활용해 병렬 전략 선택을 자동화하여 학습 성능을 최적화한다.
• 기존 TensorFlow 모델과의 후행 호환성을 확보하여 최소한의 코드 수정으로도 구현 가능하도록 한다.

제안 방법

• 복잡한 분산 학습 전략을 표현하고 조합하기 위한 기초 추상화로 Whale IR를 설계한다.
• 포괄적인 병렬 전략을 구현: 데이터 병렬, 모델 병렬, 연산자 분할, 파이프라인 병렬, 하이브리드 조합.
• 메타 기반 비용 모델을 사용해 최적의 전략을 선택하는 자동 병렬화 엔진을 통합한다.
• 최소한의 코드 수정으로 분산 학습을 배포할 수 있도록 TensorFlow와의 원활한 통합을 보장한다.
• 고수준 모델을 효율적인 분산 실행 계획으로 매핑하기 위해 통합된 IR 기반 컴파일 파이프라인을 사용한다.

실험 결과

연구 질문

• RQ1 하이브리드 조합을 포함한 다양한 분산 학습 전략을 효율적으로 지원할 수 있는 단일 프레임워크의 설계는 어떻게 이루어질 수 있는가?
• RQ2 비용 모델 기반 자동 전략 선택이 수동 또는 고정 전략 선택 대비 학습 효율을 얼마나 향상시킬 수 있는가?
• RQ3 BERT-Large와 같은 대규모 모델, 100,000개 클래스의 이미지 분류 작업에서 기존 데이터 병렬 기반 대비 상당한 속도 향상을 달성할 수 있는가?
• RQ4 기존 딥러닝 프레임워크, 특히 TensorFlow와의 호환성을 유지하면서 모델 코드 수정 없이도 어떻게 통합할 수 있는가?
• RQ5 대규모 학습 워크로드에서 연산자 분할과 하이브리드 전략을 통해 달성 가능한 성능 향상은 어느 정도인가?

주요 결과

• Whale의 파이프라인 전략은 BERT-Large 학습에서 64개 GPU에서 Horovod 데이터 병렬 기반 대비 2.32배 빠른 성능을 달성한다.
• Whale의 하이브리드 전략(연산자 분할과 데이터 병렬의 조합)은 100,000개 클래스의 이미지 분류 작업에서 64개 GPU에서 Horovod 데이터 병렬 기반 대비 14.8배 빠른 성능을 달성한다.
• 모델 병렬과 연산자 분할을 활용해 단일 GPU 메모리 용량을 초과하는 모델의 지원에 성공한다.
• 메타 기반 비용 모델을 통한 자동 병렬화로 수동 튜닝 없이도 최적의 전략 선택이 가능하다.
• Whale은 TensorFlow와 완전히 호환되며, 기존 모델에 최소한의 코드 수정으로 분산 학습을 구현할 수 있다.