用低精度计算让分子力场模型MACE提速4倍,同时保持物理准确性。
Speeding Up MACE: Low-Precision Tricks for Equivarient Force Fields
- 仅将线性层转为半精度,模型推理速度提升4倍
- 在水分子模拟中,能量与热力学量波动在可接受范围内
- 推荐使用cuEquivariance+FP32主干+线性层半精度,兼顾速度与精度
机器学习力场能以高精度实现分子动力学模拟,但计算成本高昂。针对SO(3)等变模型MACE,本文系统评估了低精度算术与GPU优化内核对性能的影响。通过端到端及模块级分析,对比e3nn与NVIDIA cuEquivariance后端,并测试FP64/FP32/BF16/FP16(FP32累加)在推理、短时NVT和长时NPT水模拟及训练任务中的表现。结果表明,cuEquivariance使推理延迟降低约3倍;仅将线性层设为BF16/FP16可在不改变主干精度下再提速约4倍,且NVT/NPT模拟的能量与热力学可观测量仍在重复实验波动范围内。训练时使用半精度权重会导致力的均方根误差上升。混合e3nn与cuEq模块而无显式适配器会引发表示不匹配。融合等变内核与混合精度推理可显著加速当前最优力场,对下游分子动力学影响极小。建议默认使用cuEquivariance + FP32,对线性层启用BF16/FP16(保留FP32累加)以达最大吞吐量,训练仍维持在FP32。Ampere/Hopper架构上预期进一步提升(支持TF32/BF16),未来可通过内核级半精度路径与流水线融合持续优化。
原文摘要 · Abstract (English)
Machine-learning force fields can deliver accurate molecular dynamics (MD) at high computational cost. For SO(3)-equivariant models such as MACE, there is little systematic evidence on whether reduced-precision arithmetic and GPU-optimized kernels can cut this cost without harming physical fidelity. This thesis aims to make MACE cheaper and faster while preserving accuracy by identifying computational bottlenecks and evaluating low-precision execution policies. We profile MACE end-to-end and per block, compare the e3nn and NVIDIA cuEquivariance backends, and assess FP64/FP32/BF16/FP16 settings (with FP32 accumulation) for inference, short NVT and long NPT water simulations, and toy training runs under reproducible, steady-state timing. cuEquivariance reduces inference latency by about $3\times$. Casting only linear layers to BF16/FP16 within an FP32 model yields roughly 4x additional speedups, while energies and thermodynamic observables in NVT/NPT MD remain within run-to-run variability. Half-precision weights during training degrade force RMSE. Mixing e3nn and cuEq modules without explicit adapters causes representation mismatches. Fused equivariant kernels and mixed-precision inference can substantially accelerate state-of-the-art force fields with negligible impact on downstream MD. A practical policy is to use cuEquivariance with FP32 by default and enable BF16/FP16 for linear layers (keeping FP32 accumulations) for maximum throughput, while training remains in FP32. Further gains are expected on Ampere/Hopper GPUs (TF32/BF16) and from kernel-level FP16/BF16 paths and pipeline fusion.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。