NVIDIA Tensor Core|ML原子間ポテンシャルの計算速度向上と精度維持の手法

Blockchain News アメリカ
概要
NVIDIAは、モデルの再学習やアーキテクチャ変更なしに、機械学習(ML)原子間ポテンシャルの計算を高速化するTensor Core FP32手法を公開しました。Blackwell GPUのBF16x9やAmpere GPUのTF32x3を用いることで、精度をFP32ベースラインに近づけたまま、行列演算(GEMM)や分子動力学ステップの速度を向上させています。原子レベルのシミュレーションで課題となる数値的な安定性と計算速度の両立を実現するアプローチとなります。
詳細

主要成果

NVIDIAが2026年10月29日に公開した開発者ブログによると、同社のTensor Core FP32手法を用いることで、機械学習原子間ポテンシャルの計算速度を大幅に向上させることが可能です。EPFLのLab COSMOが開発したTransformerベースのポテンシャル「Point Edge Transformer(PET)」を用いた検証では、モデルの再学習や構造変更を行うことなく、FP32ベースラインに近い挙動を維持したまま高速化を達成しました。

Blackwell GPUにおけるBF16x9パスでは、17,280原子のワークロードにおいて大規模なPET GEMM(一般行列乗算)が2.2倍から2.5倍に加速されました。分子動力学(MD)ステップの速度向上は、192原子の水ボックスで1.2倍〜1.4倍、17,280原子のシステムで1.4倍〜1.7倍(cuBLAS戦略が「performant」か「eager」かによる)となりました。また、BF16x9の誤差はPET自体の非決定性のレベルを下回る結果となっています。

Ampere GPUのTF32x3では、A100 Tensor Core GPUにおいて大規模PET GEMMが約1.7倍〜1.9倍に加速されました。2,600万パラメータのチェックポイントでは、大規模システムでのMDステップが1.5倍高速化し、2億5,500万パラメータのチェックポイントでは192原子システムで1.4倍の速度向上を記録しました。この際の力誤差は3.9 × 10⁻⁶ eV/Åでした。

技術・詳細

本手法の核心は、数値精度の維持と計算効率の最適化にあります。原子レベルのシミュレーションは微小な数値差が蓄積しやすいため、精度管理が極めて重要です。検証の結果、TF32x3はFP32と同様に滑らかなCO/Pt(111)ポテンシャルエネルギー曲線を描き、192原子の水ボックスを用いた30 psのNVEシミュレーションでも、FP32と同等の限定的なエネルギー変動を示しました。対照的に、1パスのTF32では点間の変動が顕著であり、エネルギーのドリフトも大きくなる傾向が見られました。

また、NVIDIA ALCHEMI ToolkitのFIRE2オプティマイザーを用いた539原子のカンター合金スラブの構造緩和テストでは、FP32が218ステップで収束したのに対し、TF32x3は3回の試行で218〜222ステップで収束し、いずれも1.0 × 10⁻³ eV/Åの力許容値に達しました。一方で、1パスのTF32は500ステップ以内に収束せず、制限を1,000ステップに倍増させても収束に失敗しました。

実装面では、BlackwellのBF16x9は各FP32オペランドを3つのBF16コンポーネントに分解し、9つの積をFP32で再結合します。これは環境変数「CUBLAS_EMULATE_SINGLE_PRECISION=1」などで有効化できます。AmpereのTF32x3は、FP32オペランドを主要なTF32コンポーネントと残差コンポーネントに分割し、3つのクロス積を評価して残差同士の項を省略する手法です。これはCUTLASSのOpMultiplyAddFastF32オペレーターを使用し、モデルインターフェースでのFP32入出力を維持しています。

背景・業界の文脈

原子間ポテンシャルの計算において、実行時間の多くは最適化済みの一般行列乗算(GEMM)に費やされます。そのため、さらなる最適化の余地(ヘッドルーム)は限られていますが、ここを効率化することがエンドツーエンドの性能向上に直結します。ただし、得られる恩恵はシステムサイズに依存します。小規模なシステムでは、カーネルの起動オーバーヘッドや近接原子の処理、グラフ操作の割合が大きいため、GEMMの高速化によるメリットは相対的に小さくなります。一方で、大規模システムほどGEMMの計算量が増えるため、本手法による加速効果が顕著に現れます。

今後の展望

NVIDIAは、TF32x3の性能がモデルの行列形状やハードウェアに対するカーネルチューニングに依存することを指摘しています。また、BF16x9およびTF32x3の導入にあたっては、個別のワークロードにおける精度および安定性の要件に照らして検証を行う必要があるとしています。本報告は製品の発売ではなく技術的なパフォーマンスレポートであり、具体的な価格や市場への影響、広範な採用見込みについては言及されていません。

元記事: https://blockchain.news/news/nvidia-blog-tensor-core-speedups-ml-interatomic-potentials

毎週の技術動向レポートを無料でお届け

各分野の分析レポートを読む価値があるかどうか一目で判断できるインフォグラフィックをメールで受け取れます。

📢 メールマガジンに無料登録(週刊・技術動向レポート)

ご登録いただくと、Troy-Technical から週刊で技術動向レポート(メールマガジン)をお届けします。

  • 取得したメールアドレス・選択分野は配信目的にのみ使用します。
  • 第三者へ提供することはありません。
  • 配信はいつでも解除できます(各メール下部のリンクから)。

詳しくはプライバシーポリシーをご覧ください。

登録は1分・いつでも解除できます

この記事は Troy-Technical が公開しています。素材メーカーで技術開発に携わってきた技術者が、個人で運営しているサイトです。
運営者について / お問い合わせ info@troy-technical.jp
よかったらシェアしてね!
  • URLをコピーしました!
  • URLをコピーしました!

この記事を書いた人

目次