32块GPU下1.5小时训完通用原子势模型,速度提升百倍。
FastCHGNet: Training one Universal Interatomic Potential to 1.5 Hours with 32 GPUs
- 拆解力与应力预测,优化图神经网络读出机制。
- 通过内核融合等技术,训练时间从8.3天缩短至1.53小时。
- 适合需要快速训练高精度材料模拟模型的研究者。
图神经网络通用原子势模型(GNN-UIPs)在材料发现和性质预测中展现出卓越的泛化与迁移能力,可将分子动力学模拟速度提升数个数量级,同时保持量子力学精度,是材料模拟的新范式。以晶体哈密顿图神经网络(CHGNet)为代表,其在MPtrj数据集上预训练,包含能量、力、应力和磁矩信息,是当前最先进的电荷感知分子动力学模型。但传统CHGNet训练耗时长(单块A100 GPU需8.3天),原因有三:(i)多层传播以获取远距离原子信息,(ii)需计算二阶导数完成权重更新,(iii)原实现未充分挖掘计算潜力。本文提出FastCHGNet,通过三项改进:首先设计力/应力读出模块分解预测任务;其次采用内核融合、冗余跳过等大规模优化,充分释放GPU算力;最后扩展支持多卡并行,提出负载均衡策略提升利用率。数值结果表明,FastCHGNet内存占用降低3.59倍,32块GPU上训练时间降至1.53小时,且不损失模型精度。
原文摘要 · Abstract (English)
Graph neural network universal interatomic potentials (GNN-UIPs) have demonstrated remarkable generalization and transfer capabilities in material discovery and property prediction. These models can accelerate molecular dynamics (MD) simulation by several orders of magnitude while maintaining \textit{ab initio} accuracy, making them a promising new paradigm in material simulations. One notable example is Crystal Hamiltonian Graph Neural Network (CHGNet), pretrained on the energies, forces, stresses, and magnetic moments from the MPtrj dataset, representing a state-of-the-art GNN-UIP model for charge-informed MD simulations. However, training the CHGNet model is time-consuming(8.3 days on one A100 GPU) for three reasons: (i) requiring multi-layer propagation to reach more distant atom information, (ii) requiring second-order derivatives calculation to finish weights updating and (iii) the implementation of reference CHGNet does not fully leverage the computational capabilities. This paper introduces FastCHGNet, an optimized CHGNet, with three contributions: Firstly, we design innovative Force/Stress Readout modules to decompose Force/Stress prediction. Secondly, we adopt massive optimizations such as kernel fusion, redundancy bypass, etc, to exploit GPU computation power sufficiently. Finally, we extend CHGNet to support multiple GPUs and propose a load-balancing technique to enhance GPU utilization. Numerical results show that FastCHGNet reduces memory footprint by a factor of 3.59. The final training time of FastCHGNet can be decreased to \textbf{1.53 hours} on 32 GPUs without sacrificing model accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。