[논문 리뷰] Boosting Distributed Training Performance of the Unpadded BERT Model

Jinle Zeng, Min Li외 5명|arXiv (Cornell University)|2022. 08. 17.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 변수 길이 입력을 활용하여 패딩 토큰에서 발생하는 중복 계산을 제거함으로써 고도로 최적화된 패딩 없는 BERT 모델을 제안한다. 이는 그룹화된 다중 스트림 FMHA, 데이터 교환을 겹쳐 처리하는 워크로드 균형 조절, 그리고 커널/연산자 최적화를 통해 달성된다. 이 방법은 8x A100 GPU에서 MLPerf Training v2.0에서 1위를 기록하며 상태 기준 최고 수준의 처리량 2578개 샘플/초를 달성했으며, 단지 16.598분 만에 수렴한다.

ABSTRACT

Pre-training models are an important tool in Natural Language Processing (NLP), while the BERT model is a classic pre-training model whose structure has been widely adopted by followers. It was even chosen as the reference model for the MLPerf training benchmark. The distributed training performance optimization of BERT models plays an important role in accelerating the solutions of most NLP tasks. BERT model often uses padding tensors as its inputs, leading to excessive redundant computations. Thus, removing these redundant computations is essential to improve the distributed training performance. This paper designs a new approach to train BERT models with variable-length inputs efficiently. Firstly, we propose a general structure for the variable-length BERT models, and accelerate the encoder layer via our grouped multi-stream FMHA (Fused Multi-Head Attention) method. Secondly, through data exchange, we address the unbalanced workload problem caused by the variable-length inputs, which overlaps highly with the training process. Finally, we optimize the overall performance of the BERT model, such as kernel fusion, and operator optimization. Our experimental results show that our highly optimized BERT model achieves state-of-the-art throughput and ranks first in MLPerf Training v2.0 within the same GPU configuration. The optimizations in this paper can be applied to more BERT-like models in our future works.

연구 동기 및 목표

• BERT 학습에서 패딩 토큰으로 인한 중복 계산을 제거하여 분산 학습 성능 저하를 방지하기 위해.
• 변수 길이 입력으로 인한 분산 학습 환경에서의 워크로드 불균형 문제를 해결하기 위해.
• 동일 하드웨어 구성에서 BERT-Large에 대해 최고 수준의 학습 처리량을 달성하기 위해.
• 어텐션 메커니즘, 커널 융합, 연산자 수준의 성능 최적화를 포함한 전체 학습 파이프라인을 최적화하기 위해.

제안 방법

• 패딩 없이도 동적 시퀀스 길이를 지원할 수 있는 변수 길이 BERT 모델을 위한 일반적인 아키텍처를 설계한다.
• 어텐션 계산을 최적화하여 인코더 레이어의 처리 속도를 향상시키기 위해 그룹화된 다중 스트림 FMHA(Fused Multi-Head Attention)를 도입한다.
• GPU 계산과 겹쳐 데이터 교환을 처리하여 유휴 시간과 통신 오버헤드를 최소화하고 워크로드 균형을 조절한다.
• 메모리 대역폭을 줄이고 GPU 활용도를 향상시키기 위해 커널 융합 및 연산자 수준 최적화를 적용한다.
• 혼합 정밀도 학습을 적용하고 LAMB 최적화기를 사용하며, 모든 장치 간 일관된 배치 처리를 보장한다.
• 딥러닝 프레임워크로 PaddlePaddle를 사용하고, 8x NVIDIA A100 400W GPU에서 벤치마크를 수행한다.

실험 결과

연구 질문

• RQ1 BERT 학습에서 패딩 토큰으로 인한 중복 계산를 효과적으로 제거할 수 있는 방법은 무엇인가?
• RQ2 변수 길이 입력이 분산 학습 환경에서 장치 간 워크로드 불균형을 초래할 경우 이를 균형 조절하기 위한 기법은 무엇인가?
• RQ3 계산과 겹쳐 데이터 교환을 처리하면 분산 학습 효율성이 어떻게 향상되는가?
• RQ4 커널 및 연산자 최적화의 조합 중 어떤 조합이 BERT-Large에 대해 가장 높은 처리량을 낳는가?
• RQ5 완전히 패딩 없는 BERT 모델은 MLPerf와 같은 표준 벤치마크에서 기존 최적화된 구현보다 성능이 뛰어나다고 할 수 있는가?

주요 결과

• 최적화된 패딩 없는 BERT 모델은 2578개 샘플/초의 처리량을 달성하여 동일한 GPU 구성에서 모든 다른 구현보다 뛰어난 성능을 보였다.
• 모델은 MLPerf Training v2.0에서 1위를 기록했으며, 16.598분 만에 72% MLM 정확도에 수렴하여 NVIDIA, HazyResearch 및 기타 제출물들을 압도했다.
• 패딩 관련 계산을 제거하기만 해도 베이스라인 패딩된 구현 대비 2.3배의 속도 향상을 달성했다.
• 커널 융합 최적화는 8.9%의 성능 향상을, 연산자 수준 최적화는 11.3%의 성능 향상을 각각 기여했다.
• 계산과 겹쳐 데이터 교환을 처리함으로써 2.8%의 성능 향상이 있었으며, 통신 병목 현상을 줄였다.
• 그룹화된 다중 스트림 FMHA 최적화는 표준 unpad FMHA 대비 3.6%의 성능 향상을 제공했다.