[論文レビュー] Using Machine Learning at Scale in HPC Simulations with SmartSim: An Application to Ocean Climate Modeling

Sam Partee, Matthew J. Ellis+5 more|arXiv (Cornell University)|Apr 13, 2021

Oceanographic and Atmospheric Processes参考文献 31被引用数 13

ひとこと要約

本論文では、スケーラブルでオンラインの機械学習推論をハイパフォーマンスコンピューティング（HPC）シミュレーション内に可能にするソフトウェアフレームワーク、SmartSimを提示する。PyTorch、TensorFlow、ONNXモデルをFortranベースの海洋気候モデル（例：MOM6）に統合し、12メンバーのグローバル海洋シミュレーションを120年間にわたり9700億回の推論を実行するなど、安定したパフォーマンスを達成し、最小限のランタイムオーバーヘッドを実現した。

ABSTRACT

We demonstrate the first climate-scale, numerical ocean simulations improved through distributed, online inference of Deep Neural Networks (DNN) using SmartSim. SmartSim is a library dedicated to enabling online analysis and Machine Learning (ML) for traditional HPC simulations. In this paper, we detail the SmartSim architecture and provide benchmarks including online inference with a shared ML model on heterogeneous HPC systems. We demonstrate the capability of SmartSim by using it to run a 12-member ensemble of global-scale, high-resolution ocean simulations, each spanning 19 compute nodes, all communicating with the same ML architecture at each simulation timestep. In total, 970 billion inferences are collectively served by running the ensemble for a total of 120 simulated years. Finally, we show our solution is stable over the full duration of the model integrations, and that the inclusion of machine learning has minimal impact on the simulation runtimes.

研究の動機と目的

• ファイルI/Oのボトル neck を引き起こさずに、Fortran/C/C++で書かれた従来のHPCシミュレーションと機械学習を統合する課題に対処する。
• 長時間にわたる気候シミュレーション中に、リアルタイムで分散型の深層ニューラルネットワーク推論を可能にする。
• HPCワークロードと現代のMLライブラリとの間の言語的・データ交換の障壁を克服する。
• 気候スケールの解像度におけるML増強型HPCシミュレーションのスケーラビリティと安定性を示す。
• 異種のHPCシステムと複数のMLフレームワークをサポートする、柔軟でライブラリに依存しないインターフェースを提供する。

提案手法

• SmartSimのインfraストラクチャおよびクライアントライブラリを活用し、ML推論をシミュレーション論理から分離する。
• SmartRedis APIを用いて、Fortranシミュレーションから低遅延でリモート実行可能なMLモデルを実現する。
• JITトレースされたPyTorchモデルをMOM6に統合し、従来の渦運動エネルギーのパrameterizationを置き換える。
• 12のアンサンブルメンバーすべてに共通のMLモデルをデプロイし、各メンバーが19ノードで実行され、各タイムステップで同期された推論を実現する。
• Redisベースのデータ交換を活用し、ファイルI/Oのオーバーヘッドを排除し、シミュレーションとMLコンponent間のリアルタイムデータフローを可能にする。
• CPUおよびGPUワークロードをサポートする異種のHPCシステムに対応し、多様なハードウェアおよびソフトウェアスタック間での相互運用性を確保する。

実験結果

リサーチクエスチョン

• RQ1 大規模でプロダクショングレードのHPC気候シミュレーションに機械学習を統合しても、シミュレーションの安定性やパフォーマンスに影響を及ぼさずに実現可能か？
• RQ2 スケール上で、1つのMLモデルが複数のHPCシミュレーションインスタンスに分散してリアルタイム推論を効率的に行えるか？
• RQ3 共有モデルアーキテクチャを用いてグローバル海洋モデルにオンラインML推論を統合した際のパフォーマンスオーバーヘッドはどの程度か？
• RQ4 SmartSimフレームワークは、FortranベースのシミュレーションとPyTorchやTensorFlowなどの現代のMLライブラリとの間で、シームレスで言語に依存しない通信を可能にするか？
• RQ5 SmartSimを介したML推論の統合は、長期間にわたる気候シミュレーションにおいても数値的安定性を保持するか？

主な発見

• SmartSim-EKEアンサンブルは、10年ずつ合計12回のグローバルかつ高解像度の海洋シミュレーションを実行し、数値的不安定性は観測されなかった。
• 120年間のシミュレーション期間中に、合計9700億回の推論操作が実行され、エクサスケールに近いスケーラビリティを実証した。
• SmartSimを介したML推論の統合により、最小限のパフォーマンスオーバーヘッドが生じ、シミュレーション実行時間はほとんど影響を受けなかった。
• フレームワークは、異種のHPCシステム上で数千プロセッサにわたる線形スケーリングを達成し、CPUおよびGPUワークロードを両方サポートした。
• すべてのアンサンブルメンバーに共通のMLモデルを用いることで、各シミュレーションでのモデルの重複を回避し、一貫性があり同期された推論が可能になった。
• ソリューションはプロダクション環境で利用可能であり、完全に再現可能で、すべてのソースコード、データセット、およびトレーニング済みモデルがGitHubで公開されている。