[論文レビュー] Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework

Xiaojie Wu, Qiming Sun+12 more|arXiv (Cornell University)|Apr 15, 2024

Machine Learning in Materials Science被引用数 9

ひとこと要約

GPU4PySCF は PySCF を基盤にした GPU 加速Pythonパッケージで、DFT、SCF、勾配、ヘシアン計算で大幅な速度向上を実現し、コスト削減とエコシステム統合が大きい。暗黙的溶媒モデルとさまざまな量子化学ワークフローをサポートし、Q-Chem と検証済み。

ABSTRACT

We describe our contribution as industrial stakeholders to the existing open-source GPU4PySCF project (https: //github.com/pyscf/gpu4pyscf), a GPU-accelerated Python quantum chemistry package. We have integrated GPU acceleration into other PySCF functionality including Density Functional Theory (DFT), geometry optimization, frequency analysis, solvent models, and density fitting technique. Through these contributions, GPU4PySCF v1.0 can now be regarded as a fully functional and industrially relevant platform which we demonstrate in this work through a range of tests. When performing DFT calculations on modern GPU platforms, GPU4PySCF delivers 30 times speedup over a 32-core CPU node, resulting in approximately 90% cost savings for most DFT tasks. The performance advantages and productivity improvements have been found in multiple industrial applications, such as generating potential energy surfaces, analyzing molecular properties, calculating solvation free energy, identifying chemical reactions in lithium-ion batteries, and accelerating neural-network methods. With the improved design that makes it easy to integrate with the Python and PySCF ecosystem, GPU4PySCF is natural choice that we can now recommend for many industrial quantum chemistry applications.

研究の動機と目的

• 量子化学計算におけるデータ駆動の研究需要を満たすための GPU 加速の必要性を動機づける。
• PySCF を基盤とした Python ベースの GPU 最適化フレームワークを開発し、DFT/SCF ワークフローの性能を向上させる。
• オープンソースのエコシステムとの統合と、PES 探索、溶媒和エネルギー、ニューラルネットワークワークフローなどの産業タスクへの適用性を実証する。

提案手法

• GPU の効率を活用するために density fitting アルゴリズムを優先し、直接的な SCF のサポートを限定的に提供。
• テンソルコアを含む現代の GPU アーキテクチャへコードベースを適応し、テンソル縮約を高速化する。
• DFT 法、基底関数、溶媒モデル、性質計算のために、オープンソースの量子化学パッケージの広範なセットを GPU4PySCF に統合する。
• 結果の信頼性と正確性を確保するために Q-Chem 6.1 とのクロスバリデーションを実施する。

実験結果

リサーチクエスチョン

• RQ1 最新の GPU と比較して、32 コア CPU と比べて SCF、勾配、ヘッセ計算で GPU4PySCF によって達成される性能向上（スピードアップ）はどの程度か？
• RQ2 PCM/SMD のような implicit solvent モデルは、GPU 加速 DFT 計算の性能と精度にどのように影響するか？
• RQ3 GPU4PySCF はねじれ回転スキャン、二量体相互作用エネルギー、溶媒和エネルギー計算のような大規模タスクを実用的な時間内に可能にするか？

主な発見

分子	C-PCM_SCF	C-PCM_勾配	C-PCM_ヘシアン	IEF-PCM_SCF	IEF-PCM_勾配	IEF-PCM_ヘシアン
Vitamin C	41.1	28.4	106.4	41.6	27.3	104.9
Inosine	65.5	38.9	147.4	65.6	40.2	146.2
Bisphenol A	76.2	43.3	162.7	76.7	45.4	161.3
Mg Porphin	84.0	49.0	162.2	87.3	49.8	160.1
Penicillin V	83.2	44.1	170.5	84.0	43.4	169.4

• GPU4PySCF は、Q-Chem を 32 CPU コアで比較した場合、SCF で約 20x、ヘッセアン計算で約 50x、勾配ではより控えめな利得を示す。
• NVIDIA A100-80G および density fitting を用いると、溶媒モデルで SCF/勾配/ヘッセアンのスピードアップがそれぞれ 40-80x、20-40x、100-170x となり、中規模分子の大規模な振動分析を1日で可能にする。
• A100GPU での CPU ベース実行と比較して多くのタスクで約 90% のコスト削減が可能で、純粋な DFT の勾配タスクでは約 70% の節約を示す。
• CHELPG 電荷の速度改善は GPU4PySCF が CPU ベースのソフトウェアより約 10 倍速い。
• NMR シールド定数とその他の性質は、GIAO ベースの GPU 加速で実装され、効率的な計算を実現。
• 暗黙溶媒モデル（PCM/SMD）を用いた溶媒化自由エネルギーは、中性種で一般的に参照法に対して誤差が 1 kcal/mol 未満、イオンで約 4 kcal/mol 程度で、より大きな基底関数と高次レベルの汎関数を用いると精度が高くなる。