[論文レビュー] CosmoFlow: Using Deep Learning to Learn the Universe at Scale

Amrita Mathuriya, Deborah Bard+15 more|arXiv (Cornell University)|Aug 14, 2018

Parallel Computing and Optimization Techniques参考文献 14被引用数 24

ひとこと要約

CosmoFlow は、TensorFlow に基づくスケーラブルなディープラーニングフレームワークを提案し、3D ダークマター・シミュレーションから宇宙論的パラメータ（$\Omega_M$, $\sigma_8$, $n_s$）を予測する。8192ノードのCoriスーパーコンピュータ上で、3D畳み込みニューラルネットワークを最適化し、完全同期型データ並列学習を実現することで、3.5 Pflop/sの持続的性能と77%の並列効率を達成し、パrameter推定において前例のない精度を実現した。

ABSTRACT

Deep learning is a promising tool to determine the physical model that describes our universe. To handle the considerable computational cost of this problem, we present CosmoFlow: a highly scalable deep learning application built on top of the TensorFlow framework. CosmoFlow uses efficient implementations of 3D convolution and pooling primitives, together with improvements in threading for many element-wise operations, to improve training performance on Intel(C) Xeon Phi(TM) processors. We also utilize the Cray PE Machine Learning Plugin for efficient scaling to multiple nodes. We demonstrate fully synchronous data-parallel training on 8192 nodes of Cori with 77% parallel efficiency, achieving 3.5 Pflop/s sustained performance. To our knowledge, this is the first large-scale science application of the TensorFlow framework at supercomputer scale with fully-synchronous training. These enhancements enable us to process large 3D dark matter distribution and predict the cosmological parameters $Ω_M$, $σ_8$ and n$_s$ with unprecedented accuracy.

研究の動機と目的

• 大規模な3D ダークマター シミュレーションからの宇宙論的パラメータ推定における計算ボトルネックを解決すること。
• エクサスケール HPC プラットフォーム上で、ディープラーニングを用いて宇宙論的パラメータ（$\Omega_M$, $\sigma_8$, $n_s$）の高精度な回帰を可能にすること。
• スケーラブルなディープラーニングのためのソフトウェアスタック全般（ネットワークアーキテクチャ、I/O、通信、プリミティブ）を最適化し、CPUベースのスーパーコンピュータ上で効率を発揮すること。
• TensorFlow を用いて、スーパーコンピュータ上で大規模な完全同期型データ並列学習を実証すること。
• スケールの大きなディープラーニングを用いて、迅速かつ高精度に宇宙論的データの科学的探索を可能にすること。

提案手法

• Ravanbakhsh ら（2017）が提案した3D畳み込みニューラルネットワークアーキテクチャを、$128^3$ボクセル入力データに適応し、3つの宇宙論的パラメータを予測する。
• Intel MKL-DNN を用いて TensorFlow フレームワーク内に最適化された3D畳み込みおよびプーリングプリミティブを実装し、CPU最適化性能を実現した。
• Cray PE マシンラーニングプラグインを活用し、数千ノードにわたる効率的なMPIベースのデータ並列性を実現した。
• I/Oボトルネックを軽減するため、Lustreとは対照的に高帯域幅のSSDベースのファイルシステムを採用し、スケーリング効率を向上させた。
• 全パラメータの平均化を伴う同期的確率的勾配降下法（SSGD）を用い、大規模スケールでの収束を保証した。
• Intel Xeon Phi（KNL）プロセッサに最適化されたHPCのベストプラクティスを用いて、データI/O、モデル学習、ノード間通信に至るまで、スタック全体を最適化した。

実験結果

リサーチクエスチョン

• RQ1 3D ダークマター シミュレーションで学習したディープラーニングモデルは、宇宙論的パラメータ $\Omega_M$, $\sigma_8$, および $n_s$ を正確に予測できるか？
• RQ2 8192ノードを有するスーパーコンピュータ上で、宇宙論のためのディープラーニングモデルを学習する際、どの程度のパフォーマンスとスケーリング効率が達成できるか？
• RQ3 I/Oシステムの選択（例：Lustre 対 SSDベース）が、大規模ディープラーニングワークロードにおけるスケーリング効率とパフォーマンスにどのように影響するか？
• RQ4 完全同期型データ並列学習が、極めて大規模（8192ノード）なスケールでも収束性と正確性を維持できるか、その程度はどの程度か？
• RQ5 最適化されたCPUベースのディープラーニングフレームワークは、科学的応用分野でエクサスケールレベルのパフォーマンスを達成できるか？

主な発見

• CosmoFlow は、Intel Xeon Phi プロセッサを搭載したCoriスーパーコンピュータの8192ノードで、3.5 Pflop/sの持続的単精度パフォーマンスを達成した。
• スケールで77%の並列効率を示し、TensorFlow を用いたスーパーコンピュータ上での大規模、完全同期型データ並列学習の初の実証となった。
• 8192ノードでの学習は収束に近い結果に達し、高い計算負荷にもかかわらず、科学的に意味のある予測を生成した。
• 2048ノードでモデルは完全に収束し、ダークマター分布からの宇宙論的パラメータ推定において、最高水準の科学的精度を達成した。
• 高帯域幅のSSDベースのファイルシステムの採用により、I/Oパフォーマンスが著しく向上し、Lustreで観察されたスケーリング劣化が緩和された。
• 単一ノードでのパフォーマンスはKNLノードで535 Gflop/sに達し、フルスタック最適化により1.4TBを超える宇宙論的シミュレーションデータを効率的に処理できた。