分布式推理平台让分子模拟快3.4倍,支持超大规模原子系统计算。
DistMLIP: A Distributed Inference Platform for Machine Learning Interatomic Potentials
- 基于图划分实现多设备并行,兼容多种图神经网络结构。
- 在8张GPU上实现近百万原子系统秒级模拟,速度提升最高8倍。
- 可直接接入现有模型,适合材料与药物研发的大规模模拟需求。
大规模原子模拟对连接计算材料与化学到真实材料和药物发现至关重要。近年来,机器学习势(MLIPs)的发展为扩展量子力学计算提供了新路径。将这些势函数在多个设备上并行化是进一步拓展模拟尺度的挑战性但前景广阔的方法。本文提出DistMLIP,一种基于零冗余、图级别并行的高效分布式推理平台。相比传统的空间分区并行,DistMLIP通过图划分实现高效并行,支持如多层图神经网络等灵活的MLIP架构。该平台提供易用、灵活的插件接口,可直接用于已有MLIP模型的分布式推理。我们在四种主流先进模型——CHGNet、MACE、TensorNet和eSEN上验证了其性能。结果显示,与以往多GPU方法相比,DistMLIP能处理3.4倍更大的原子体系,最快可达8倍加速;在8张GPU上,现有基础势函数可在数秒内完成近百万原子的模拟。
原文摘要 · Abstract (English)
Large-scale atomistic simulations are essential to bridge computational materials and chemistry to realistic materials and drug discovery applications. In the past few years, rapid developments of machine learning interatomic potentials (MLIPs) have offered a solution to scale up quantum mechanical calculations. Parallelizing these interatomic potentials across multiple devices poses a challenging, but promising approach to further extending simulation scales to real-world applications. In this work, we present DistMLIP, an efficient distributed inference platform for MLIPs based on zero-redundancy, graph-level parallelization. In contrast to conventional spatial partitioning parallelization, DistMLIP enables efficient MLIP parallelization through graph partitioning, allowing multi-device inference on flexible MLIP model architectures like multi-layer graph neural networks. DistMLIP presents an easy-to-use, flexible, plug-in interface that enables distributed inference of pre-existing MLIPs. We demonstrate DistMLIP on four widely used and state-of-the-art MLIPs: CHGNet, MACE, TensorNet, and eSEN. We show that DistMLIP can simulate atomic systems 3.4x larger and up to 8x faster compared to previous multi-GPU methods. We show that existing foundation potentials can perform near-million-atom calculations at the scale of a few seconds on 8 GPUs with DistMLIP.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。