[논문 리뷰] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Dmitry Lepikhin, HyoukJoong Lee외 7명|arXiv (Cornell University)|2020. 06. 30.

Topic Modeling참고 문헌 75인용 수 348

한 줄 요약

GShard는 일반적인 주석 기반 접근법과 XLA 기반 SPMD 컴파일러를 도입하여 Sparsely-Gated MoE 계층을 가진 거대 Transformer 모델을 학습시키고, 계산 증가를 부분선형으로 달성하며 2048 TPUs에서 4일 만에 600B 파라미터 다국어 번역을 가능하게 한다.

ABSTRACT

Neural network scaling has been critical for improving the model quality in many real-world machine learning applications with vast amounts of training data and compute. Although this trend of scaling is affirmed to be a sure-fire approach for better model quality, there are challenges on the path such as the computation cost, ease of programming, and efficient implementation on parallel devices. GShard is a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler. It provides an elegant way to express a wide range of parallel computation patterns with minimal changes to the existing model code. GShard enabled us to scale up multilingual neural machine translation Transformer model with Sparsely-Gated Mixture-of-Experts beyond 600 billion parameters using automatic sharding. We demonstrate that such a giant model can efficiently be trained on 2048 TPU v3 accelerators in 4 days to achieve far superior quality for translation from 100 languages to English compared to the prior art.

연구 동기 및 목표

• 신경망의 확장이 모델 품질을 향상시키는 필요성을 촉진하고, 계산, 프로그래밍 용이성, 병렬 배치의 실용적 도전을 해결한다.

제안 방법

• Transformer를 Position-wise Sparsely-Gated Mixture-of-Experts (MoE) 계층을 추가하여 부분선형 계산 규모 확장을 달성하도록 확장한다.
• 가벼운 주석 API 모듈과 자동 병렬화를 위한 XLA 컴파일러 확장이 포함된 GShard를 도입한다.
• 장치 수에 관계없이 컴파일 시간을 일정하게 유지하기 위한 SPMD (Single Program Multiple Data) 파티셔닝 전략을 채택한다.
• 모델 개발자가 단일 거대한 장치에서 작성하는 것처럼 자동 파티셔닝이 컴파일러에 의해 적용되는 설계를 제공한다.
• 수천 개의 전문가에 걸친 부하를 균형 있게 분배하기 위해 전문가 용량 제약이 있는 게이팅 메커니즘과 보조 손실을 사용한다.
• 100개 언어쌍의 다국어 기계 번역에서 엔드-투-엔드 학습 및 확장을 시연한다.

실험 결과

연구 질문

• RQ1 초거대 Transformer 모델을 수천 대의 디바이스에 걸쳐 비대칭적 컴파일이나 통신 오버헤드 없이 어떻게 효율적으로 학습시킬 수 있는가?
• RQ2 Sparsely-Gated MoE 계층을 통한 조건부 계산이 모델 용량이 증가함에 따라 부분선형 계산 증가를 제공할 수 있는가?
• RQ3 주석 기반의 GShard 접근법이 XLA에서 자동적이고 확장 가능한 파티셔닝을 가능하게 하면서 모델 개발을 단순화하는가?
• RQ4 다국어 설정에서 수백 억 개의 파라미터로 확장할 때 번역 품질의 실제 향상은 무엇인가?

주요 결과

• 600B-parameter Sparsely-Gated MoE Transformer가 2048 TPU v3 디바이스에서 4일간 학습되어 100개 언어를 English로 번역하는 데 있어 기존 연구보다 우수한 품질을 달성했다.
• 학습 비용은 모델 크기에 따라 부분선형으로 증가하여, 증가하는 용량에 비해 계산이 부분선형으로 확장됨을 보여준다.
• 밀집 기반 Transformer(2.3B params)는 235.5 TPU v3 core-years가 필요했고, MoE 접근법의 효율 향상을 강조한다.
• GShard는 자동 파티셔닝을 가능하게 하고 SPMD 접근법으로 수천 대의 디바이스에 확장하며 컴파일 시간을 O(1)로 유지한다.
• MoE 게이팅은 전문가 용량 제약, 부하를 균형 있게 만들기 위한 보조 손실, 용량을 효과적으로 활용하기 위한 두 번째 최선 전문가를 위한 무작위 라우팅을 포함한다.