不均衡・多忠実度の原子論的データ向けエクサスケール・マルチタスク・グラフ基盤モデル

Abstract

本稿では、HydraGNN上に構築された原子的グラフ基盤モデルを用いて、材料探索のためのエクサスケール・ワークフローを提案します。16の公開された第一原理データセット（544+百万の構造、85+元素をカバー）に対し、データセットごとのヘッドを備えたマルチタスク・アーキテクチャと、スケーラブルなADIOS2/DDStoreデータ・パイプラインを用いて共同で学習します。Frontier上では、FP64で6件の大規模DeepHyperハイパーパラメータ最適化キャンペーンを実行し、最高性能のメッセージパッシング・モデルを、持続的な2,048ノード学習へと昇格させ、PaiNNベースのリードモデルを得ます。その結果得られたモデルにより、10億スケールのスクリーニングが可能となり、50秒で11億の原子的構造を評価できます。これにより、第一原理計算だけで何年も要するであろう作業負荷が圧縮され、データが乏しい状況でも多様な下流タスクに対する微調整（fine-tuning）が可能になります。精度と性能のトレードオフ（BF16/FP32/FP64）を定量化し、化学的に多様な12の下流タスク間での転移を実証し、Frontier、Aurora、Perlmutterにわたるシームレスな強スケーリングおよび弱スケーリングを確立します。本研究は、第一原理手法では到達できない広大な化学設計空間を、迅速かつ信頼性高く探索することを可能にします。