Equivariant materials models run faster with new low precision method
Accelerating SO(2)-Equivariant Materials Models through Low-Precision Emulation
Machine LearningDistributed, Parallel, and Cluster Computing
Summary
Machine learning models that understand materials often need very precise calculations, which makes them slow. The authors introduce a new technique called FlashSO(2) that speeds up these calculations by using a special low-precision method without losing accuracy. This lets these models run up to 4 times faster in training and more than 3 times faster in some simulations, while also using less memory. The approach is especially helpful for models that rely on rotations and symmetries to better predict material behavior.
What this means in practice
- •For materials simulation teams: Train and run ML models of materials faster and with less memory using optimized low-precision tensor operations for rotationally symmetric architectures.
- •For molecular dynamics developers: Accelerate molecular dynamics simulations by integrating low-precision equivariant graph neural networks that maintain accuracy with higher throughput and lower memory.
Authors
Denghui Lu, Alexander Maeder, Alexandros Nikolaos Ziogas, Mathieu Luisier, Manasa Kaniselvan
Abstract
Rotationally equivariant architectures improve the physical fidelity of machine-learned (ML) universal materials models based on graph neural networks (GNNs). However, their scalability is limited by the use of non-trivial, high-precision tensor operations during both training and inference. Adopting the reduced-precision formats of matrix multiplication units (MMUs), e.g., FP8 instead of FP32, provides massive increases in throughput, up to 30x (60x) on NVIDIA Hopper (Blackwell) GPUs, but may lead to drastic loss of accuracy. Here, we introduce tailored 'emulated' low-precision operations under the integer Ozaki scheme to accelerate the 'SO(2)-equivariant' processing layers of GNNs, which dominate the runtime of (1) many state-of-the-art ML interatomic potentials (MLIPs) and (2) most universal ML Hamiltonian (MLH) models. Our framework, FlashSO(2), outperforms TF32, the most accurate native low-precision MMU format, in throughput, accuracy, and memory footprint at once. At the SO(2) layer level, it reaches 40x (7x) the throughput of MALOQ's (UMA's) reference implementation, an MLH (MLIP) model, sustaining up to 236 (232) effective Tflop/s on an NVIDIA GH200, 3.5x the FP32 peak. At the model level, FlashSO(2) helps train MALOQ 3.9x faster than its pure-FP32 version while reducing peak memory by 58% at no loss in accuracy. In molecular dynamics simulations with UMA, the speedup factor is 3.3x, and the peak memory is reduced by 51%. FlashSO(2) thus pushes the speed-accuracy Pareto frontier of equivariant materials models beyond the native low-precision formats and substantially lowers the resources required to exploit their scaling laws.