Mat3raが NVIDIA ALCHEMI に対応：材料シミュレーションを40〜60倍高速化し、導入も容易に

Takuya Shibayama と Mitsunori Nakamoto（NVIDIA）によるゲスト投稿。編集：Erica Tsai（NVIDIA）。

（Mat3ra ブログ上で、Timur Bazhirov のアカウントから公開予定）·

AI は今や、誰にもテストしきれない速さで新しい材料を提案できます。モデルが提案した候補は、作る価値があるかどうかが分かる前に、安定構造へと緩和し、加熱し、歪みを加え、欠陥の有無を確認しなければなりません。このテストこそが、AI 駆動型材料探索のボトルネックです。Mat3ra は、この課題を両面から解消することを目指しています。Mat3ra プラットフォームがシミュレーションの実行を簡単にし、NVIDIA ALCHEMI がシミュレーションそのものを高速化します。

この取り組みの中核にあるのが NVIDIA ALCHEMI Toolkit 0.3.0 です。1 枚の GPU で、数百の構造を 1 つずつではなく同時にシミュレーションできます。Mat3ra プラットフォームでの最初のベンチマークでは、NVIDIA H100 NVL GPU 1 基で、従来の 1 構造ずつ処理する方式の 40〜60 倍の速さで結晶構造を緩和しました。4,000 構造を約 3 分で処理できたのに対し、従来の方法では 3 時間近くかかります。さらに Mat3ra プラットフォームが、使いやすさとアクセスのしやすさを高めます。Mat3ra プラットフォームのユーザーは、ソフトウェアのインストールも GPU の購入も不要で、同じワークフローをノートブックから実行できます。NVIDIA ALCHEMI チームが最適化したセットアップの上に、ユーザー自身の構造、組成、ユースケースを組み込むことができます。

1 つの GPU で、数百の構造を同時に

データセンター向け GPU は、大規模な並列処理のために作られています。小さな結晶 1 個だけでは能力を使い切れないため、広く使われている Atomic Simulation Environment（ASE）[4] のように 1 度に 1 構造ずつ扱う従来のツールでは、GPU の大半が遊んだままになります。

NVIDIA のオープンソースの AI 原子シミュレーション向けフレームワークである ALCHEMI Toolkit [1] を使うと、数百の構造を 1 回の処理でまとめて評価し、一緒に更新できるため、GPU をフルに活用できます。内部では、NVIDIA cuEquivariance [2] が機械学習原子間ポテンシャル MACE [3] を高速化します。その結果、1 基の GPU から大幅に高いスループットが得られます。

緩和は出発点にすぎません。同じバッチ処理エンジンで、温度・圧力を設定した分子動力学も実行できます。バージョン 0.3.0 では、より高速な最適化手法、反応経路・拡散経路の探索、拡張サンプリング、そして生成 AI モデルからシミュレーションへ直接つなぐパイプラインが加わりました。AI モデルがどんな候補を提案しても、ALCHEMI でまとめてテストできます。

数値で見る結果

[図 1：mat3ra-alchemi-figure-1-throughput.webp をここに挿入]

図 1. Mat3ra プラットフォーム上の NVIDIA H100 NVL GPU 1 基における、バッチサイズごとの構造緩和スループット。ALCHEMI Toolkit 0.3.0 と cuEquivariance で高速化した MACE を使用。塗りつぶしの四角：最適化フェーズ。白抜きの四角：エンドツーエンド。ひし形：ASE による 1 構造ずつの処理。

[図 2：mat3ra-alchemi-figure-2-speedup.webp をここに挿入]

図 2. NVIDIA H100 NVL GPU 1 基、バッチサイズ 400 における、1 構造ずつ処理する ASE に対する ALCHEMI の高速化倍率。左：最適化フェーズのみ。右：エンドツーエンド。

スループットは、数百構造で GPU が埋まるまで、バッチサイズとともに上昇します。バッチサイズ 400 では、ALCHEMI のスループットは最適化フェーズで毎秒 44 構造、エンドツーエンドで毎秒 21 構造でした。1 構造ずつ処理した場合は、それぞれ毎秒 1.1 構造と 0.34 構造です。これは、およそ 40 倍と 60 倍のエンドツーエンドの高速化に相当します。バッチ処理した結果と 1 構造ずつ処理した結果は一致しており、速度のために精度を犠牲にすることはありません。

表 1. ベンチマークの概要

システム
歪みを加え、原子をランダムに変位させた 64 原子のシリコンスーパーセル。0.03 eV/Å まで緩和

ハードウェアとモデル
Mat3ra プラットフォーム上の NVIDIA H100 NVL 1 基。cuEquivariance を使用した MACE-MPA-0 medium

比較
バッチサイズ 10〜800 の ALCHEMI Toolkit 0.3.0 と、1 構造ずつ処理する ASE の比較。エンドツーエンドのスループットは、バッチサイズごとに総構造数を変えて測定

ノートブックから研究キャンペーンまで、1 日で

Mat3ra プラットフォームでは、ベンチマーク全体が 1 つの Jupyter ノートブックになっています。プラットフォームの材料ライブラリから構造を作成してセットとして保存し、手法ごとに 1 つのジョブを GPU キューに投入します。プラットフォームが構造をワークフローに渡し、ノード上にツールキットと GPU 依存ライブラリをインストールして、結果を表として返します。ターミナルは一切不要です。ノートブックは、プラットフォーム組み込みの JupyterLite を通じてブラウザ上で、または Mat3ra API を通じてローカルで実行できます。

研究キャンペーンも同じ流れで進められます。ご自身の構造や、動力学・反応経路のスクリプトに差し替えるだけで、ワークフローの残りはそのままです。構造が 50 個でも 5 万個でも、ノートブックは変わらず、変わるのはセットだけです。

共同で作り、チームのために作る

どちらか片方だけでは成り立ちません。40 倍の高速化も、それを設定できるラボにしか役立ちませんし、使いやすいプラットフォームも、そこで動くものが速くなければ意味がありません。今回の取り組みでは、ツールキットとベンチマークに関する NVIDIA のエンジニアリングと、Mat3ra のプラットフォームおよびワークフローのエンジニアリングを組み合わせ、Mat3ra 自身の H100 NVL ノード上で実行することで、ユーザーが両方を同時に得られるようにしました。Mat3ra プラットフォームは、AI 駆動型材料 R&D のための物理に基づくインフラを 4 万人を超える登録ユーザーに提供しており、こうした協働を 1 か所にまとめます。構造、ワークフロー、ジョブ、結果はチームの組織アカウントに属するため、全員が同じデータを使い、同じ結果を見ることができます。ノートブックは Mat3ra の API サンプルリポジトリ [5] で公開予定で、すぐに再実行したり、応用したりできます。

「構造緩和、温度依存性、欠陥や不純物の振る舞い、そして実際の材料や製造で生じるさまざまな効果など、原子レベルのさまざまな経路を探索するには、膨大な位相空間を探索する必要があります」と、Mat3ra の CEO、Timur Bazhirov 氏は述べています。「Mat3ra プラットフォーム上の ALCHEMI なら、数千通りの探索をバッチモードで数分のうちに実行できます。ワークフローはモジュール式で、ユーザーは NVIDIA ALCHEMI チームによる最適なセットアップを土台に、自分の構造、組成、ユースケースを組み込めます。」

「Mat3ra プラットフォームは直感的に使え、ALCHEMI を素早く導入できる柔軟性があります」と、NVIDIA のプロダクトマネジメント担当シニアディレクター、Geetika Gupta 氏は述べています。「環境構築なしで、GPU の性能をエンドユーザーの手の届くところに届けます。ALCHEMI は非常に活発に開発が進められており、Mat3ra では各リリース直後に最新機能を試すことができます。AI モデルがこれまで以上に多くの候補を提案するなか、研究者はそれらをまとめてテストし、より早く検証済みの材料にたどり着けます。」

始めるには

ALCHEMI Toolkit は GitHub [1] で公開されており、ドキュメントは https://nvidia.github.io/nvalchemi-toolkit/ にあります。https://platform.mat3ra.com でサインインまたは無料アカウントを作成し、Mat3ra の API サンプルリポジトリ [5] でノートブックを見つけてください。

参考文献

[1] NVIDIA ALCHEMI Toolkit. https://github.com/NVIDIA/nvalchemi-toolkit

[2] NVIDIA cuEquivariance. https://github.com/NVIDIA/cuEquivariance

[3] I. Batatia, D. P. Kovács, G. N. C. Simm, C. Ortner and G. Csányi, “MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields,” Advances in Neural Information Processing Systems (2022). https://arxiv.org/abs/2206.07697.

[4] A. Hjorth Larsen et al., “The atomic simulation environment—a Python library for working with atoms,” J. Phys.: Condens. Matter 29, 273002 (2017). https://doi.org/10.1088/1361-648X/aa680e

[5] Mat3ra API examples. https://github.com/mat3ra/api-examples

注：構造緩和の測定は 1 つの構成で、ケースごとに 1 回の実行です。スループットは、演算、構造、モデル、設定、ハードウェアによって変わります。商標はそれぞれの所有者に帰属します。

キーワード：#NVIDIA #ALCHEMI #cuEquivariance #MACE #MLIP #BatchedSimulation #GPU #MaterialsInformatics #Mat3ra