分子動力学シミュレーション「GROMACS」を最大2倍高速化――NVIDIA、NVSHMEMを用いたGPU主導の通信最適化手法を発表
要点
-
分子動力学シミュレーションの「GROMACS」において、従来のCPUを介したMPI通信が並列計算スケーリングのボトルネックとなっていた。
-
NVIDIAの「NVSHMEM」ライブラリを活用し、CPUを介さずにGPU自身が直接遠隔メモリへアクセスするGPU主導の通信システムを実装した。
-
依存関係を考慮したカーネル融合などの最適化により、通信と計算の重ね合わせを最大化し、1ステップあたりのカーネル起動を6回から1回に削減した。
-
スーパーコンピューター「Eos」や「NVIDIA GB200 NVL72」での検証により、従来のGPU-aware MPIと比べて強スケーリング性能が最大2倍向上することを示した。
-
NVIDIAは2026年7月9日、自社の技術ブログにおいて、分子動力学(MD)シミュレーションパッケージである「GROMACS」の並列スケーリング性能を大幅に向上させる、GPU主導の新しい通信手法を公表した。この手法は、GPU用の通信ライブラリ「NVIDIA NVSHMEM」を用いて、従来のCPUが通信を制御する方式に起因するボトルネックを解消するものだ。検証の結果、最新のGPUクラスターにおいて、従来手法と比較して最大2倍の強スケーリング性能向上が確認されたという。
分子動力学シミュレーションと通信のボトルネック
分子動力学(MD)シミュレーション(原子や分子の物理的な動きをコンピューター上で追跡して観察する技術)は、タンパク質の折り畳み現象の解明から創薬、新材料の開発に至るまで、計算科学の分野で非常に重視されている。これらは数十万から数千万個の原子挙動をフェムト秒(1000兆分の1秒)単位で追跡し、数十億ステップにわたって計算するため、極めて膨大な計算リソースを必要とする。
通常、こうしたシミュレーションでは、問題のサイズ(計算対象の原子数など)を固定したまま、利用可能なGPUなどのリソース全体に処理を分散させて並列計算する「強スケーリング(計算規模を固定したまま演算装置を増やして高速化を図る並列計算手法)」が適用される。
代表的なMDソフトウェアであるGROMACSは、CPUとGPUを組み合わせたヘテロジニアス(混在型)並列化により、複数のGPUにわたって1タイムステップあたり100〜200マイクロ秒(1万分の1〜2秒)という極めて高速な領域に達している。しかし、このように並列規模の拡大を進めるなかで、GPU同士の通信速度がボトルネック(処理全体の速度を低下させる要因)となっていた。
多くの高性能計算(HPC)アプリと同様に、GROMACSも従来はプロセス間の通信に「MPI(異なる計算ノード間でデータをやり取りするための標準的な通信規格)」を使用してきた。しかし、MPIは本質的にCPU中心の実行設計となっている。そのため、GPU上でGROMACSを実行する際、データのやり取りを行う「ハローエクスチェンジ(並列計算において隣り合う計算領域の境界データを互いに共有するアルゴリズム)」が発生すると、CPUが通信を調整し終わるまでGPUが処理を一時停止しなければならない。このCPU-GPU間のハンドオフ(処理の受け渡し)は、3次元空間の分割において1タイムステップあたり最大12回ものブロッキング(待機が発生する)同期を引き起こす。計算自体が高速化するにつれてこの同期のオーバーヘッドが重くのしかかり、ピーク動作時にはCPUによる通信制御だけで全計算時間の50%以上を消費してしまっていた。
NVSHMEMによるGPU主導の直接通信
この課題を克服するため、NVIDIAの開発チームは、CPUを通信のクリティカルパス(処理の遅延に直結する経路)から排除するアプローチを開発した。ここで用いられたのが、OpenSHMEMの「分割グローバルアドレス空間(PGAS、物理的に分散したメモリをプログラム上は単一の共有メモリのように扱える仕組み)」モデルに基づく通信ライブラリ「NVSHMEM」である。NVSHMEMを導入することで、GPU内のプログラム(カーネル)が直接、他のGPUのメモリへアクセスできるようになり、CPUを介した余分な待機が解消された。
GROMACSのハローエクスチェンジは、空間の各次元(Z方向→Y方向→X方向)の順にデータを中継する「パルス」と呼ばれるステップで構成されている。従来は、各パルスの開始時に原子データを送信用の連続バッファにまとめる「パック」処理をカーネルとして起動し、シリアル(直列)に実行していた。
今回の改良では、こうした段階的な依存関係を、1つの「融合カーネル(複数の計算や処理を統合して効率的に実行するプログラム)」の中で細分化された信号(シグナル)を用いて制御する設計に変更した。
具体的には、原子のインデックス(識別情報)マップを「他のデータに依存しない部分」と「依存する部分」に分割し、大半のデータについてはパルスの依存関係を待たずに即座にバッファに詰めて転送できるようにした。依存関係がある残りの部分についても、パルス単位で細かくシグナルを送信する仕組みをとっている。これにより、驚くべきことに1タイムステップあたりのカーネル起動回数は、従来の6回からわずか1回へと削減された。
相互接続に応じた転送の最適化と性能実証
さらに本手法では、接続されたGPUの構成(インターコネクト)に応じて通信方式を自動で切り替える賢いトランスポート(輸送)ロジックが組み込まれている。高速接続規格の「NVLink」で接続された隣接GPUに対しては、ポインタを直接操作して書き込みを行い、データの整合性を保証する「システムスコープのリリーストア(データ書き込み完了を全体に通知する機能)」を適用する。
一方、NVLinkで接続されていない離れたGPUに対しては、InfiniBandなどのネットワークを介して直接メモリアクセスを行う「RDMA(CPUを経由せずにネットワーク越しで別マシンのメモリへアクセスする技術)」上のNVSHMEM put操作へと自動的に切り替える。さらに、NVIDIAの最新GPUアーキテクチャ「Hopper」に搭載されている「TMA(Tensor Memory Accelerator、大量のデータ転送を高速に行うための専用のハードウェアエンジン)」を活用し、大量のリモートストア(遠隔地への書き込み)を極めて効率的に処理できる設計とした。
このNVSHMEMを適用した最適化モデルを、NVIDIAのスーパーコンピューター「Eos」および「NVIDIA GB200 NVL72」のクラスター環境で検証した。その結果、従来のGPU-aware MPIを用いたシステムと比較して、特に遅延(レイテンシ)の影響を受けやすいシステムにおいて、ノード内およびノード間の双方で最大2倍の強スケーリング性能(計算規模を固定したまま並列処理を増やした際の性能効率)向上が実証された。
標準化への課題
NVIDIAは、この手法がハローエクスチェンジのパターンを用いる他のHPCアプリケーションにも広く一般化できると説明している。一方で、GPUネイティブな通信プリミティブ(基本的な通信操作)の標準化は、今後の普及に向けた継続的な課題として残されているという。