[論文レビュー] FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver

Chirag Modi, François Lanusse+1 more|arXiv (Cornell University)|Oct 22, 2020

Galaxies: Formation, Evolution, Phenomena参考文献 53被引用数 46

ひとこと要約

FlowPMは、Mesh-TensorFlow上の新規マルチグリッド・マルチレゾリューション金字塔スキームを用いたTensorFlowでのGPU加速・微分可能・分散型Particle-Mesh N体解算機を実装し、Python FastPMに対して約10倍の実行時間短縮を実現するとともに、宇宙論推論の微分可能なフォワードモデリングを可能にする。

ABSTRACT

We present FlowPM, a Particle-Mesh (PM) cosmological N-body code implemented in Mesh-TensorFlow for GPU-accelerated, distributed, and differentiable simulations. We implement and validate the accuracy of a novel multi-grid scheme based on multiresolution pyramids to compute large scale forces efficiently on distributed platforms. We explore the scaling of the simulation on large-scale supercomputers and compare it with corresponding python based PM code, finding on an average 10x speed-up in terms of wallclock time. We also demonstrate how this novel tool can be used for efficiently solving large scale cosmological inference problems, in particular reconstruction of cosmological fields in a forward model Bayesian framework with hybrid PM and neural network forward model. We provide skeleton code for these examples and the entire code is publicly available at https://github.com/modichirag/flowpm.

研究の動機と目的

• 分散GPU上で、微分可能なPMフレームワークを用いて高速でスケーラブルな宇宙論N体シミュレーションを実現する。
• 新規のマルチレゾリューションピラミッドに基づくマルチグリッド力推定を開発・検証し、プロセス間通信を削減する。
• 分散型のモデルパラレルN体計算のためのMesh-TensorFlowベースのアーキテクチャを提供する。
• ハイブリッドPM/MLフォワードモデルによるフォワードモデリングおよび宇宙論場再構成への適用性を示す。

提案手法

• GPU加速と微分可能性のために、TensorFlowでFastPMスタイルのPM進化を実装する。
• 長距離力と近距離力を粗いグローバル格子と細いローカル格子に分割するため、マルチレゾリューションピラミッドを用いた二レベルのマルチグリッド力推定を導入する。
• Mesh-TensorFlow内で分散グリッドチャンクに対してスライス単位の演算を可能にするため、パディングを伴うハロー交換プロトコルを用いる。
• グリッド次元の分配を制御したプロセッサメッシュへテンソル計算をマップするため、Mesh-TensorFlowを採用する。
• bsplineカーネルを用いた3D畳み込みによる滑らかなダウンサンプリング/アップサンプリング（REDUCE/EXPAND）操作を構築・評価してピラミッド構築を行う。
• 標準のFastPMコードと精度を比較し、 grid/mesh レイアウトを変えたCori GPUs上のスケーリングを評価する。

実験結果

リサーチクエスチョン

• RQ1 微分可能でGPU加速されたフレームワークで、PM宇宙論N体解法をどのように実装できるか？
• RQ2 マルチグリッド・マルチレゾリューションピラミッド方式は、精度を保ちながら分散FFT通信を削減できるか？
• RQ3 FlowPMを大規模GPUメッシュ上で実行した場合の性能とスケーリング特性は、Python FastPM実装と比較してどうか？
• RQ4 FlowPMは前方モデル推定ワークフロー（例: ハイブリッド PM/ML コンポーネントによる再構成）に効果的に統合できるか？

主な発見

• FlowPMは、テスト設定に対してFastPMと比較してサブパーセントレベルの精度を達成する（トランスファ関数と交差相関は0.01%の範囲内）。
• 二レベルのマルチグリッドピラミッド方式は、長距離力を粗いグローバル格子で、短距離力を局所格子で処理することで分散FFT通信を削減する。
• FlowPMは8-V100ノードのCori GPUs上でスケールし、同等の設定でPython FastPMと比べて平均で約10xのwallclockスピードアップを示す。
• Mesh-TensorFlowベースの実装は、プロセススライス間の境界領域をハロー交換で管理する分散モデルパラレルシミュレーションを可能にする。
• このフレームワークは、ニューラルネットワーク前方モデルを統合したおもちゃの例を通じて示されたように、フォワードモデリングと再構成タスクに適した微分可能なシミュレーションをサポートする。