[논문 리뷰] Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

Xianyan Jia, Shutao Song외 12명|arXiv (Cornell University)|2018. 07. 30.

Advanced Neural Network Applications참고 문헌 24인용 수 313

한 줄 요약

이 논문은 LARS를 사용한 혼합-정밀도 학습 시스템을 통해 64K 미니배치와 최적화된 모든-감소를 가능하게 하여 ImageNet 모델(AlexNet 및 ResNet-50)을 몇 시간 대신 몇 분 안에 학습시키고, 이전 시스템보다 우수한 성능을 달성합니다.

ABSTRACT

Synchronized stochastic gradient descent (SGD) optimizers with data parallelism are widely used in training large-scale deep neural networks. Although using larger mini-batch sizes can improve the system scalability by reducing the communication-to-computation ratio, it may hurt the generalization ability of the models. To this end, we build a highly scalable deep learning training system for dense GPU clusters with three main contributions: (1) We propose a mixed-precision training method that significantly improves the training throughput of a single GPU without losing accuracy. (2) We propose an optimization approach for extremely large mini-batch size (up to 64k) that can train CNN models on the ImageNet dataset without losing accuracy. (3) We propose highly optimized all-reduce algorithms that achieve up to 3x and 11x speedup on AlexNet and ResNet-50 respectively than NCCL-based training on a cluster with 1024 Tesla P40 GPUs. On training ResNet-50 with 90 epochs, the state-of-the-art GPU-based system with 1024 Tesla P100 GPUs spent 15 minutes and achieved 74.9\% top-1 test accuracy, and another KNL-based system with 2048 Intel KNLs spent 20 minutes and achieved 75.4\% accuracy. Our training system can achieve 75.8\% top-1 test accuracy in only 6.6 minutes using 2048 Tesla P40 GPUs. When training AlexNet with 95 epochs, our system can achieve 58.7\% top-1 test accuracy within 4 minutes, which also outperforms all other existing systems.

연구 동기 및 목표

• 대형 미니배치 학습의 일반화 위험을 다루면서 처리량을 높인다.
• 매우 큰 미니배치 크기에서도 정확도를 유지하기 위해 LARS를 활용한 혼합-정밀도 학습 개발.
• 수천 대의 GPU에서 확장 가능하도록 통신을 최적화한 all-reduce 알고리즘 설계.
• AlexNet와 ResNet-50에 대해 ImageNet에서 최첨단 학습 속도 시연.
• 실제 하드웨어를 갖춘 대규모 GPU 클러스터에서 수렴성과 확장성 평가.

제안 방법

• 대형 미니배치 크기를 가능하게 하는 LARS 혼합-정밀도 학습 도입.
• 전방/역전 파에 FP16 사용, 업데이트의 안정성을 위한 FP32 마스터 가중치 사용.
• BN 바이어스 및 BN 매개변수에 대한 가중치 감소를 제거하고 AlexNet의 수렴 개선을 위해 추가 BN 계층 도입.
• 텐서 합치기(tensor fusion) 및 계층적 링 기반 접근을 결합한 하이브리드 all-reduce 전략 개발.
• 통신 지연을 줄이고 대역폭을 개선하기 위해 RoCEv2 및 GPUDirect RDMA를 1024- 및 2048-GPU 구성에서 활용.

실험 결과

연구 질문

• RQ1 LARS와 혼합-정밀도 학습이 64K까지의 미니배치에서 ImageNet 정확도를 유지할 수 있는가?
• RQ2 극도로 큰 미니배치에서 수렴을 보장하기 위해 어떤 구조적/최적화 조정이 필요한가?
• RQ3 대규모 GPU 클러스터에서 높은 확장성을 위한 all-reduce 전략 최적화 방법은 무엇인가?
• RQ4 통신 최적화가 AlexNet 및 ResNet-50의 전체 학습 시간에 미치는 영향은 무엇인가?

주요 결과

• LARS를 포함한 혼합-정밀도 학습은 64K 미니배치에서 90 에포크의 ResNet-50 상단-1 정확도(76.2% with LARS)를 보존합니다.
• 64K 미니배치와 목표된 구조 조정으로 AlexNet은 95 에포크 후 상단-1 정확도 58.8%에 도달합니다.
• 시스템은 1024 및 2048 Tesla P40 GPU에서 각각 4분(AlexNet, 95 에포크) 및 6.6분(ResNet-50, 90 에포크) 학습 시간을 달성합니다.
• 이 접근법은 1024-GPU 클러스터에서 NCCL 기반 학습 대비 AlexNet 및 ResNet-50에 대해 최대 3배 및 11배의 속도 향상을 보여줍니다.
• 전반적으로 ResNet-50 학습은 2048 GPU에서 6.6분 만에 75.8% 상단-1 정확도에 도달하고, 75.8%는 기존 연구에 비해 경쟁력 있는 효율로 달성됩니다.
• 최적화된 all-reduce 및 텐서 합성을 사용할 때 1024 GPU에서 스케일링 효율이 9.0%에서 99.2%로 개선됩니다.