輝達B300測試：PET分子動力學完整步驟最高快1.7倍

輝達（NVIDIA，$NVDA）以B300 GPU搭配BF16x9技術，在由17,280個原子組成的水系統中執行PET分子動力學計算，速度較FP32最高快1.7倍。這項數據衡量的是完整一個分子動力學步驟，而非只有矩陣乘法。

輝達於29日在技術部落格公布，將兩種以Tensor Core為基礎的矩陣乘法技術應用於機器學習原子間位能模型PET（Point Edge Transformer）。BF16x9在B300（Blackwell）GPU上測試，TF32x3則在A100 GPU上測試。性能數據是輝達以特定模型和設備自行測得的結果。

機器學習原子間位能模型會根據原子結構估算能量與原子間作用力，用於計算分子和材料的運動。此次測試檢視以Tensor Core執行運算量較大的FP32矩陣乘法時，效能和結果會如何變化。

PET由瑞士洛桑聯邦理工學院（EPFL）COSMO實驗室開發。輝達以參數量2,600萬個的PET-MAD小型模型，在含192個與17,280個原子的水系統中進行評估；參數量2億5,500萬個的大型模型則在含192個原子的系統中測試。

BF16x9會將FP32數值拆分成三個BF16分量，計算九個部分乘積，再將結果加總為FP32。這種方法的設計目的，是讓大型矩陣乘法能在Tensor Core上執行。

在含17,280個原子的系統中，BF16x9的主要矩陣乘法比FP32快2.2至2.5倍；若計入矩陣乘法以外的處理，一個分子動力學步驟依運作方式不同，快1.4至1.7倍。

在含192個原子的系統中，BF16x9的主要矩陣乘法快1.2至1.4倍，分子動力學步驟快1.1至1.2倍。輝達表示，以此方法計算單一時點的能量與作用力，其誤差小於PET模型本身的非確定性。

A100上的TF32x3會把FP32數值拆成主成分與殘差成分，執行三次TF32乘法，再將結果加總為FP32。這種方法採用的精度格式與BF16x9不同。

在含17,280個原子的系統中，TF32x3的主要矩陣乘法快1.7至1.9倍，分子動力學步驟最高快1.5倍。在含192個原子的系統中，小型模型的分子動力學步驟快1.2倍，大型模型則快1.4倍。

兩種方法都只改變矩陣乘法的執行方式，沒有調整已訓練模型的權重或架構，因此不必重新訓練模型；但套用方法不同。BF16x9使用cuBLAS支援的功能，TF32x3則在PET的線性層套用獨立的核心程式。

輝達以TF32x3呈現長時間能量守恆與結構最佳化測試結果。在含192個原子的水系統中，能量變動幅度與FP32相近；在鉑表面移動一氧化碳分子時，測得的能量曲線也與FP32一樣平滑。

在含539個原子的Cantor合金表面結構最佳化測試中，FP32經218次迭代、TF32x3經218至222次迭代後，達到1.0×10⁻³ eV/Å的作用力容許值。單次TF32在500次迭代內未收斂，計算至1,000次仍未達到相同標準。

輝達未提供BF16x9在同類長時間能量守恆與結構最佳化測試中的結果，因此不能據此認定兩種方法都已在相同範圍內驗證物理一致性。

性能差異會隨系統規模，以及矩陣乘法在整體計算中的占比而變化。輝達指出，這些數據適用於此次列出的PET模型、GPU與核心程式設定。

NVIDIA cuBLAS文件列出的BF16x9 FP32模擬支援對象，為Compute Capability 10.0與10.3的GPU，並要求使用CUDA 12.9以上版本。