[論文レビュー] High Performance Monte Carlo Simulation of Ising Model on TPU Clusters

Kun Yang, Yifan Chen+3 more|arXiv (Cornell University)|Mar 27, 2019

Theoretical and Computational Physics参考文献 17被引用数 8

ひとこと要約

この論文では、Cloud TPUクラスタ上でTensorFlowを用いて2次元イジング模型の高性能モンテカルロシミュレーションを提示している。TPUの行列演算および高速な相互接続機能を活用することで、先行するベンチマークと比較してマルチコア性能が250%向上した。実装は1つのJupyter Notebookに収まり、bfloat16精度を維持しながらも、2048コアまで強い線形スケーリングを示した。

ABSTRACT

Large-scale deep learning benefits from an emerging class of AI accelerators. Some of these accelerators' designs are general enough for compute-intensive applications beyond AI and Cloud TPU is one such example. In this paper, we demonstrate a novel approach using TensorFlow on Cloud TPU to simulate the two-dimensional Ising Model. TensorFlow and Cloud TPU framework enable the simple and readable code to express the complicated distributed algorithm without compromising the performance. Our code implementation fits into a small Jupyter Notebook and fully utilizes Cloud TPU's efficient matrix operation and dedicated high speed inter-chip connection. The performance is highly competitive: it outperforms the best published benchmarks to our knowledge by 60% in single-core and 250% in multi-core with good linear scaling. When compared to Tesla V100 GPU, the single-core performance maintains a ~10% gain. We also demonstrate that using low precision arithmetic---bfloat16---does not compromise the correctness of the simulation results.

研究の動機と目的

• AIアクセラレータ（Cloud TPU）をディープラーニング以外の科学計算にも応用可能かどうかを検討すること。
• TensorFlowを用いたTPU上で、パフォーマンスが高く、スケーラブルで読みやすい2次元イジング模型のモンテカルロシミュレーションを開発すること。
• 統計物理学におけるモンテカルロシミュレーションへの低精度算術（bfloat16）の影響を評価すること。
• 数千のTPUコアにわたるシミュレーションの強い線形スケーリングを示すこと。

提案手法

• Cloud TPU上でTensorFlowを用いて、2次元イジング模型のSIMD分散マルコフ連鎖モンテカルロ（MCMC）アルゴリズムを実装する。
• TPUの高帯域幅メモリおよび2次元トロイダルメッシュ相互接続を活用し、効率的なデータ移動と低遅延通信を実現する。
• 最近接スピン間のエネルギー寄与を効率的に計算するためにtf.nn.conv2dを使用し、MXUの利用を最適化するためのバッチ行列乗算を置き換える。
• リソース利用率を最大化するために、ローブラック、ドライバーパックド、スーパードライバーパックドの3つの構成でワークロードの分配を最適化する。
• TensorFlow r1.15とTPU v3ハードウェアを用いて、最小限のコード複雑性で高いパフォーマンスを達成する。
• Google Colab経由で最小限のセットアップで、Jupyter Notebook内で全シミュレーションを実行可能にし、TPU上でインタラクティブな科学計算を可能にする。

実験結果

リサーチクエスチョン

• RQ1 Cloud TPUクラスタは、ディープラーニング以外の高パフォーマンスな科学的シミュレーションに効果的に利用可能だろうか？
• RQ2 TPU上でTensorFlowを用いたイジング模型シミュレーションのパフォーマンスは、公表済みのGPUおよびCPUベンチマークと比較してどうなるか？
• RQ3 bfloat16精度を用いることで、統計物理学におけるモンテカルロシミュレーションの正しさが損なわれるだろうか？
• RQ4 シミュレーションは、何千ものTPUコアにわたってどの程度強い線形スケーリングを示すだろうか？
• RQ5 複雑な分散科学的アルゴリズムを、パフォーマンスを犠牲にすることなく、TensorFlowで明確かつ効率的に表現できるだろうか？

主な発見

• TPUベースの実装は、シングルコアでは60%、マルチコアでは250%も、最高の公表済みベンチマークを上回り、優れた線形スケーリングを示した。
• 32,768コアを備えたフル TPU v3 ポッド上では、1ナノ秒あたり20,460.92回のスピン反転が達成され、1,000コアまで強いスケーリングが実証された。
• 1,000コアを超えると通信オーバーヘッドが顕著になり、強スケーリングテストにおけるさらなるスケーリング効率が制限された。
• bfloat16精度を用いても、シミュレーション結果の正しさに影響がなく、物理的正確性が維持された。
• 全シミュレーションは1つの読みやすいJupyter Notebookに収まり、最小限のセットアップでGoogle Colabから実行可能だった。
• バッチ乗算の代わりにtf.nn.conv2dを最適化に使用したバージョンは、初期実装と比較して約80%のパフォーマンス向上を達成した。