[논문 리뷰] Maximizing Parallelism in Distributed Training for Huge Neural Networks

Zhengda Bian, Qifan Xu외 2명|arXiv (Cornell University)|2021. 05. 30.

Advanced Neural Network Applications참고 문헌 14인용 수 20

한 줄 요약

이 논문은 선형 레이어의 계산을 텐서, 시퀀스, 헤드의 세 가지 차원에 걸쳐 분산시켜 완벽한 로드 밸런싱을 달성함으로써, 특히 트랜스포머 언어 모델을 위한 기존의 1차원 및 2차원 병렬화 방식보다 더 높은 성능을 보이는 3차원 모델 병렬화 기법을 제안한다. 이는 메모리 및 통신 비용을 감소시키며, 64개의 V100 GPU에서 1-D 병렬화 대비 2.32배, 2-D 병렬화 대비 1.57배의 속도 향상을 보였다.

ABSTRACT

The recent Natural Language Processing techniques have been refreshing the state-of-the-art performance at an incredible speed. Training huge language models is therefore an imperative demand in both industry and academy. However, huge language models impose challenges to both hardware and software. Graphical processing units (GPUs) are iterated frequently to meet the exploding demand, and a variety of ASICs like TPUs are spawned. However, there is still a tension between the fast growth of the extremely huge models and the fact that Moore's law is approaching the end. To this end, many model parallelism techniques are proposed to distribute the model parameters to multiple devices, so as to alleviate the tension on both memory and computation. Our work is the first to introduce a 3-dimensional model parallelism for expediting huge language models. By reaching a perfect load balance, our approach presents smaller memory and communication cost than existing state-of-the-art 1-D and 2-D model parallelism. Our experiments on 64 TACC's V100 GPUs show that our 3-D parallelism outperforms the 1-D and 2-D parallelism with 2.32x and 1.57x speedup, respectively.

연구 동기 및 목표

• 단일 장치의 메모리 및 계산 능력이 초과하는 초거대 언어 모델을 훈련하는 데 발생하는 점점 커지는 과제를 해결하기 위해.
• 선형 레이어의 3차원 분해를 도입하여 분산 모델 병렬화에서의 통신 및 메모리 오버헤드를 줄이기 위해.
• 초거대 트랜스포머의 분산 훈련에서 GPU 간 완벽한 로드 밸런싱을 달성하기 위해.
• 약한 스케일링 및 강한 스케일링 시나리오 모두에서 기존의 1-D 및 2-D 모델 병렬화 기법을 능가하기 위해.

제안 방법

• 선형 레이어를 텐서, 시퀀스, 헤드 차원에 따라 분해하여 계산 및 통신을 균형 잡는 3차원 모델 병렬화 알고리즘을 제안한다.
• 무작위 대기 시간을 최소화하고 모든 장치에서 GPU 활용도를 극대화하기 위해 로드 밸런스가 맞춰진 데이터 레이아웃을 사용한다.
• 가중치 행렬과 활성화 값을 3차원 GPU 격자에 나누어 저장함으로써 통신 부담을 줄이기 위해 3D 타일링 전략을 적용한다.
• PyTorch의 분산 통신 백엔드와 통합하고, 메모리 점유율을 줄이기 위해 혼합 정밀도 훈련을 지원한다.
• 계산과 통신의 오버랩을 통해 지연 시간을 숨기기 위해 고유의 통신 스케줄링 전략을 구현한다.
• 기존 PyTorch 구현을 확장하여 3D 텐서 병렬화를 지원하는 3D 병렬 트랜스포머를 구현한다.

실험 결과

연구 질문

• RQ1 3차원 모델 병렬화 설계는 대규모 트랜스포머의 분산 훈련에서 1-D 또는 2-D 접근 방식보다 더 나은 로드 밸런싱을 달성할 수 있는가?
• RQ2 기존의 1-D 및 2-D 방법에 비해 3-D 병렬화는 통신 및 메모리 비용을 줄일 수 있는가?
• RQ3 대규모 GPU 클러스터에서 3-D 접근 방식은 약한 스케일링 및 강한 스케일링 시나리오에서 어떻게 스케일링되는가?
• RQ4 3-D 병렬화는 모델 정확도를 유지하면서도 더 높은 훈련 스루풋을 달성할 수 있는가?

주요 결과

• 64개의 V100 GPU에서 3-D 병렬화는 강한 스케일링에서 1-D 모델 병렬화 대비 2.32배의 속도 향상을 기록했다.
• 동일한 하드웨어 환경에서 3-D 접근 방식은 2-D 모델 병렬화 대비 1.57배의 속도 향상을 보였다.
• 약한 스케일링에서 3-D 방법은 평균 스텝 시간의 증가율이 가장 낮아, 뛰어난 확장성과 낮은 통신 오버헤드를 나타낸다.
• 모든 GPU 수에서 3-D 병렬화는 가장 작은 평균 스텝 시간을 유지하여 최적의 로드 밸런싱을 입증한다.
• 계산 및 데이터를 GPU 간에 더 균형 잡힌 분포로 분배함으로써 메모리 및 통신 비용을 감소시킨다.
• 특히 대규모 환경에서 3-D 방법은 전방 및 역전파 단계의 시간 모두에서 1-D 및 2-D 접근 방식을 능가한다.