[論文レビュー] Scale MLPerf-0.6 models on Google TPU-v3 Pods

Sameer Kumar, Victor Bitorff+10 more|arXiv (Cornell University)|Sep 21, 2019

Advanced Data Storage Technologies参考文献 11被引用数 35

ひとこと要約

論文は、分散評価、勾配和算の加速、モデル並列化、ウェイト更新のシャーディングなどの最適化を用いて、Google TPU-v3 Pod（1024チップ）へMLPerf-0.6トレーニングベンチマークをスケールさせ、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ポッド（1024チップ）でMLPerf-0.6トレーニングのスケーラブル性をデモンストレーションし、大規模でのボトルネックを特定する。
• スループットを向上させつつ目標精度を維持する最適化技術を開発・検証する。
• モデル並列性とコンパイラ支援最適化を活用して、多様なMLワークロードのスケーリングを可能にする。

提案手法

• XLAコンパイラを使用してTensorFlowグラフを低減化し、TPU-v3の実行を最適化する。
• 計算には混合精度（bfloat16）を使用し、精度を維持するために重要な非畳み込み演算を32-bitのままにする。
• 分散評価、2-D勾配和算、ウェイト更新シャーディングなど一連の最適化を適用し、大規模ポッドでのボトルネックを削減する。
• モデル並列性技術（空間分割とハロー交換）を用いてコア間の同時実行性を高める。
• TransformerとGNMTに対してデータ並列とインメモリ評価を活用し、インフラのオーバーヘッドを緩和する。
• 評価オーバーヘッドを低減し、入力パイプラインを最適化して高いスループットを維持する。

実験結果

リサーチクエスチョン

• RQ1 Google TPU-v3ポッド（最大1024チップ）でMLPerf-0.6ベンチマークはどのようにスケールできるか？
• RQ2 スケールでMLPerfの精度目標を維持しつつピークスループットを達成するために必要な最適化は何か？
• RQ3 ResNet-50、SSD、Mask-RCNN、Transformer、GNMTの大規模スケーリングを最も効果的に実現するモデル特有の並列化とオプティマイザ戦略は何か？

主な発見

• TPU-v3ポッドを最大1024チップまで大規模スケーリングを達成し、複数のベンチマークで記録的な性能を達成。
• 評価を分散させ、トレーニング/評価ループを密結合することでボトルネックを削減し、スケーラビリティを向上させる。
• パイプライン化されたデータ転送を伴う2-D勾配和算は大幅な速度向上をもたらし、ResNet-50で1.5倍超。
• モデル並列性技術（空間分割とウェイト更新シャーディング）は、コアあたりの作業量が小さい場合と大規模パラメータを持つTransformerのスケーリングを可能にする。
• SSDは4コアで1.6xの速度アップを達成し、2048 TPUコアへのスケーリングを実現。Mask-RCNNは第一段と第二段のマルチコア分散の恩恵を受けた。
• TransformerとGNMTのスケーリングには、インメモリ評価、最大シーケンス長の削減、メモリ/計算最適化を含む専用の最適化が必要で、大規模バッチでの目標精度達成へとつながった。