NVIDIA、PET原子シミュレーションでTensor Core FP32を検証、カーネルと全工程の倍率は別

NVIDIAは10月29日、機械学習による原子間ポテンシャルで精度を保ちながら計算を速める実験を公表した。11月1日は本サイトの技術号の日付である。対象は特定のPoint Edge Transformer、すなわちPETモデルとGPUであり、あらゆる科学モデルが同じ倍率で速くなるという報告ではない。行列計算のカーネルと、分子動力学の一歩全体の倍率を分ける必要がある。

原子間ポテンシャルは配置からエネルギーを予測し、その勾配から得た力が多くの時刻の位置と速度を動かす。小さな数値誤差でも軌跡に蓄積し得るため、推論にFP32を使うモデルが多い。NVIDIAの方法はPETを再学習するものではない。FP32の値を低精度の成分に分解し、Tensor Coreで積を計算してFP32で加算する。Blackwell向けBF16x9は各値を三成分にして九つの積を使い、Ampere向けTF32x3は主成分と残差から三つの積を使って最小の残差同士の積を省く。二方式の誤差を同じものとして扱えない。

試験には2600万パラメータのPET-MAD小モデルを使い、192原子の水箱と1万7280原子の超胞を調べた。2億5500万パラメータの大モデルは192原子の系で試した。各条件の比較対象は固有の通常FP32経路である。Blackwell B300のBF16x9では、大きい系の主要GEMMが2.2～2.5倍、分子動力学の一歩全体はperformant方式で1.7倍、eager方式で1.4倍だった。192原子の一歩全体はそれぞれ1.2倍と1.1倍にとどまる。カーネルの2.5倍をシミュレーション全体の倍率に使ってはならない。

Ampere A100でのTF32x3はCUTLASSの`OpMultiplyAddFastF32`と独自の`TF32Linear`で互換の線形層を置き換える。入力と出力はFP32のままである。大きい系の主要GEMMはFP32比で約1.7～1.9倍、一歩全体は約1.5倍速い。192原子の大モデルでは一歩が1.4倍、小モデルでは1.2倍だった。単純な一回のTF32はさらに速い場合がある一方、精度と軌跡安定性で不利になるとNVIDIAは報告する。算術の最高速度と有用な物理計算は同義ではない。

物理的な振る舞いも調べられた。CO分子をPt(111)表面から動かす試験ではTF32x3のエネルギー曲線はFP32に近く滑らかで、一回のTF32には揺れが見えた。192原子の水の30ピコ秒NVE計算でも、FP32とTF32x3は有界のエネルギー変動を示し、一回のTF32にはより大きな系統的ドリフトが現れた。これは選んだ材料と時間幅の結果であり、すべての物質の長期保存則を保証しない。

539原子の合金板の構造緩和では、FP32の三試行は218ステップで収束し、TF32x3は218～222ステップで1.0×10⁻³ eV/Åの力の基準に到達した。一回のTF32の三試行は500ステップ以内に収束せず、上限を1000にしても達しなかったという。一歩当たりではTF32x3が1.5倍、一回のTF32は2.0倍速かったが、後者は必要な収束を得ていない。速度だけの比較では実用上の結論を取り違える。

利用者は自分の負荷で大きなGEMMの時間割合を測り、エネルギー、力、長い軌跡と収束を確認すべきだ。小さい系では近傍処理や起動費用が全体の利益を薄める。この記事の倍率はNVIDIAがPETと指定GPUで報告したもので、本サイトによる再計測ではない。モデル規模、原子数、測定項目、FP32基準を添えたまま判断する必要がある。