新媒网 · 跨境数据社区

各位跨境的实战精英们，大家好！

咱们今天聊点硬核的，如何在高性能计算的浪潮中，让咱们的分子动力学（MD）模拟跑得更快、更准、更远。分子动力学模拟，这可是计算化学和材料科学领域的一把“瑞士军刀”，无论是研究化学反应机理、材料的微观性质，还是探索复杂的生物相互作用，都离不开它。但大家也知道，这玩意儿计算量巨大，想要处理超大规模的系统，光靠蛮力可不行，得用巧劲，比如引入机器学习势函数（MLIPs），才能实现效率和精度的双丰收。

新媒网跨境获悉，最近国际上的一项重要进展，就是将基于PyTorch的MLIPs模型与LAMMPS这款主流的MD模拟软件巧妙结合，通过一个名为ML-IAP-Kokkos的接口，让原子系统的模拟变得既快速又可扩展。这个接口可不是凭空出现的，它凝聚了像美国英伟达（NVIDIA）这样的科技巨头、美国洛斯阿拉莫斯国家实验室以及美国桑迪亚国家实验室等多方顶尖科学家们的智慧结晶。它的诞生，就是为了简化外部开发者将社区模型与LAMMPS对接的流程，让MLIPs能更顺畅地融入LAMMPS，从而推动更大规模的MD模拟。

ML-IAP-Kokkos接口特别支持消息传递型的MLIP模型，它巧妙地利用了LAMMPS内置的通信机制，确保了多颗GPU之间数据传输的高效性。这对于咱们进行大规模模拟、充分榨取多GPU的计算潜力来说，简直是雪中送炭。同时，这个接口还通过Cython技术，搭建起了Python和C++/Kokkos LAMMPS之间的桥梁，实现了端到端的GPU加速，让整个模拟工作流都得到了优化提升。

接下来，我将以实战专家的视角，一步步带大家，手把手地教你如何将自己的PyTorch模型无缝连接到LAMMPS，实现可扩展的模拟。准备好了吗？咱们这就开始！

磨刀不误砍柴工：准备你的工具与环境

首先，咱们得把“装备”都备齐。如果你有LAMMPS或其他MD模拟工具的使用经验，那最好不过。另外，熟悉Python语言，对PyTorch和机器学习模型有所了解是基础，如果还能掌握一些Cython那就更上一个台阶了。

实战演练：构建你的专属ML-IAP-Kokkos接口

第一步：环境搭建

工欲善其事，必先利其器。确保你的系统上已经安装好了所有必要的软件：

• LAMMPS：编译时需要包含Kokkos、MPI和ML-IAP模块。
• Python环境：并且安装了PyTorch。
• 训练好的PyTorch MLIP模型：如果模型支持cuEquivariance那就更棒了。

第二步：安装支持ML-IAP-Kokkos/Python的LAMMPS

这部分是基础，确保你的LAMMPS版本和配置能够正确识别并调用ML-IAP-Kokkos接口。

第三步：开发ML-IAP-Kokkos接口

为了让大家更好地理解LAMMPS是如何将数据传递给咱们的模型，我们先从一个简单的例子入手。我们创建一个类，它会打印出LAMMPS传递过来的一些关键信息：

from lammps.mliap.mliap_unified_abc import MLIAPUnified
import torch
class MLIAPMod(MLIAPUnified):
    def __init__(
        self,
        element_types = None,
    ):
        super().__init__()
        self.ndescriptors = 1
        self.element_types = element_types
        self.rcutfac = 1.0 # Half of radial cutoff
        self.nparams = 1
    def compute_forces(self, data):
        print(f"Total atoms: {data.ntotal}, Local atoms: {data.nlocal}")
        print(f"Atom indices: {data.iatoms}, Atom types: {data.elems}")
        print(f"Neighbor pairs: {data.npairs}")
        print(f"Pair indices and displacement vectors: ")
        print("\n".join([f" ({i}, {j}), {r}" for i,j,r in zip(data.pair_i, data.pair_j, data.rij)]))

然后，我们将这个模型对象用PyTorch的保存函数存起来：

mymodel = MLIAPMod(["H", "C", "O"])
torch.save(mymodel, "my_model.pt")

接下来，我们准备一个LAMMPS的输入文件sample.in：

2 sample.in
units metal
atom_style atomic
atom_modify map array
boundary p p p

read_data sample.pos

# Load the ML-IAP model
pair_style mliap unified my_model.pt 0
pair_coeff * * C O

dump        dump0 all custom 1 output.xyz id type x y z fx fy fz
run 0

这个示例中，我们模拟了一个简单的二氧化碳分子。它由一个碳原子和两个氧原子组成，沿X轴线性排列。碳原子位于x=1.0 Å处，两个氧原子分别位于x=-0.1 Å和x=2.0 Å处。这样的最小构型，能让我们清晰地观察到LAMMPS是如何处理局部原子、幽灵原子（ghost atoms）、对相互作用以及周期性边界条件的，这些都是理解接口数据流动的关键。下图展示了原子位置及其与模拟单元边界的关系。

我们还需要明确我们感兴趣的元素子集（C和O），以及它们与位置文件中的索引映射关系。

在命令行中，咱们用Kokkos在单个GPU上运行LAMMPS：

lmp -k on g 1 -sf kk -pk kokkos newton on neigh half -in sample.in

运行后，我们来仔细看看输出中与咱们的Python类相关的那部分：

Total atoms: 6, Local atoms: 3
Atom indices: [0 1 2], Atom types: [2 1 1 2 1 1]
Neighbor pairs: 4
Pair indices and displacement vectors:
 (0, 5), [-1.1 0. 0. ]
 (0, 1), [1. 0. 0.]
 (1, 0), [-1. 0. 0.]
 (2, 3), [1.1 0. 0. ]

对照咱们的MLIAPMod类，大家应该能更好地理解这里的数据结构了：

• nlocal代表局部原子数量，ntotal则是总原子数量。
• iatoms记录了原子索引，elems则对应原子类型。
• npairs告诉我们有多少对相互作用。
• 而pair_i、pair_j以及rij则详细列出了每一对原子的索引以及它们之间的位移矢量。

接着，咱们来构思一个简单的势函数。想象一下，一个模型可能在局部原子上计算一些“节点特征”，然后根据邻居特征和键长来更新能量，这是MLIPs中常见的设定。为了简化，我们假设所有特征都为1，径向函数就是原子间距（在实际应用中，这些都会是更复杂的、由模型计算出的对象）。这个势函数的部分实现可能会是这样：

def compute_forces(self, data):
    rij = torch.as_tensor(data.rij)
    rij.requires_grad_()
    features = torch.zeros((data.ntotal,1)).cuda()
    for i in range(data.nlocal):
        features[i,0] = 1.0
    Ei = torch.zeros(data.nlocal).cuda()

    # Message passing step (j → i)
    for r,i,j in zip(rij, torch.as_tensor(data.pair_i), torch.as_tensor(data.pair_j)):
        Ei[i] += features[j,0] * torch.norm(r)

    # Compute total energy and forces
    Etotal = torch.sum(Ei)
    Fij = torch.autograd.grad(Etotal, rij)[0]

    # Copying total energy and forces
    data.energy = Etotal
    data.update_pair_forces_gpu(Fij)

这段代码里有几个小细节，咱们得留意一下：

• 所有在ML-IAP和LAMMPS之间传递的数据，默认都是fp64精度，并且直接位于设备（GPU）上。当然，允许在模型内部进行类型转换以及主机与设备之间的数据传输。
• 总能量的更新是直接在咱们的模型中完成的。
• 力的计算是通过autograd函数对rij进行求导来实现的。这意味着力是“按键”计算的，而不是“按原子”计算。那些直接计算原子力的势函数，需要咱们将其重新调整为这种形式。
• update_pair_forces_gpu这个函数，是把计算出来的力回传给LAMMPS的关键。

作为参考，如果咱们手工计算只有两对键、两个距离（1和1.1）的情况，总能量应该是4.2。每对键贡献1的力，那么碳原子会受到大小为2、方向相反的力，而氧原子则不受力。

在咱们这个简化的例子中，可以把幽灵原子（ghosts）的特征初始化为1。但在真实场景下，这可能是一个更复杂的环境函数。这时候，咱们就可以利用ML-IAP-Kokkos内置的消息传递钩子，把幽灵原子的特征更新成它们对应的真实原子的特征值。这样就能巧妙处理周期性边界条件产生的幽灵粒子，以及那些位于不同GPU上的幽灵粒子了。

• forward_exchange ghost physical：用于将物理原子信息传递给幽灵原子。
• reverse_exchange ghost physical：用于在反向传播时，将幽灵原子的梯度信息回传给物理原子。

第四步：实现ML-IAP-Kokkos的消息传递支持

def compute_forces(self, data):
    rij = torch.as_as_tensor(data.rij)
    rij.requires_grad_()
    features = torch.zeros((data.ntotal,1)).cuda()
    for i in range(data.nlocal):
        features[i,0] = 1.0
    Ei = torch.zeros(data.nlocal).cuda()

    # Update of the features
    new_features = torch.empty_like(features)
    data.forward_exchange(features, new_features, 1)
    features = new_features

    # Local message passing step
    for r,i,j in zip(rij, torch.as_as_tensor(data.pair_i), torch.as_as_tensor(data.pair_j)):
        Ei[i] += features[j,0] * torch.norm(r)

    # Compute total energy and forces
    Etotal = torch.sum(Ei)
    Fij = torch.autograd.grad(Etotal, rij)[0]

    # Copying total energy and forces
    data.energy = Etotal
    data.update_pair_forces_gpu(Fij)

再次运行这段代码，咱们就能得到正确的能量和力了。不过，如果是在多颗GPU上进行计算，这种力计算方式可能就不灵了，因为不同设备之间的梯度是无法直接计算的。

为了解决多GPU下的梯度计算问题，我们来看一个实现完整消息传递的简单例子：

import torch
class LAMMPS_MP(torch.autograd.Function):
    @staticmethod
    def forward(ctx, *args):
        feats, data = args  # unpack
        ctx.vec_len = feats.shape[-1]
        ctx.data = data
        out = torch.empty_like(feats)
        data.forward_exchange(feats, out, ctx.vec_len)
        return out

    @staticmethod
    def backward(ctx, *grad_outputs):
        (grad,) = grad_outputs  # unpack
        gout = torch.empty_like(grad)
        ctx.data.reverse_exchange(grad, gout, ctx.vec_len)
        return gout, None

有了这个新类，咱们插件的最终版本就变成了这样：

def compute_forces(self, data):
    rij = torch.as_as_tensor(data.rij)
    rij.requires_grad_()
    features = torch.zeros((data.ntotal,1)).cuda()
    for i in range(data.nlocal):
        features[i,0] = 1.0
    Ei = torch.zeros(data.nlocal).cuda()

    # Update of the features
    features = LAMMPS_MP().apply(features, data)

    # Local message passing step
    for r,i,j in zip(rij, torch.as_as_tensor(data.pair_i), torch.as_as_tensor(data.pair_j)):
        Ei[i] += features[j,0] * torch.norm(r)

    # Compute total energy and forces
    Etotal = torch.sum(Ei)
    Fij = torch.autograd.grad(Etotal, rij)[0]

    # Copying total energy and forces
    data.energy = Etotal
    data.update_pair_forces_gpu(Fij)

现在，咱们就拥有了能够正确支持自动微分的多GPU消息传递功能了。这种结构对于消息传递型的MLIP模型来说，理论上都应该适用。在任何消息传递步骤之前调用LAMMPS_MP，能够确保特征数据始终保持同步，从而保证模型工作的正确性。当然，有些情况下咱们可以进行优化，跳过一些通信步骤。由于LAMMPS在调用ML-IAP之前会先收集幽灵原子，并且之后会将幽灵原子上的力传播到其他进程，因此在多层MLIP中，第一层的前向交换和最后一层的反向交换通常是可以跳过的。

前面提到，力的计算需要对键长进行能量梯度求导。此外，将所有项都传递给LAMMPS，还能让LAMMPS汇总计算应力。这意味着应力计算无需在咱们的插件中完成。

下图展示了通信钩子带来的性能提升。上方图表是每颗GPU每秒原子步数随GPU数量的变化。下方图表则清晰展示了使用通信钩子后带来的加速比。

上方图表中展示的通信钩子带来的加速，直接得益于幽灵原子数量的减少。当幽灵原子层数为1、2、3、4层时，真实原子占总原子数（真实原子加幽灵原子）的比例分别为54%、38%、26%和18%。而如果启用了通信钩子，我们就可以只使用一层半径的幽灵原子，这样在所有情况下真实原子比例都能保持在54%。这使得在两层、三层和四层幽灵原子的情况下，总原子数分别减少了1.4倍、2.1倍和3倍（对于一层幽灵原子，钩子不发挥作用），这与我们观察到的加速效果是完全吻合的。

上图展示了对于不同LAMMPS插件，模拟一个水盒子系统时，单次MD步长所需时间随系统原子数量的变化。

专家总结与未来展望

LAMMPS的ML-IAP-Kokkos接口，无疑是当前实现多GPU、多节点MD模拟，并有效利用MLIPs的最重要工具之一。它能够以极高的效率模拟超大规模系统，成功地弥合了现代机器学习势场与高性能计算基础设施之间的鸿沟。

通过将基于PyTorch的MLIPs与LAMMPS通过ML-IAP-Kokkos接口集成，开发者们可以轻松实现可扩展且高效的MD模拟。今天的教程详细概述了设置和执行模拟的必要步骤，并强调了LAMMPS中GPU加速和消息传递功能的巨大优势。各项基准测试结果也进一步印证了该接口在大规模模拟方面的巨大潜力，使其成为科研人员和开发者们手中的一把利器。

风险前瞻与合规性提示

各位同行，任何技术都有其两面性，咱们在享受ML-IAP-Kokkos接口带来的高效与便捷时，也需保持清醒，预见潜在的风险。

• 学习曲线与技术门槛： 尽管我们力求深入浅出，但涉及LAMMPS、PyTorch、高性能计算以及分布式编程，对于初学者而言，其学习曲线依然陡峭。需要投入足够的时间和精力去掌握。
• 硬件依赖性： 这种大规模模拟对计算硬件，尤其是GPU数量和性能有较高要求。软硬件投入成本不容忽视，对于资源有限的团队来说，可能需要谨慎评估。
• 模型训练与数据质量： MLIP模型的精度直接决定了模拟结果的可靠性。模型的训练需要大量高质量的量子力学计算数据，数据的获取和处理是一项挑战。
• 知识产权与开源合规： 在使用开源库和模型时，务必注意其许可协议，避免潜在的知识产权风险。如果涉及商业应用，更要仔细审查。
• 系统稳定性与调试难度： 跨框架、跨语言的集成（如Python与C++）在分布式环境下，系统稳定性可能面临挑战，调试复杂问题会比单机模拟更困难。

时效性说明

科技发展日新月异，尤其是在高性能计算和人工智能领域。本文所介绍的技术和方法，是基于当前（2025年）最前沿的实践。请大家注意，随着LAMMPS、PyTorch以及GPU硬件和相关库的不断更新迭代，未来可能会出现更优的实现方式或接口。建议大家在实际操作中，始终参考最新的官方文档，并保持对技术趋势的关注，以便及时调整和优化工作流。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/md-sim-mlip-link-30-min-3x-speedup.html