Yitian 710 instance helps ABACUS new practice

01 AI for Science が切り拓く材料研究開発の新パラダイム

はじめに、AI for Science が材料研究開発プロセスにおいてどのような新パラダイムをもたらしているかについてお話しします。材料イノベーションは、創薬設計や新エネルギーなどの分野の発展を牽引する原動力です。材料研究開発のパラダイムは、従来の大量な反復実験に基づく試行錯誤から、理論シミュレーションで候補材料をスクリーニングしてから研究開発の検証を行う計算駆動型プロセスへと進化してきました。

しかし、理論シミュレーションのプロセスでは、次元の呪いが高精度計算の進展を妨げています。ポール・ディラックの言葉を借りれば、「量子力学により、すべての化学問題およびほとんどの物理問題について、物理学の基本的法則は既に解明されている。困難なのは、これらの法則を正確に適用すると数学方程式が複雑になりすぎて解けなくなることだけである」と述べています。

次元の呪いの問題を解決するため、科学者は最も基本的で正確な物理モデルを段階的に抽象化し、異なる系に対応する物理モデルを選択することで、計算結果を合理的な時間内に導き出せるようにしています。ただし、異なるスケールの物理モデルは時間と空間において数桁の差があり、結果の精度にも小さくない差が生じます。

AI for Science のアプローチは上記の問題を解決することを目的としています。高精度手法で学習した特徴を機械学習を通じてより大きな系に適用でき、高精度な計算結果と効率的な求解時間を両立できます。

原子スケールの分子シミュレーション手法には依然として課題が残っています。従来の分子動力学法では、科学者が力場の経験的パラメータを与える必要があり、ポテンシャル関数の開発サイクルは非常に長いものでした。

密度汎関数理論（DFT）については、DFT ソフトウェアのコード分岐が非常に複雑で研究開発サイクルが長く、交換相関汎関数近似を使用していますが、近似の精度を高めるほど計算負荷が大きくなります。

ディープポテンシャル法は機械学習に基づく分子動力学法であり、科学計算、機械学習、高性能コンピューティングを良好に融合させた手法です。

左図はディープポテンシャル法のトレーニングプロセスを示しています。DFT で原子ポテンシャルを計算し、ニューラルネットワークでポテンシャル関数を学習させ、最後に分子動力学に応用することで、効率的かつ高精度な計算を実現します。

上図中央の結果は、ディープポテンシャル法と従来の DFT 法で計算した分子動径分布関数の比較を示しており、両者の結果はよく一致しています。

DFT はシュレーディンガー方程式を解くことで物質の性質を直接計算するアルゴリズムです。

結晶胞パラメータなどの情報を与えると、DFT 計算により導電率やシステム密度などの基本的な物理特性を得ることができます。これは経験パラメータをほとんど必要としない第一原理計算法です。この手法により、研究者は 1998 年にノーベル化学賞を受賞しました。

上記の数式に示すように、密度汎関数理論の核心は、系の総エネルギー E を電子密度 ρ の汎関数として表現することにあります。

左下は DFT のヤコブの梯子を示しており、最も簡単な局域密度近似から始まり、計算量を代償として正確な量子力学手法による計算結果に段階的に近づいていきます。DeePKS 法はニューラルネットワークモデルを用いて高精度手法と低精度手法の差を表現します。研究者は DeePKS ソフトウェアを使用し、まず低精度汎関数で比較的高効率に結果を導出し、そこに DeePKS 法による補正値を加えることで、計算結果を高精度汎関数に近づけることができます。右図は水分子中の酸素元素間距離の動径分布関数です。DeePKS 法と高精度 DFT 法で計算した結果はよく一致しています。対照的に、PBE 汎関数を用いた計算結果は大きな偏差を生じます。

AI + 材料科学パラダイムでは、AI モデル トレーニング用のデータを DFT ソフトウェアで生成する必要があります。トレーニング中のモデルは DFT ソフトウェアの計算結果にも影響を与え、このプロセスは収束するまで反復する必要があり、その間に大量の DFT 計算が求められます。

ABACUS は中国発のオープンソース密度汎関数理論ソフトウェアです。ABACUS は当初、中国科学技術大学の何立新教授の研究グループによって開発され、2021 年 3 月に DeepModeling オープンソースコミュニティに参加しました。

科学計算ソフトウェアとして、ABACUS の研究開発は従来の研究グループによるコード開発の枠組みを打破し、GitHub プラットフォーム上でコードをホストしています。オープンソースコントリビューターが新機能の開発やエラーの修正に共同で取り組むことを歓迎しています。

DeepModeling コミュニティ参加後、ABACUS は中国科学技術大学、北京大学、物理研究所、北京智能科学研究院（AISI）など、さまざまな機関からのコントリビューターを迎えました。

特筆すべきは、北京智能科学研究院（AISI）は 2021 年にエビアン院士によって設立された、AI for Science をミッションとする最初の科学研究機関であることです。

02 Yitian 710 への適合と最適化

さて、ここでは ABACUS を Yitian 710 クラウドプラットフォームへ移行した方法について紹介します。Yitian 710 のオペレーティングシステムは Alibaba Cloud が提供する Ali Linux3 で、既存の Linux ソフトウェアエコシステムを十分にサポートできます。ソフトウェアのインストール時に、ユーザーはパッケージマネージャから必要な依存関係を直接ダウンロードでき、手動で再コンパイルする必要はありません。

Yitian 710 チップは ARM v9 アーキテクチャをベースに開発され、SVE、INT8mm、BF16mm などの命令セットアクセラレーションをサポートしています。さらに、ARM プラットフォームは高性能な数学ライブラリを提供しており、研究者は行列計算の実装を気にすることなくアルゴリズム開発に集中できます。

上図は ABACUS 付属の計算例による性能比較を示しています。右端の列が Yitian 710 の計算時間で、第 7 世代 x86 アーキテクチャ CPU および第 6 世代高周波数インスタンスと同等の結果を示しています。

なお、上記テストで使用した Yitian 710 インスタンスは 4xlarge 仕様であり、これは x86 インスタンスの半分の規模です。この結果は、Yitian 710 が独立した物理コア、独立したキャッシュ、優れた ALU 性能を備え、過度なスレッド損失が生じないことに起因しています。

03 E-HPC を活用した 1 万核規模の検証

研究チームは Alibaba Cloud E-HPC 上で、実際の環境に近い計算能力の検証を実施しました。Alibaba Cloud が提供する Elastic High Performance Computing (E-HPC) サービスは、ECS 仮想マシンイメージからコンピューティングノードを作成することで、コンピューティングクラスターリソースのエラスティックスケーリングを実現し、ローカルスーパーコンピューターと同等の操作性を保ちながら効率的な計算を確保します。

E-HPC で ECS インスタンスを作成する際、入札インスタンスを選択できます。これにより、時間要件に敏感でない科学計算タスクを、クラウドプラットフォームリソースの利用率が低い時間帯に非常に低価格で実行できます。

ABACUS は Yitian 710 インスタンスのクラスター上で、350 eV の超高温条件下において 32 個のホウ素原子系に対して新開発の確率的 DFT 法による計算を実行しました。

上図に示すように、研究者は Yitian 710 の 11,008 個の CPU コア、すなわち 86 台の 128 コアノードを使用しました。これはウィークスケーリングタスクであり、各コアの計算量は一定で、計算リソースの消費量はコア数の増加に伴い線形に増加します。

ABACUS が実装する軌道並列と K 点並列の 2 種類のタスク分割モードにおいて、計算時間の増加は限定的であり、並列性能は非常に良好です。ソフトウェアの精度面でも、異なるコア数で計算したエネルギーと圧力は一致しており、計算結果の正確性が確認されました。

各データポイントは 10 個のランダムシードで初期化して計算し、系統的なランダム誤差を回避しています。計算量の増加に伴い、圧力とエネルギーの標準偏差は理論的最適値に向かって収束し続けています。実際の計算では、研究者はダウンストリームタスクが必要とする精度に応じて対応するコア数を選択できます。

この実験により、Yitian ECS インスタンスの以下の利点が検証されました。

まず、Yitian 710 は安定したメインクロック周波数を備え、高密度計算の科学計算シナリオでも周波数の低下を防ぎ、性能出力の一貫性を維持できます。

次に、Yitian 710 インスタンスのスケーラビリティは優れており、1 万コア規模で線形に近い加速を実現できます。

最後に、x86 と比較して Yitian インスタンスは非常にコストパフォーマンスが高く、研究機関は約 70% のコストを削減できます。

Yitian は、従来の科学計算タスクをローカルスーパーコンピューティングからクラウドプラットフォームへデプロイするプロセスにおいて、コスト削減と効率向上を実現します。