[논문 리뷰] Scale MLPerf-0.6 models on Google TPU-v3 Pods

Sameer Kumar, Victor Bitorff외 10명|arXiv (Cornell University)|2019. 09. 21.

Advanced Data Storage Technologies참고 문헌 11인용 수 35

한 줄 요약

이 논문은 distributed evaluation, gradient summation acceleration, 모델 병렬성, 가중치 업데이트 샤딩과 같은 최적화를 사용하여 MLPerf-0.6 훈련 벤치마크를 Google TPU-v3 Pods (1024 칩)으로 확장하고 ResNet-50, SSD, Transformer에서 기록적 성능을 달성하는 방법을 보여준다.

ABSTRACT

The recent submission of Google TPU-v3 Pods to the industry wide MLPerf v0.6 training benchmark demonstrates the scalability of a suite of industry relevant ML models. MLPerf defines a suite of models, datasets and rules to follow when benchmarking to ensure results are comparable across hardware, frameworks and companies. Using this suite of models, we discuss the optimizations and techniques including choice of optimizer, spatial partitioning and weight update sharding necessary to scale to 1024 TPU chips. Furthermore, we identify properties of models that make scaling them challenging, such as limited data parallelism and unscaled weights. These optimizations contribute to record performance in transformer, Resnet-50 and SSD in the Google MLPerf-0.6 submission.

연구 동기 및 목표

• TPU-v3 Pods(1024 칩)에서 MLPerf-0.6 훈련의 확장 가능성을 입증하고 대규모에서 병목 현상을 식별한다.
• 처리량을 향상시키고 목표 정확도를 유지하기 위한 최적화 기법을 개발하고 검증한다.
• 다양한 ML 워크로드의 확장을 가능하게 하는 모델 병렬성과 컴파일러 보조 최적화를 선보인다.

제안 방법

• XLA 컴파일러로 TensorFlow 그래프를 최적화하여 TPU-v3 실행을 향상시킨다.
• 계산에 혼합 정밀도(bfloat16)를 사용하되 중요 비합성 연산은 32비트로 유지하여 정확도를 보존한다.
• 분산 평가, 2-D 그래디언트 합산, 가중치 업데이트 샤딩 등 일련의 최적화를 적용하여 대형 포드에서 병목을 줄인다.
• 코어 간 동시성을 높이기 위해 공간 분할(spatial partitioning)과 halo 교환과 같은 모델 병렬화 기법을 적용한다.
• Transformer와 GNMT에 대해 데이터 병렬성과 메모리 내 평가를 사용하여 인프라 오버헤드를 완화한다.
• 평가 오버헤드를 줄이고 입력 파이프라인을 최적화하여 높은 처리량을 유지한다.

실험 결과

연구 질문

• RQ1 Google TPU-v3 Pod(최대 1024 칩)에서 학습할 때 MLPerf-0.6 벤치마크는 어떻게 확장될 수 있는가?
• RQ2 대규모에서 MLPerf 정확도 목표를 유지하면서 피크 처리량을 달성하기 위해 필요한 최적화는 무엇인가?
• RQ3 ResNet-50, SSD, Mask-RCNN, Transformer, GNMT의 대규모 확장을 가장 잘 가능하게 하는 모델 특정 병렬성 및 최적화 전략은 무엇인가?

주요 결과

• TPU-v3 Pods에서 최대 1024칩까지 대규모 확장을 달성했고, 여러 벤치마크에서 기록적 성능을 기록했다.
• 평가 분산 및 학습/평가 루프의 밀착 결합은 병목을 줄이고 확장성을 향상시킨다.
• 피이프라인 데이터 전송이 포함된 2-D 그래디언트 합산은 상당한 속도 향상—ResNet-50에서 1.5배 이상—을 가져온다.
• 모델 병렬화 기법(공간 분할 및 가중치 업데이트 샤딩)은 코어당 작업이 작아지는 경우와 매개변수가 큰 트랜스포머에 대한 확장을 가능하게 한다.
• SSD는 4개 코어에서 1.6배 속도향상을 보였고 2048 TPU 코어 확장을 가능하게 했으며, Mask-RCNN의 확장은 1단계와 2단계의 다중 코어 분산에서 이익을 얻었다.
• Transformer 및 GNMT 확장은 대형 배치에서 목표 정확도에 도달하기 위한 메모리 내 평가, 최대 시퀀스 길이 축소, 메모리/계산 최적화를 포함한 전용 최적화가 필요했다.