NVIDIA 用 Tensor Core FP32 加速机器学习原子间势推理
原标题:Unlocking Tensor Core Performance for Machine Learning Interatomic Potentials
AI 摘要
NVIDIA 技术博客介绍如何用 Tensor Core FP32 方法加速机器学习原子间势(MLIP)推理。文章以 EPFL 的 PET 模型为测试对象,在 Blackwell GPU 上用 BF16x9、在 Ampere GPU 上用 TF32x3 分解 FP32 运算,在保持接近 FP32 精度的同时提升矩阵乘法吞吐。实测 17,280 原子体系大 GEMM 加速 2.2–2.5 倍,MD 单步最高加速 1.7 倍,且无需重训练或修改模型结构。
正文节选
Atomistic simulations help researchers understand how molecules and materials behave, from chemical reactions and phase transitions to mechanical failure. These simulations come with a cost-accuracy trade-off. Quantum-mechanical methods such as density functional theory (DFT) can provide accurate energies and forces, but their computational cost often limits their use for large systems or long simulations. Classical interatomic potentials are much faster, but their fixed functional forms limit a