[논문 리뷰] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Yanping Huang, Youlong Cheng외 9명|arXiv (Cornell University)|2018. 11. 16.

Digital Imaging for Blood Diseases인용 수 881

한 줄 요약

GPipe는 배치 분할 파이프라인 병렬화를 도입하여 여러 가속기에서 동기화된 그래디언트로 거대한 신경망을 학습시키며, 계층 시퀀스로 표현될 수 있는 아키텍처에 대해 거의 선형 속도 증가와 확장 가능한 메모리 효율을 가능하게 한다.

ABSTRACT

Scaling up deep neural network capacity has been known as an effective approach to improving model quality for several different machine learning tasks. In many cases, increasing model capacity beyond the memory limit of a single accelerator has required developing special algorithms or infrastructure. These solutions are often architecture-specific and do not transfer to other tasks. To address the need for efficient and task-independent model parallelism, we introduce GPipe, a pipeline parallelism library that allows scaling any network that can be expressed as a sequence of layers. By pipelining different sub-sequences of layers on separate accelerators, GPipe provides the flexibility of scaling a variety of different networks to gigantic sizes efficiently. Moreover, GPipe utilizes a novel batch-splitting pipelining algorithm, resulting in almost linear speedup when a model is partitioned across multiple accelerators. We demonstrate the advantages of GPipe by training large-scale neural networks on two different tasks with distinct network architectures: (i) Image Classification: We train a 557-million-parameter AmoebaNet model and attain a top-1 accuracy of 84.4% on ImageNet-2012, (ii) Multilingual Neural Machine Translation: We train a single 6-billion-parameter, 128-layer Transformer model on a corpus spanning over 100 languages and achieve better quality than all bilingual models.

연구 동기 및 목표

• 단일 가속기의 메모리 한계를 넘는 네트워크의 확장 가능한 학습 가능성 확보.
• 유연하고 태스크 및 아키텍처에 구애받지 않는 모델-병렬화 라이브러리 제공.
• 안정성을 위한 동기화 그래디언트 업데이트가 있는 배치 분할 파이프라인 알고리즘 개발.
• 거대한 모델에 맞춘 재-구성(re-materialization)을 통한 활성화 메모리 감소.
• 이미지 분류 및 다국어 기계번역 태스크에서 확장성 시연.

제안 방법

• L개의 층으로 구성된 순차 신경망을 각기 다른 가속기에 배치된 K개의 셀로 분할한다.
• 크기가 N인 각 미니 배치를 M개의 마이크로 배치로 분할하고 K개의 가속기에 걸쳐 파이프라인한다.
• 매개변수를 업데이트하기 전에 M개의 마이크로 배치에 걸쳐 그래디언트 축적과 함께 동기식 미니 배치 그래디언트 하강을 사용한다.
• 역전파 동안 순전파 활성화를 재계산하도록 재-구성(re-materialization)을 적용하여 피크 메모리를 O(N + (L/K) * (N/M))로 감소시킨다.
• 파이프라인 효율을 균형 있게 구성하고 버블 오버헤드를 줄이기 위해 추정 셀 비용의 분산 최소화(K-1)/(M+K-1).
• 추가 확장을 위해 GPipe를 데이터 병렬성과 선택적으로 결합할 수 있다.

실험 결과

연구 질문

• RQ1 GPipe가 단일 가속기의 메모리 한계를 넘어 층의 시퀀스로 표현될 수 있는 임의의 깊은 네트워크를 확장할 수 있는가?
• RQ2 마이크로 배치 분할(M)과 파티션 수(K)이 학습 효율, 메모리, 확장성에 미치는 영향은 무엇인가?
• RQ3 재-구성(re-materialization)이 거대한 모델에서 메모리 사용량과 학습 안정성에 미치는 영향은 무엇인가?
• RQ4 합성곱 기반(AmoebaNet) 및 트랜스포머 기반(다국어 Transformer) 아키텍처에서 어떤 성능 및 처리량 향상을 달성할 수 있는가?
• RQ5 층 분포가 불균형한 네트워크 아키텍처가 GPipe 아래에서 속도에 어떤 영향을 미치는가?

주요 결과

• GPipe는 8개의 가속기에서 최대 1.8B AmoebaNet 매개변수를 학습할 수 있게 하여 단일 가속기 한계를 크게 확장한다.
• 6B 매개변수 다국어 Transformer를 128개의 가속기에 걸쳐 훈련할 수 있으며, 100개 언어에 걸쳐 이중언어 베이스라인보다 향상을 달성한다.
• M이 최소한 4×K일 때 버블 오버헤드가 무시 가능해져 Transformer 워크로드에서 거의 선형 속도 증가를 얻는다.
• 재-구성(re-materialization)을 통해 메모리를 크게 감소시켜 주어진 하드웨어에서 훨씬 더 큰 모델이 가능하도록 한다.
• 실험에서 557M 매개변수 AmoebaNet으로 ImageNet에서 강력한 결과를 보였고, 데이터셋 간 전이 학습도 경쟁력 있다.
• 다국어 MT에서 깊이-너비 트레이드오프가 존재한다; 더 깊은 모델이 저자원 언어에서 우수하고, 매우 큰 모델이 고자원 언어를 개선한다.