NVIDIA 在 PET 原子模拟测试 Tensor Core FP32 方法：须分开看矩阵核与完整步骤加速

NVIDIA 于 10 月 29 日发表机器学习原子间势（MLIP）的精度与速度实验，本站编入 11 月 1 日技术刊期。它展示在特定 PET 模型和 GPU 上用 Tensor Core 执行近似 FP32 矩阵乘法的结果，并未宣称任意科学模型都能获得相同加速。分清矩阵核的加速与完整分子动力学步骤的加速，是阅读这篇厂商实验的第一道门槛。

MLIP 用学习到的能量预测原子系统的力，力又影响多步模拟的速度和位置。微小数值误差可能在轨迹中积累，因此许多模型保留 FP32 推理。NVIDIA 的办法不是重训 PET 或换模型结构，而是把 FP32 操作数拆成低精度部分，用 Tensor Core 计算交叉乘积，再以 FP32 累加。Blackwell 上的 BF16x9 使用三个 BF16 分量形成九个乘积；Ampere 上的 TF32x3 计算三个主要乘积，省去最小的残差相乘项。两种方法的误差与性能不能当成同一组数字。

测试对象是 EPFL Lab COSMO 的 Point Edge Transformer。NVIDIA 在 192 原子水盒与 17,280 原子超胞上测试 2600 万参数的小检查点，并在 192 原子系统测试 2.55 亿参数的大检查点，基线是各自的原生 FP32 路径。BF16x9 在 Blackwell B300 的大矩阵核上，对 17,280 原子负载给出 2.2–2.5 倍核速度；完整分子动力学步骤在同一大系统上，按 performant 策略快 1.7 倍，eager 策略快 1.4 倍。在 192 原子系统，完整步骤分别只有 1.2 与 1.1 倍。核加速显著高于完整步骤，不能互换引用。

TF32x3 的实验在 Ampere A100 上，用 CUTLASS 的 `OpMultiplyAddFastF32` 和自定义 `TF32Linear` 替换兼容线性层，输入输出仍是 FP32。17,280 原子工作负载的主要 GEMM 核约快 1.7–1.9 倍；较大系统的完整 MD 步骤约快 1.5 倍。大检查点在 192 原子系统的完整步骤约快 1.4 倍，而小检查点同尺寸约 1.2 倍。厂商还说明更简单的一次 TF32 虽可能更快，但数值精度与轨迹稳定性付出代价。

速度之外，NVIDIA 检查了物理一致性。它让 CO 分子离开 Pt(111) 表面，观察能量曲线；又在 192 原子水盒做 30 皮秒 NVE 模拟。TF32x3 接近 FP32 的平滑曲线和有界能量波动，一次 TF32 则出现较明显的抖动或漂移。对 539 原子合金板的三次结构松弛实验，FP32 用 218 步收敛，TF32x3 用 218–222 步达到 1.0×10⁻³ eV/Å 的力阈值；一次 TF32 的三次实验在 500 步内均未收敛，即使上限增至 1,000 步仍未达标。这些结果只覆盖所测材料和设定。

工程决策因此不能只看吞吐量。大系统若大部分时间花在 GEMM，收益可能较明显；小系统的邻居处理、图操作与启动开销会压低端到端增益。选择 BF16x9 或 TF32x3 后，应在自己的模型上比对能量、力、长轨迹与收敛，不把单点误差当成完整物理验证。上述倍数均为 NVIDIA 在其硬件与 PET 负载下的报告，不是本站复测；正文保留模型、GPU、系统规模和 FP32 基线，供读者判断可迁移性。

比较数值之前，须固定“快”的分母。矩阵乘核耗时只覆盖模型计算的一部分；完整分子动力学步骤还包括邻居列表、数据搬运、图操作和积分。若在大型系统中 GEMM 占比高，核收益更容易传导到整步；小系统则可能由其他环节主导。NVIDIA 给出的 B300 与 A100 结果也不能合并成单一倍数，因为两种硬件、Tensor Core 精度模式和实现方式不同。

精度验证同样需要分层。一次前向的能量误差小，不保证力的导数稳定；单步力相近，不保证多步轨迹或结构松弛仍收敛。厂商因此同时展示了 CO/Pt(111) 能量曲线、30 皮秒水盒和合金板松弛。特别是一次 TF32 未在 500 步收敛，增加到 1,000 步仍失败，说明追求更快的矩阵运算可能改变下游科学结论。TF32x3 在这些例子里接近 FP32，不代表它对所有材料和模型都有同等精度。

团队若复现，应记录 PET 检查点大小、原子数、GPU 型号、软件版本、基线 FP32 实现和性能计时范围，再保存能量、力及收敛阈值。先在自己的任务上设定可接受的物理误差，再比较吞吐量，避免反过来为某一加速数字选择宽松的误差条件。若自己的结果不一致，应按相同口径报告差异，而不是使用厂商表格替代本地测量。本站的数字归因于 NVIDIA 原公告。

如果报告自己的结果，应将失败或不收敛的试验也保留在样本中，说明每种精度路径的重复次数和停止条件。只挑选成功轨迹会高估精度与速度的兼容性。材料模拟的目标是可信的物理结论，性能指标必须服从已事先确定的误差界限，而不能在看到结果后再改标准。