机器学习原子间势不再只跑FP32：PET换用张量核心低精度分解，MD步最高快1.7倍
核心概要
机器学习原子间势（MLIP）的推理长期停留在FP32，因为分子动力学对力的小误差很敏感。针对EPFL Lab COSMO开发的PET模型，在Blackwell GPU上用BF16x9、在Ampere GPU上用TF32x3把FP32矩阵乘分解为低精度分量在张量核心上计算，无需重训练或改架构。在17,280原子水盒上，BF16x9使PET的大矩阵乘快2.2–2.5倍、完整MD步最高快1.7倍；TF32x3在A100上使主要GEMM快约1.7–1.9倍、MD步最高快1.5倍，并保持与FP32接近的能量守恒与结构弛豫行为。
深度剖析
机器学习原子间势的推理可以继续用FP32的数值行为，同时吃到张量核心的低精度吞吐，而不必重训练模型或改动其transformer架构。 此前许多MLIP坚持用FP32推理，尽管现代NVIDIA GPU在BF16和TF32上的张量核心吞吐明显高于FP32矩阵乘。 在PET上，BF16x9带来的误差低于模型自身的非确定性水平；PET的运行时本就大量花在高度优化的GEMM上，仍出现了性能余量。
做法是把每个FP32操作数拆成若干低精度分量，在张量核心上算选定的交叉乘积，再用FP32累加重组。 Blackwell上的BF16x9把每个FP32操作数拆成三个BF16分量并重组九个两两乘积，可通过cuBLAS环境变量开启；Ampere上的TF32x3用前导与残差两个TF32分量算三个乘积，省略残差乘残差项，由CUTLASS的OpMultiplyAddFastF32提供。 评测使用26M参数的PET-MAD小检查点（pet-mad-s-v1.6.0-nn-32）在192原子水盒与17,280原子超胞上，以及255M参数的大检查点（pet-mad-m-v1.6.0）在192原子体系上，与原生FP32路径比较能量、力和应力。
除单点精度外，TF32x3在势能面平滑性、能量守恒和结构弛豫这些随时间累积的检验上也保持了FP32式的行为。 单点能量和力吻合只是数值保真的第一道检验，力是势能面的梯度并在许多时间步上积分，微小数值不规则会粗糙化势能面或造成能量漂移。 把单个CO分子从Pt(111)表面的顶位刚性拖到60个碳高度上，TF32x3跟随FP32的平滑能量曲线，而单遍TF32出现可见的逐点波动；192原子水盒的NVE模拟中两者能量涨落相当有界，单遍TF32出现更大系统性漂移；539原子Cantor合金板的弛豫中，FP32三步均为218步收敛，TF32x3为218–222步，均达到1.0×10⁻³ eV/Å力容差，而单遍TF32三次均在500步上限内未收敛。
启示与展望
演示范围限于PET-MAD的26M与255M检查点，体系为192原子水盒、17,280原子超胞和539原子Cantor合金板，硬件为Blackwell（计算能力10.0或10.3）与Ampere A100；端到端收益取决于体系大小和GEMM在运行时的占比，小体系因启动开销、近邻处理和图操作占比更高而受益更少。
TF32x3的精度轮廓与BF16x9略有不同，需要针对模型自身的精度与物理稳定性要求单独验证；TF32x3的性能还依赖针对模型矩阵形状与硬件的核函数调优；与图构建优化、算子融合、自定义核函数和CUDA图的组合增益未在本次评测中测量，需按负载端到端确定。