面向大规模分子动力学的高效消息传递神经网络势函数并行化

《JACS Au》：Efficient Parallelization of Message Passing Neural Network Potentials for Large-Scale Molecular Dynamics

　　

机器学习势（Machine Learning Potentials, MLPs）在加速原子尺度模拟方面取得了巨大成功，其中消息传递神经网络（Message Passing Neural Networks, MPNNs）因其卓越的精度而日益普及。然而，由于MPNN势的半局部架构特性，其高效并行化面临挑战。本研究开发了一种高效的MPNN势并行化方案，该方案仅在每层中对局部原子进行附加数据通信且避免冗余计算，从而实现与消息传递深度呈线性关系的可扩展性。该方案首先在多个块体体系上进行测试，包括银、液态水和高熵合金，展示了高达超过1亿原子的优异强可扩展性和弱可扩展性性能，将MPNN的适用性扩展至前所未有的规模。此外，研究人员开发了一个基于C、H、O、N综合数据集的通用势函数，并利用此并行算法高效执行了乙炔-氧气爆轰制备石墨烯的反应性分子动力学模拟，这是碳纳米材料合成中的核心过程。模拟以原子分辨率揭示了氧如何主动调控反应网络，并确定了最有利于石墨烯前驱体形成的O₂/C₂H₂比例，为实验观察到氧对控制石墨烯质量和产量的必要性提供了机制性见解。所提出的并行化框架可 readily 扩展至其他MPNN势函数。

**研究背景与问题**

近年来，机器学习方法在基于从头算计算表示复杂体系势能面（Potential Energy Surfaces, PESs）方面经历了革命性发展。机器学习势（MLPs）比经验力场具有更高的灵活性，同时计算成本远低于即时从头算分子动力学（Ab Initio Molecular Dynamics, AIMD）模拟。其中一类特别成功的MLPs基于原子wise的势能面表示，系统总能量表示为各原子能量之和，每个原子能量由截断半径（r_c）内原子中心环境的多体特征决定。这类原子wise MLPs随系统规模呈线性扩展，确保了对扩展体系的高可扩展性，其局部特征本质也使其易于并行化。

更近期，消息传递神经网络（MPNN）方法在学习势能面方面日益流行。在此框架中，系统表示为由节点（原子）和边（节点间连接）构成的原子图，原子特征通过边特征从邻近节点迭代传播几何信息形成。经过多次消息传递（Message Passing, MP）迭代（T，通常为2 ≤ T ≤ 6）后，MPNN能有效学习局部高阶多体相互作用及超出原始r_c的某些半局部信息，因此通常比传统局部MLP模型实现更优的精度。MPNN势已在包括数万至数十万原子的蛋白质尺度体系等近期应用中取得成功。

尽管具有上述优势，MPNN模型在并行可扩展性方面仍面临固有限制，这阻碍了其应用于涉及数百万原子及更大规模的模拟。这是因为消息通信发生在T个顺序步骤中，有效将邻域半径扩展至Tr_c。因此，该有效T依赖截断r_c,T内的所有原子都对中心原子的最终状态有贡献。通常认为MPNN的并行计算成本随r_c,T或等效地随T不利扩展，因为相互作用原子数呈立方增长。为规避此困难，Kozinsky等提出了Allegro——一种严格局部等变架构，将信息交换限制在原始截断半径内；Csányi等的MACE模型则通过利用高阶等变消息将T最大降至2来限制并行成本。Allegro和MACE的并行化方案主要依赖传统域分解，这对具有半局部消息传递的MPNN架构并非最优。Park等基于NequIP架构提出了SevenNet空间分解算法，将处理器间通信范围限制在r_c内以减少冗余计算，但该限制使其无法使用额外的表皮距离来降低邻居列表更新频率，可能降低高迁移率体系中MD模拟的效率。

**研究开展与核心发现**

本研究提出了一种概念新颖、通用且高效的MPNN模型并行化方案。核心概念是附加通信仅在每层MP中的局部邻近原子间发生，实现与T的线性扩展而非立方扩展。该方案通过简单且物理启发的MPNN模型——递归嵌入原子神经网络（Recursively Embedded Atom Neural Network, REANN）方法实现。

该并行算法的关键在于：邻居列表仅在第一层计算一次，由于构型在消息传递期间不变，计算后可跨后续MP层复用而无需修改，无需随着T增加维护扩展邻居列表。因此，MPNN模型的整体计算成本和通信开销可有利地按~O(T)扩展，这是MPNN模型高效并行化的关键前提。

在技术实现上，针对第t层MP的前向评估，各进程中类原子（Ghost Atoms, GAs）的消息从其被归类为局部原子（Local Atoms, LAs）的进程同步，且其特征正确计算。同时，每层中GAs的特征存储于相关进程以保持计算子图完整，实现反向评估中的数据通信和梯度计算。原子力基于张量微分链式法则计算，仅需同步与邻近进程相关的梯度分量，无需同步整个梯度张量或相关计算子图。实现中进一步采用向量-雅可比积（Vector-Jacobian Product）而非显式计算完整雅可比矩阵，降低梯度的计算和通信开销。该并行算法基于MPI实现，称为REANN-MPI。

数值验证在多种周期性体系上进行，包括银（Ag）、液态水（H₂O）和高熵合金（High-Entropy Alloys, HEAs）。结果表明，REANN-MPI的计算成本随MP层数呈线性扩展，推翻了此前关于MPNN架构~O(T³)（考虑累积成本则为~O(T⁴)）不利扩展的论断。

强可扩展性测试显示，对于192万原子的液态水体系，REANN-MPI在500核时保持约90%效率，1000核时约80%。弱可扩展性测试中，随节点数从1增至32（每节点约100万至300万原子），实际速度几乎保持恒定。对于超过1亿原子的Ag体系，速度达7.4 × 10^-8 s/step/atom，显著超越此前半局部MPNN架构达到的最大规模。

GPU性能测试中，REANN-MPI在单GPU上比单CPU核心快约2个数量级。与严格局部的Allegro模型相比，REANN-MPI速度高出约2个数量级。在更具挑战性的五元过渡金属HEA体系（576原子/GPU）中，REANN-MPI比SevenNet快2.6-4.7倍，比MACE模型快14.8-19.3倍，并行效率从2 GPU的0.94降至32 GPU的0.69，优于其他模型。

**关键技术方法**

研究人员采用的核心技术方法包括：（1）基于REANN架构的MPI并行算法设计，核心为每层仅局部原子间最小化数据通信的前向-反向评估方案；（2）向量-雅可比积技术降低梯度计算和通信开销；（3）分别在CPU集群（Intel Xeon 9242，每节点96核）和GPU集群（NVIDIA A100/H200）上进行强可扩展性和弱可扩展性测试；（4）基于C、H、O、N综合数据集开发通用REANN势函数，该数据集包含26,560个构型，由主动学习过程生成，能量和力通过DFT（BLYP泛函、TZV2P基组、GTH赝势及D3色散校正）计算；（5）将并行算法应用于乙炔-氧气爆轰制备石墨烯的大规模反应性MD模拟，系统规模达27,000原子。

**研究结果详述**

**线性扩展与消息传递验证**：通过Ag和液态水体系的广泛MD模拟，证明REANN-MPI的相对成本随T线性增加，无论系统规模或组成如何。对于1000原子以下的体系，通信开销占主导导致轻微偏离理想关系。

**强可扩展性**：液态水（192万原子）测试中，REANN-MPI与REANN-JIT（原生即时编译实现）在T=1时性能相当，均实现近线性加速至1000 CPU核。当T>1时，REANN-JIT仅支持单节点多线程计算，多核效率极低；而REANN-MPI保持强可扩展性和高效率。两种实现计算的能量和力完全相同，径向分布函数与实验吻合。

**弱可扩展性**：固体Ag和液态水体系中，每节点原子数固定（约100-300万），随节点数增至32，实际速度几乎恒定。绝对速度很快，如超过1亿Ag原子使用32节点时达7.4 × 10^-8 s/step/atom。相同条件下比深度势能（Deep Potential, DP）模型更快。

**GPU性能**：REANN-MPI单GPU速度（5.2 × 10^-1 steps/s）显著快于单CPU核心（1.1 × 10^-2 steps/s）。与Allegro相比快约2个数量级，两者均显示至8 GPU的线性加速。HEA体系（576原子/GPU）中，REANN-MPI比SevenNet快2.6-4.7倍，比MACE快14.8-19.3倍。REANN-MPI并行效率最高，从2 GPU的0.94降至32 GPU的0.69；SevenNet降至0.38；MACE-p因单GPU基线性能人为膨胀表现效率，但相对于原生串行基线效率降至0.34。

**石墨烯形成反应性模拟应用**：乙炔与氧气爆轰制备石墨烯是重要的工业化高产量路线。研究人员训练的高度并行化通用REANN模型在变化的O₂/C₂H₂比条件下进行了广泛MD模拟。

模拟揭示了关键过程：碰撞诱导的C₂H₂解离产生氢原子、乙炔基等自由基；随后形成中间体长碳链，部分分支化；最终形成由多个稠合碳环组成的石墨烯前驱体。不同O₂/C₂H₂比下，低氧浓度时六元环为主导产物，随后五元和七元环；随氧浓度增加，所有碳环形成数量减少。

尺寸效应研究显示：小周期胞（1 K原子）中，最大石墨烯前驱体尺寸随O₂/C₂H₂比单调降低；当模拟胞尺寸加倍（8 K原子）或三倍（27 K原子）时，最大前驱体尺寸统一增加并在O₂/C₂H₂=0.08处达到峰值，之后随氧浓度升高而降低。即使O₂/C₂H₂=0.86，仍可获得约20个碳原子的最大石墨烯前驱体尺寸，与此前ANI-1xnr和DFTB的小胞模拟结果形成对比，更符合实验观察到的高至0.8 O₂/C₂H₂比下石墨烯纳米片形成的实验现象。T=3的深层半局部MPNN模型比T=1的局部模型预测更低的石墨烯环产率，可能与长碳链中间体的半局部相互作用相关。

**讨论与结论**

研究人员总结了所提出的高效MPNN势函数大规模并行化方案。该方案的核心概念是在每层MP中最小化局部邻近原子间的数据通信，确保其优越效率。REANN模型上的数值测试证明，该算法使MPNN势函数随MP层数线性扩展，而非此前认为的多项式扩展。REANN-MPI模型展示了随CPU核数和GPU数量的强可扩展性，以及随原子数的优异弱可扩展性，能够在各种基准体系上实现涉及超过1亿原子的快速并行MD模拟。

该研究的意义在于显著扩展了MPNN势函数的适用性。REANN-MPI在石墨烯合成中的成功应用证明，充分描述石墨烯前驱体形成需要多层消息传递和大规模模拟尺寸的结合，凸显了所提出可扩展MPNN并行化方案在模拟真实体系中的优势。该算法不限于特定MP形式，可容易地适配于其他MPNN势函数的并行化而无需改变其原始结构。未来可通过实现GPU-to-GPU直接数据传输进一步加速。研究人员预期该方案将使MPNN模型能够高效模拟此前仅严格局部模型可及的超大复杂体系。