cuNNQS-SCI: 高パフォーマンスなニューラルネットワーク量子状態を用いた選択的構成相互作用（SCI）のための完全 GPU 加速フレームワーク

本論文は、複雑な多体系のシュレーディンガー方程式を高精度に解くための AI 駆動型手法である「ニューラルネットワーク量子状態（NNQS）」と「選択的構成相互作用（SCI）」を組み合わせた手法の性能を大幅に向上させた、cuNNQS-SCI という完全 GPU 加速フレームワークを提案するものです。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

---

1. 背景と課題（Problem）

量子化学計算における電子構造計算は、多体系のシュレーディンガー方程式の数値解法に帰着されます。しかし、状態空間がシステムサイズに対して指数関数的に増大するため、完全な構成相互作用（FCI）は小規模な系にしか適用できません。

近年、ニューラルネットワーク量子状態（NNQS）を用いた変分法が注目されていますが、特に高精度かつスケーラブルな手法としてNNQS-SCIが注目されています。しかし、既存の NNQS-SCI 実装には以下の重大なボトルネックが存在し、大規模系への適用を阻害していました。

1. 中央集権的な CPU による重複排除（Global De-duplication）のボトルネック:
   • 生成された膨大な構成（configuration）から重複を除去する際、従来の手法は単一の CPU ノード（ルートノード）にすべてのデータを収集して処理していました。
   • これにより、通信帯域の飽和と CPU メモリ容量の限界が生じ、GPU の計算リソースが遊ばれる状態（GPU 利用率の低下）を引き起こします。大規模クラスターでは、この処理が全体の稼働時間の 30% 以上を占めることがありました。
2. CPU 依存の結合状態生成（Coupled-State Generation）:
   • 選択された構成から新しい候補構成を生成する処理が CPU 上で実行されており、GPU 推論の高速さに比べて著しく遅い非 AI 部分として機能していました。
   • ホスト（CPU）とデバイス（GPU）間の頻繁なデータ転送がパフォーマンスのボトルネックとなっていました。
3. GPU メモリ容量の壁:
   • 完全な GPU 加速を実現しようとしても、中間データの量が単一 GPU の高帯域幅メモリ（HBM）の容量（例：A100 の 40GB）を超えてしまい、大規模な化学系（例：Cr2 分子など）のシミュレーションが物理的に不可能でした。

2. 手法とアーキテクチャ（Methodology）

cuNNQS-SCI は、上記の課題を解決するために、ワークフロー全体を GPU 上で完結させる「完全 GPU 常駐（GPU-resident）」フレームワークを設計しました。主な技術的アプローチは以下の 3 つです。

A. 分散ソートベースのグローバル重複排除アルゴリズム

• ハッシュではなくソート: 量子構成データが密な整数配列であることを利用し、ハッシュベースではなくソートベースのアプローチを採用しました。
• 正規サンプリング（Regular Sampling）: 各 GPU がローカルデータをソートし、一定間隔でピボット（代表値）をサンプリングしてルートノードに集約します。ルートノードでグローバルな分割点（スプリッター）を決定し、これを全ノードにブロードキャストします。
• AlltoAll 通信: 各ノードは、分割点に基づいてデータを宛先ノードへ直接転送（MPI_AlltoAllv）します。これにより、重複排除が分散処理され、負荷分散が保証され、通信オーバーヘッドが最小化されます。

B. 細粒度の GPU カーネルによる結合構成生成

• CUDA カーネルの最適化: 従来の CPU での生成処理を、GPU 上で実行する専用 CUDA カーネルに移植しました。
• 圧縮されたデータレイアウト: ハミルトニアンの非ゼロ要素（スレーター・コンドン則に基づく単一・二重励起）を事前計算し、圧縮された「励起テーブル（Excitation Tables）」として GPU メモリに保持します。これにより、巨大な行列を保持することなく、必要な積分値に効率的にアクセスできます。
• 並列フィルタリング: 生成された候補の中から有効なものをフィルタリングする際、アトミック操作の競合を避けるため、共有メモリ内でのプレフィックス和（prefix sum）を用いた並列ストリーム圧縮を実装し、スループットを最大化しました。

C. GPU メモリ中心の実行モデル（GPU Memory-Centric Execution）

• ミニバッチ処理と非同期オーバーラップ: 単一 GPU のメモリ容量を超えても計算を継続するために、データをミニバッチに分割して処理します。
• 3 段階パイプライン:
  1. 生成・重複排除: 生成されたデータを非同期で CPU メモリにオフロードし、必要なデータのみを GPU に保持します。
  2. 推論・選択: 重複排除されたユニークな構成をバッチ単位で GPU にストリーミングし、NNQS による波動関数振幅の推論を行います。
  3. エネルギー計算・復元: 推論結果を元の構成空間にマッピングし、エネルギーを計算します。
• 通信隠蔽: データ転送（H2D/D2H）と計算（Kernel）を別々の CUDA ストリームで並行実行し、PCIe 帯域のボトルネックを隠蔽します。

3. 主要な貢献（Key Contributions）

1. 完全 GPU 加速 NNQS-SCI フレームワークの提案:
   • 結合構成生成、グローバル重複排除、エネルギー計算までを含むワークフロー全体を GPU 上に移行し、CPU 側のスケーラビリティボトルネックを排除しました。
2. 高性能カーネルと分散重複排除アルゴリズム:
   • 細粒度の並列化を用いた専用 CUDA カーネルと、負荷分散を保証する分散ソートベースの重複排除アルゴリズムを実装しました。
3. GPU メモリ中心の実行モデル:
   • GPU メモリを「キャッシュ」として扱い、ミニバッチ処理と非同期オフロードを組み合わせることで、物理的な GPU メモリ容量を超えた大規模シミュレーションを可能にしました。
4. 広範な実験的評価:
   • 多様な分子系（C2, N2, Cr2 など）および大規模クラスター（64 GPU）での評価を行い、既存手法との比較を通じて有効性を証明しました。

4. 実験結果（Results）

実験は、64 個の NVIDIA A100 GPU を搭載したクラスター環境で行われました。

• エンドツーエンドの高速化:
  • 最適化された NNQS-SCI ベースラインと比較して、最大 2.32 倍の高速化を達成しました（N2 分子の cc-pVDZ 基底関数系で 2.32 倍、Cr2 分子で 1.88 倍）。
  • 高速化の要因は、結合生成処理が 432 倍、重複排除処理が 40 倍それぞれ高速化したことにあります。
• スケーラビリティ:
  • 強スケーリング（Strong Scaling）: 64 GPU 使用時でも91% 以上の並列効率を維持しました。
  • 弱スケーリング（Weak Scaling）: 問題サイズを GPU 数に比例して増大させた場合、重複排除アルゴリズムの効率により、有効な計算負荷が相対的に減少するため、**163.7% という超線形加速（Super-linear speedup）**が観測されました。
• 精度の維持:
  • 化学的精度（1.6 × 10⁻³ Hartree）を達成し、既存の CPU ベースの NNQS-SCI と同等の精度を維持していることを確認しました。GPU 並列計算による浮動小数点演算の順序の違いによる誤差は極めて小さく（MAE: 3.91 × 10⁻⁴ Hartree）、問題ないレベルです。
• メモリ制約の突破:
  • 理論上、単一 GPU のメモリ容量（40GB）を超えてしまう系（例：C2H4O で理論上 70GB 必要）であっても、cuNNQS-SCI のメモリ管理により、ピークメモリ使用量を 40GB 以内に抑え、計算を成功させました。

5. 意義と結論（Significance）

cuNNQS-SCI は、AI 駆動型科学計算（AI for Science）におけるシステムレベルの課題を解決する重要なステップです。

• ボトルネックの転換: 従来の「CPU 側の通信・管理コスト」がボトルネックだった状態から、「デバイス上の推論コスト」へとボトルネックをシフトさせることに成功しました。これにより、現代の GPU クラスターが持つ膨大な並列計算能力を最大限に活用できるようになりました。
• 大規模系への適用可能性: 単一 GPU のメモリ容量の壁を破ることで、以前は計算不可能だった大規模な化学系（強相関電子系など）の高精度シミュレーションが可能になりました。
• 将来への展望: このフレームワークは、量子化学計算の限界を押し広げ、より複雑で科学的に重要な大規模系への応用を可能にする基盤技術となります。

本論文は、HPC（高性能計算）と AI の融合において、アルゴリズムの最適化だけでなく、システムアーキテクチャ（メモリ管理、通信パターン、並列化戦略）の再設計が不可欠であることを示す重要な事例となっています。