[논문 리뷰] Multi-node Bert-pretraining: Cost-efficient Approach

Jiahuang Lin, Xin Li외 1명|arXiv (Cornell University)|2020. 08. 01.

Topic Modeling참고 문헌 26인용 수 7

한 줄 요약

이 논문은 256개의 NVIDIA T4 GPU를 탑재한 32노드 대학 클러스터에서 10 Gbps 네트워크를 사용해 12일 만에 BERT-large를 다중노드로 효율적으로 미리 훈련할 수 있는 비용 효율적인 프레임워크를 제안한다. 알고리즘적 및 소프트웨어 최적화를 통해 단일 장치의 처리량을 향상시키고, 훈련 워크로드를 효율적으로 분배하며, 네트워크 통신 병목 현상을 완화하여 70%의 약한 스케일링 효율을 달성했으며, 산업용 시스템에 비해 비용은 훨씬 낮게 모델 성능은 유사하게 유지한다.

ABSTRACT

Recently, large scale Transformer-based language models such as BERT, GPT-2, and XLNet have brought about exciting leaps in state-of-the-art results for many Natural Language Processing (NLP) tasks. One of the common trends in these recent models is a significant increase in model complexity, which introduces both more weights and computation. Moreover, with the advent of large-scale unsupervised datasets, training time is further extended due to the increased amount of data samples within a single training epoch. As a result, to train these models within a reasonable time, machine learning (ML) programmers often require advanced hardware setups such as the premium GPU-enabled NVIDIA DGX workstations or specialized accelerators such as Google's TPU Pods. Our work addresses this limitation and demonstrates that the BERT pre-trained model can be trained within 2 weeks on an academic-size cluster of widely available GPUs through careful algorithmic and software optimizations. In this paper, we present these optimizations on how to improve single device training throughput, distribute the training workload over multiple nodes and GPUs, and overcome the communication bottleneck introduced by the large data exchanges over the network. We show that we are able to perform pre-training on BERT within a reasonable time budget (12 days) in an academic setting, but with a much less expensive and less aggressive hardware resource requirement than in previously demonstrated industrial settings based on NVIDIA DGX machines or Google's TPU Pods.

연구 동기 및 목표

• 제한된 하드웨어 예산을 가진 대학 환경에서 효율적이고 비용 효율적인 BERT-large 미리 훈련을 가능하게 하기 위해.
• BERT-large와 같은 대규모 트랜스포머 모델을 훈련하는 데 높은 계산 및 재정적 장벽을 해결하기 위해.
• 일반 소비자용 하드웨어에서 단일 GPU 처리량과 다중노드 분산 훈련 성능을 최적화하기 위해.
• 기울기 누적과 시스템 수준의 튜닝을 통해 분산 훈련의 통신 오버헤드를 감소시키기 위해.
• 산업 기준 대비 훨씬 낮은 하드웨어 비용으로도 최신 수준의 모델 품질을 달성할 수 있음을 입증하기 위해.

제안 방법

• 노드당 8개의 NVIDIA T4 GPU를 탑재한 32노드 클러스터를 사용하여 총 256개의 GPU를 확보하고, 10 Gbps 네트워크로 연결.
• 노드 간 통신 빈도와 네트워크 트래픽을 줄이기 위해 4단계에 걸쳐 기울기 누적 기법을 적용.
• 혼합 정밀도 훈련과 효율적인 메모리 관리를 통해 단일 GPU 훈련 처리량을 최적화.
• 이중 단계 미리 훈련: 1단계는 문장 길이 128(36 에포크), 2단계는 문장 길이 512(6 에포크)로, 각각 다른 배치 크기와 학습률을 사용.
• 노드 간 모델 병렬화와 데이터 병렬화를 구현하고, 파arameter 업데이트에 동기식 확률적 경사 하강법을 적용.
• 수렴 문제 진단과 훈련 안정성 향상을 위해 손실 곡선 분석과 하이퍼파ram터 튜닝을 수행.

실험 결과

연구 질문

• RQ1 학술용 하드웨어와 10 Gbps 네트워크만을 사용해 BERT-large를 12일 이내로 미리 훈련시킬 수 있는가?
• RQ2 일반 소비자용 다중 GPU 클러스터에서 높은 스케일링 효율을 달성하기 위해 필요한 시스템 수준 최적화는 무엇인가?
• RQ3 기울기 누적은 분산 BERT 미리 훈련에서 통신 오버헤드와 훈련 처리량에 어떤 영향을 미치는가?
• RQ4 왜 BERT 미리 훈련의 2단계에서 수렴 불안정성이 발생하는가? 하드웨어 비용을 늘리지 않고도 이를 해결할 수 있는가?
• RQ5 고비용의 DGX나 TPU 시스템 대비 저비용 T4 GPU에서 훈련했을 때 모델 성능는 어느 정도 유지될 수 있는가?

주요 결과

• 저자들은 256개의 T4 GPU를 사용한 32노드 대학 클러스터에서 BERT-large 미리 훈련을 12일 만에 완료했으며, 70%의 약한 스케일링 효율을 달성했다.
• 총 하드웨어 비용은 약 $624K로, 동등한 DGX나 TPU Pod 설정의 $15M–$12.8M에 비해 24배 낮았다.
• 微조정 후 SQuAD v1.1에서 F1 스코어가 81%에서 83%를 기록했으며, 이는 Google의 90.9%와 NVIDIA의 90%–91%에 비해 9%–10% 이내의 성능였다.
• 2단계에서의 수렴 문제 원인은 하이퍼파ram터 잘못 설정이었으며, 시스템 최적화 때문이 아니었다. 더 나은 튜닝을 통해 2단계는 6 에포크에서 4 에포크로 단축 가능했다.
• 최적화된 실행과 비최적화 실행의 손실 곡선 분석 결과, 최적화가 수렴에 악영향을 주지 않았고, 오히려 2단계에서 더 안정적인 성능을 보였다.
• 12일 동안 256개의 T4 GPU를 Google 클라우드에서 렌털하는 데에는 $25,739가 소요되었으며, 이는 하드웨어 구매 자본 비용의 24배 낮은 금액이지만, 하드웨어 수명이 3년이므로 약 90회의 실험을 수행할 수 있었다.