Machine-learning interatomic potentials leave FP32-only inference: PET runs on Tensor Cores, up to 1.7x faster MD steps
Synopsis
Machine-learning interatomic potentials (MLIPs) have largely kept inference in FP32 because molecular dynamics is sensitive to small force errors. For PET, an MLIP from EPFL's Lab COSMO, BF16x9 on Blackwell GPUs and TF32x3 on Ampere GPUs decompose each FP32 matrix-multiply operand into lower-precision components evaluated on Tensor Cores, with no retraining or architecture change. On a 17,280-atom water box, BF16x9 sped PET's large GEMMs by 2.2–2.5x and the full MD step by up to 1.7x; TF32x3 sped major GEMMs on A100 by about 1.7–1.9x and the MD step by up to 1.5x, while keeping energy conservation and structural relaxation close to FP32.
Interpretation
MLIP inference can keep FP32-like numerical behavior while using the higher low-precision throughput of Tensor Cores, without retraining the model or changing its transformer architecture. Many MLIPs had continued to run inference in FP32 even though modern NVIDIA GPUs offer substantially higher Tensor Core throughput for BF16 and TF32 than for FP32 matrix multiplication. On PET, the error from the BF16x9 path stayed below the model's own level of non-determinism, and the path exposed performance headroom even though PET already expresses much of its runtime as optimized GEMMs.
The approach splits each FP32 operand into lower-precision components, evaluates selected cross-products on Tensor Cores, and recombines them with FP32 accumulation. On Blackwell, BF16x9 decomposes each FP32 operand into three BF16 components and recombines nine pairwise products, enabled through a cuBLAS environment variable; on Ampere, TF32x3 uses a leading and a residual TF32 component, evaluates three products and omits the residual-by-residual term, provided by CUTLASS's OpMultiplyAddFastF32. Evaluation used the 26M-parameter PET-MAD small checkpoint (pet-mad-s-v1.6.0-nn-32) on a 192-atom water box and a 17,280-atom supercell, plus the 255M-parameter large checkpoint (pet-mad-m-v1.6.0) on the 192-atom system, comparing energies, forces and stresses with the native FP32 path.
Beyond single-point accuracy, TF32x3 preserved FP32-like behavior in potential-energy smoothness, energy conservation and structural relaxation, tests that accumulate over time. Matching isolated energy and force predictions is only the first test, because forces are gradients of the potential-energy surface and are integrated over many timesteps, so small numerical irregularities can roughen the surface or cause energy drift. Dragging a single CO molecule rigidly from an on-top site of a Pt(111) slab across 60 carbon heights, TF32x3 followed the smooth FP32 energy curve while one-pass TF32 showed visible point-to-point variations; in an NVE simulation of a 192-atom water box, FP32 and TF32x3 showed comparable bounded energy fluctuations while one-pass TF32 drifted more; relaxing a 539-atom Cantor-alloy slab, FP32 converged in 218 steps and TF32x3 in 218–222 steps, all reaching the 1.0×10⁻³ eV/Å force tolerance, while all three one-pass TF32 replicates failed within the 500-step limit.
Perspective
The demonstrated scope is PET-MAD's 26M and 255M checkpoints on a 192-atom water box, a 17,280-atom supercell and a 539-atom Cantor-alloy slab, on Blackwell (compute capability 10.0 or 10.3) and Ampere A100 hardware; end-to-end gains depend on system size and the share of runtime spent in GEMMs, and small systems benefit less because launch overhead, neighbor processing and graph operations take a larger share.
TF32x3 has a slightly different accuracy profile from BF16x9 and should be validated against the model's accuracy and physical-stability requirements; TF32x3 performance also depends on kernel tuning for the model's matrix shapes and hardware; the combined gain with graph-construction optimization, fused operations, custom kernels and CUDA graphs was not measured and should be established end to end for each workload.