NVIDIA Tests Tensor Core FP32 Methods on PET Atomistic Simulation; Kernel and Full-Step Gains Differ
NVIDIA published an experiment on precision and speed for machine-learning interatomic potentials on October 29. This site covers it in a November 1 technical issue without changing the original date. The work evaluates Tensor Core methods for FP32-like matrix multiplication in selected Point Edge Transformer, or PET, workloads.
NVIDIA published an experiment on precision and speed for machine-learning interatomic potentials on October 29. This site covers it in a November 1 technical issue without changing the original date. The work evaluates Tensor Core methods for FP32-like matrix multiplication in selected Point Edge Transformer, or PET, workloads. It is a vendor experiment on particular models and GPUs, not evidence that every scientific model receives the same acceleration. Kernel-level speedups and whole molecular-dynamics-step speedups must be kept separate throughout.
An interatomic potential predicts an energy from atomic positions; forces derived from that energy drive successive updates of positions and velocities. Small numerical errors can accumulate across a long trajectory, so many such models still use FP32 inference. NVIDIA's approach does not retrain PET or change its architecture. It decomposes FP32 operands into lower-precision components, computes selected cross-products on Tensor Cores, and accumulates in FP32. BF16x9 on Blackwell uses three BF16 components per operand and nine products. TF32x3 on Ampere uses leading and residual TF32 components and evaluates three products, omitting the smallest residual-by-residual term. Their accuracy and performance profiles cannot be combined into one generic “Tensor Core FP32” number.
The test cases make the denominators clear. NVIDIA used a 26-million-parameter PET-MAD small checkpoint on a 192-atom water box and a 17,280-atom supercell, and a 255-million-parameter large checkpoint on the 192-atom system. Predicted energy, forces and stresses were compared with each workload's native FP32 path. For Blackwell BF16x9, large GEMM kernels in the 17,280-atom workload were 2.2–2.5 times faster. The complete molecular-dynamics step on that large system improved by 1.7 times under the “performant” strategy and 1.4 times under the “eager” strategy. On the smaller 192-atom case, those full-step improvements were 1.2 and 1.1 times. A 2.5-times kernel result is therefore not a 2.5-times simulation result.
The Ampere experiment uses TF32x3 through a CUTLASS `OpMultiplyAddFastF32` kernel and a custom `TF32Linear` replacement for compatible linear layers. Inputs and outputs remain FP32. On an A100, major GEMMs for the 17,280-atom case were approximately 1.7–1.9 times faster than the corresponding FP32 baseline, while the complete step improved by about 1.5 times on the larger system. For the 255-million-parameter checkpoint on 192 atoms, the complete-step gain was 1.4 times, compared with 1.2 times for the small checkpoint on that size. NVIDIA notes that one-pass TF32 can deliver a higher raw speedup but with a greater accuracy and stability cost. The fastest arithmetic choice was not automatically the useful simulation choice.
NVIDIA therefore tested physical behavior beyond one predicted energy or force. In one test it moved a CO molecule away from a Pt(111) surface and examined the energy curve; in another it ran a 30-picosecond NVE simulation of a 192-atom water box. TF32x3 followed a smooth FP32-like curve and comparable bounded energy fluctuations, whereas one-pass TF32 showed visible point-to-point variation or a larger systematic drift. These are tests of selected physical properties in selected systems, not a proof of universal conservation for all materials or time scales.
Structural relaxation provides a further check. For a 539-atom Cantor-alloy slab, three FP32 replicates converged in 218 steps. TF32x3 reached the force target of 1.0 × 10⁻³ eV/Å in 218–222 steps. All three one-pass TF32 replicates failed to converge within 500 steps, and the source reports failure even when its limit was doubled to 1,000. TF32x3 sped each optimization step by 1.5 times while retaining the tested convergence behavior; one-pass TF32 could run a step about twice as fast but missed the target. Comparing only per-step speed would invert the practical conclusion.
An adopter should measure the fraction of runtime spent in large GEMMs and validate energy, force, trajectory stability and convergence on its own model. Small systems may see less benefit because graph construction, neighbor processing and launch overhead occupy a larger share of time. Kernel tuning and hardware generation also matter. NVIDIA explicitly says it did not measure a combined gain with other possible optimizations. This publication has not independently rerun the PET benchmarks. Every quoted speed ratio here remains tied to the vendor, hardware, model size, atom count, metric and native FP32 baseline that produced it.
Sources
FAQ
How do BF16x9 and TF32x3 differ in arithmetic and GPU scope?
NVIDIA's BF16x9 on Blackwell splits each FP32 operand into three BF16 components and evaluates nine cross-products. TF32x3 on Ampere evaluates three TF32 products and omits the smallest residual product. Both accumulate in FP32 without retraining PET, but they use different hardware, implementation paths and accuracy profiles.
Can the reported 2.5-times and 1.7-times figures both describe full simulation speed?
No. On the 17,280-atom PET workload, NVIDIA reports up to 2.5 times faster large GEMM kernels with B300 BF16x9, but 1.7 times for the complete MD step under its performant strategy. A100 TF32x3 has separate roughly 1.7–1.9-times kernel and 1.5-times full-step results. Keep metric, baseline, GPU and system size attached.
Why did the study test trajectories and structural convergence as well as single predictions?
Forces derive from predicted energy, so small numerical errors can accumulate across simulation steps. NVIDIA checked a CO/Pt(111) energy curve, a 30-picosecond water trajectory and relaxation of a 539-atom alloy. TF32x3 converged in 218–222 steps in that relaxation; faster one-pass TF32 did not converge within 500. Other models still need their own validation.