NVIDIA Tensor Cores: Speedups for ML Interatomic Potentials

Blockchain News USA

Overview

NVIDIA’s Tensor Core FP32 methods, specifically BF16x9 for Blackwell and TF32x3 for Ampere, accelerate machine learning interatomic potentials without requiring model retraining. In tests using the Point Edge Transformer (PET), these methods achieved up to 2.5x GEMM acceleration while maintaining numerical stability close to FP32 baselines. This allows for faster molecular dynamics simulations without the convergence failures seen in standard one-pass TF32.

In Depth

Key Findings

NVIDIA reports that its Tensor Core FP32 methods significantly accelerate the Point Edge Transformer (PET), a transformer-based machine learning interatomic potential developed by EPFL’s Lab COSMO. On Blackwell GPUs using the BF16x9 path, large PET general matrix multiplications (GEMMs) were accelerated by 2.2x to 2.5x for a 17,280-atom workload. End-to-end molecular-dynamics step speedups ranged from 1.2x to 1.4x for a 192-atom water box and 1.4x to 1.7x for the 17,280-atom system. On Ampere A100 GPUs using TF32x3, large GEMMs saw speedups of 1.7x to 1.9x, with molecular-dynamics step speedups of 1.5x for a 26-million-parameter checkpoint on the larger system and 1.4x for a 255-million-parameter checkpoint on the 192-atom system.

Technical Details

• BF16x9 (Blackwell): Decomposes each FP32 operand into three BF16 components and recombines nine products in FP32. This can be enabled via the CUBLAS_EMULATE_SINGLE_PRECISION=1 environment variable.
• TF32x3 (Ampere): Splits each FP32 operand into a leading TF32 component and a residual component, evaluating three cross-products and omitting the residual-by-residual term. This implementation uses CUTLASS’s OpMultiplyAddFastF32 operator.

Numerical Stability and Validation

Atomistic simulations are sensitive to numerical drift. In a 30 ps NVE simulation of a 192-atom water box, TF32x3 showed bounded energy fluctuations comparable to FP32, whereas one-pass TF32 exhibited more drift. In a structural-relaxation test on a 539-atom Cantor-alloy slab using the FIRE2 optimizer, FP32 converged in 218 steps and TF32x3 in 218 to 222 steps. Conversely, one-pass TF32 failed to converge within 1,000 steps.

Significance and Outlook

The degree of end-to-end gain depends on system size; larger systems expose more GEMM work, while smaller systems are limited by launch overhead and graph operations. NVIDIA notes that these methods allow for acceleration without changing model architecture or retraining, though they should be validated against specific accuracy and stability requirements for each workload.

Get our weekly technology intelligence — free

Receive an infographic that lets you judge at a glance whether each field’s analysis report is worth reading.

Subscribe Free — Weekly Tech Intelligence

By subscribing, you’ll receive Troy-Technical’s weekly technology intelligence newsletter.

• Your email and selected fields are used only to deliver the newsletter.
• We never share your information with third parties.
• You can unsubscribe anytime via the link in each email.

See our Privacy Policy for details.

Agree & Continue ▶

Takes about a minute · Unsubscribe anytime