[論文レビュー] Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

Xianyan Jia, Shutao Song+12 more|arXiv (Cornell University)|Jul 30, 2018

Advanced Neural Network Applications参考文献 24被引用数 313

ひとこと要約

この論文は、LARSを用いた混合精度トレーニングをスケーラブルに実装し、64Kミニバッチと最適化されたall-reduceにより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を使用し、安定更新のためマスターペ weightsをFP32で保持。
• BNのバイアスとBNパラメータのウェイト減衰を排除し、AlexNetの収束を改善するため追加のBN層を導入。
• テンサーフュージョンと階層型とリング型を組み合わせたハイブリッドなall-reduce戦略を開発し、勾配集約のスケーラビリティを向上。
• RoCEv2とGPUDirect RDMAを活用して1024-および2048-GPU設定で通信遅延を削減し帯域を改善。

実験結果

リサーチクエスチョン

• RQ1 LARSを用いた混合精度トレーニングは64KまでのミニバッチサイズでImageNetの精度を維持できるか?
• RQ2 極めて大きなミニバッチで収束を保つにはどのような構造的・最適化の調整が必要か?
• RQ3 大規模GPUクラスターで高いスケーラビリティを発揮するために全-reduce戦略をどう最適化できるか?
• RQ4 AlexNetとResNet-50の全体的な学習時間に対する通信最適化の影響はどれほどか?

主な発見

• LARSを組み込んだ混合精度トレーニングは64KミニバッチでResNet-50のトップ1精度を90エポックで保持（LARSありで76.2%）。
• AlexNetは64Kミニバッチと対象のアーキテクチャ調整で95エポック後にトップ1精度58.8%を達成。
• 1024および2048 Tesla P40 GPUsで、それぞれ4分（AlexNet、95エポック）、6.6分（ResNet-50、90エポック）の学習時間を実現。
• このアプローチは1024-GPUクラスターでのAlexNetとResNet-50のNCCLベースの学習に対して最大3倍および11倍のスピードアップをもたらす。
• 全体として、ResNet-50の学習は2048GPUで6.6分で75.8%トップ1精度に達し、前技術と競合する効率で75.8%を達成。
• optimized all-reduceとtensor fusionを使用した場合、1024GPUでのスケーリング効率が9.0%から99.2%へ改善。