NVIDIA published results accelerating PET, a neural network that predicts atomic energies and forces, without retraining. Its tests sped up both matrix operations and full molecular dynamics steps that update atomic positions and velocities. This showed potential to shorten simulations with existing models.
In these calculations, small force errors can accumulate at each step and disrupt atomic motion. This is why using Tensor Cores, GPU units dedicated to matrix operations, requires both speed and numerical stability (original post).
Measurement methods
NVIDIA tested the 26-million-parameter PET-MAD on a 192-atom water box and an expanded 17,280-atom structure. BF16x9 was measured on Blackwell B300 and TF32x3 on A100. Matrix multiplication used each GPU's native FP32 baseline, while BF16x9 molecular dynamics steps used compiled FP32 (methods).
Results
Matrix multiplication (GEMM) speedups must be considered separately for each device.
Table: BF16x9 GEMM speedup on B300
Conditions: PET-MAD 26M, relative to native FP32 on Blackwell B300, bar values from NVIDIA Figure 3
Unit: x
| GEMM shape MxNxK | BF16x9 B300 |
|---|---|
| 17280x4096x1024 | 2.49 |
| 17280x1024x2048 | 2.38 |
| 967680x768x256 | 2.23 |
| 192x4096x1024 | 1.21 |
| 10752x768x256 | 1.43 |
| 4096x4096x4096 | 2.41 |
| 8192x8192x8192 | 2.28 |
Table: TF32x3 GEMM speedup on A100
Conditions: PET-MAD 26M, relative to native FP32 on A100, bar values from NVIDIA Figure 5
Unit: x
| GEMM shape MxNxK | TF32x3 A100 |
|---|---|
| 17280x4096x1024 | 1.81 |
| 17280x1024x2048 | 1.83 |
| 967680x768x256 | 1.65 |
| 192x4096x1024 | 1.09 |
| 10752x768x256 | 1.87 |
| 4096x4096x4096 | 1.94 |
| 8192x8192x8192 | 1.72 |
NVIDIA measured notable BF16x9 gains for large matrices on B300 (Figure 3). On A100, TF32x3 beat its own FP32 baseline for every listed shape (Figure 5). Different devices mean these results do not rank the two methods against each other.
The BF16x9 performant strategy is used only when cuBLAS expects it to be faster. NVIDIA measured full-step speedups over compiled FP32 of 1.2x for 192 atoms and 1.7x for 17,280 atoms. Its measured TF32x3 speedup for the large system was 1.5x over its own FP32 baseline (Figures 4 and 6 and discussion).
Smaller systems gain less because GPU launch overhead and non-matrix work, such as neighbor processing, take a larger share of time. TF32x3 requires kernel tuning for matrix shapes and GPUs, plus validation against model accuracy and physical stability requirements (limits and selection criteria).
In NVIDIA's accuracy tests, BF16x9 errors were smaller than the variability seen when running the same model. Its measured TF32x3 force error was about 3x the difference seen after recompiling FP32 and running it (precision analysis).
Structural relaxation of a 539-atom alloy surface used PET, FIRE2, and NVIDIA ALCHEMI Toolkit to find a stable arrangement. Across NVIDIA's 3 runs, FP32 reached the force tolerance of 1.0×10⁻³ eV/Å in 218 steps, and TF32x3 in 218–222 steps. The single-conversion TF32 path was faster per step but failed to converge even after 1,000 steps (Figure 8 and discussion).
In NVIDIA's test detaching a carbon monoxide molecule from a platinum surface, TF32x3 followed FP32's smooth energy curve. A separate 30 ps simulation of 192-atom water at fixed particle count, volume, and energy showed energy fluctuations similar to FP32. In NVIDIA's tests, single-conversion TF32 showed curve oscillations and energy drift, respectively (Figure 7 and discussion).
Why the results differ
BF16x9 splits FP32 inputs into three BF16 components, computes nine pairwise products, and sums them in FP32 (cuBLAS documentation). TF32x3 splits inputs into a large component and an approximation residual. It accumulates only three products in FP32, omitting the product of the small residuals (algorithm).
PET obtains forces by differentiating predicted energies. TF32x3 therefore also applies to computing derivatives passed to neural network layers, so energy and forces use the same computation path (implementation).
Selection criteria
BF16x9 is supported with CUDA 12.9 or later and compute capabilities 10.0 and 10.3, which identify GPU architecture support. Enabling CUBLAS_EMULATE_SINGLE_PRECISION=1 allows compatible PyTorch matrix operations to use this path (requirements).
TF32x3 on A100 needs a separate implementation. NVIDIA replaced PET's nn.Linear with TF32Linear, which calls CUTLASS kernels. The key operator setting in the original post also appears in CUTLASS example 27.
cutlass::arch::OpMultiplyAddFastF32The practical value of accelerating PET lies in the total time of a simulation that finishes reliably, not in peak matrix multiplication speed.