启用可扩展的AI驱动分子动力学模拟

引言

分子动力学 (MD) 模拟是计算化学和材料科学中的强大工具，对于在微观层面研究化学反应、材料特性和生物相互作用至关重要。然而，其复杂性和计算需求通常需要借助机器学习原子间势 (MLIPs) 等先进技术来实现可扩展性、效率和准确性。

通过 ML-IAP-Kokkos 接口将基于 PyTorch 的 MLIPs 集成到 LAMMPS MD 软件包中，可以实现化学和材料科学研究中原子系统的快速、可扩展模拟。

该接口是 NVIDIA、洛斯阿拉莫斯国家实验室 (Los Alamos National Lab) 和桑迪亚国家实验室 (Sandia National Lab) 的科学家们合作开发的成果。它旨在简化社区模型与外部开发者的接口流程，使他们能够无缝地将 MLIPs 与 LAMMPS 连接起来，从而实现可扩展的 MD 模拟。

ML-IAP-Kokkos 接口支持消息传递 MLIP 模型，利用 LAMMPS 内置的通信功能促进 GPU 之间的高效数据传输——这对于利用多个 GPU 的计算能力进行大规模模拟至关重要。该接口使用 Cython 桥接 Python 和 C++/Kokkos LAMMPS，确保端到端的 GPU 加速，优化整体模拟工作流程。

本文将逐步指导您如何连接自己的 PyTorch 模型以进行可扩展的 LAMMPS 模拟。让我们开始吧。

准备工具和环境

在开始之前，请确保您具备以下经验和软件：

• 经验要求
  • 有使用 LAMMPS 或其他 MD 模拟工具的经验。
  • 有使用 Python 的经验，熟悉 PyTorch 和机器学习模型（了解 Cython 更佳）。
  • 可选：为了获得额外加速，需要 cuEquivariance 支持的模型和 NVIDIA GPU。
• 软件要求
  • 构建了 Kokkos、MPI 和 ML-IAP 的 LAMMPS。
  • 包含 PyTorch 的 Python 环境。
  • 训练好的 PyTorch MLIP 模型（可选支持 cuEquivariance）。

构建您自己的 ML-IAP-Kokkos 接口

第 1 步：设置环境

首先，请确保您已安装所有必要的软件。为方便起见，我们准备了一个包含预编译 LAMMPS 的容器。

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

下载并安装 LAMMPS（2025 年 9 月或更新版本），并确保启用了 Kokkos、MPI、ML-IAP 和 Python 支持。安装 LAMMPS 时启用 Python 支持也会在您的 Python 环境中安装 lammps 模块。

第 3 步：为您的首选模型开发 ML-IAP-Kokkos 接口

使用 ML-IAP-Kokkos 接口时，LAMMPS 会在模拟期间调用 Python 解释器，执行您的原生 Python 模型，并提供所有 Python 功能，而无需编译您的模型。这需要一个可以访问您模型类的工作 Python 环境。

要将您的模型与 LAMMPS 连接，您需要实现 LAMMPS 中定义的抽象类 MLIAPUnified（位于 mliap_unified_abc.py 中）。具体来说，您需要实现 compute_forces 函数，使用从 LAMMPS 传递给它的数据来推断成对的力和能量。该类还需要 compute_gradients 和 compute_descriptors 函数，但这些函数可以为空，因为在此设置中不会使用它们。

让我们从一个简单的示例开始，以了解 LAMMPS 传递给您的类的数据。我们创建一个打印一些信息的类：

# 导入必要的库
from lammps.mliap.mliap_unified_abc import MLIAPUnified
import torch

# 定义一个继承自 MLIAPUnified 的新类
class MLIAPMod(MLIAPUnified):
    """一个简单的 MLIAP 插件示例，用于演示数据流。"""
    def __init__(self, element_types=None):
        """初始化函数，设置模型参数。"""
        super().__init__()
        # 描述符数量，这里设为1
        self.ndescriptors = 1
        # 模型能处理的元素类型列表
        self.element_types = element_types
        # 邻居列表截断半径的一半
        self.rcutfac = 1.0
        # 模型参数数量，这里设为1
        self.nparams = 1

    def compute_forces(self, data):
        """计算力的函数，LAMMPS 将在此处传递数据。"""
        # 打印从 LAMMPS 接收到的原子和邻居对信息
        print(f"总原子数: {data.ntotal}, 本地原子数: {data.nlocal}")
        print(f"原子索引: {data.iatoms}, 原子类型: {data.elems}")
        print(f"邻居对数量: {data.npairs}")
        print(f"邻居对索引和位移向量: ")
        # 逐行打印每个邻居对的索引和相对位移
        print("\n".join([f"  ({i}, {j}), {r}" for i, j, r in zip(data.pair_i, data.pair_j, data.rij)]))

请注意，__init__ 函数指定了我们的模型可以处理的元素类型和邻居列表的截断半径（的一半）。这些参数是 LAMMPS 在模型设置中所必需的。

我们使用 PyTorch 的 save 函数保存我们的模型对象：

# 创建 MLIAPMod 类的实例，指定可处理的元素
mymodel = MLIAPMod(["H", "C", "O"])
# 使用 torch.save 将模型对象序列化到文件
torch.save(mymodel, "my_model.pt")

这将创建文件 my_model.pt，这是一个“pickled”的 Python 对象，LAMMPS 可以加载它。在实际场景中，这会包含我们模型的权重，但在这里它只是我们类的序列化。注意：在最新版本的 PyTorch 中，默认情况下禁用加载包含任意 Python 代码的类作为安全保护。设置环境变量 TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 允许加载该类。请仅使用受信任的 *.pt 文件。

为了测试，我们将在文件 sample.pos 中定义一个单个 CO2 分子的系统：

# 第一行为注释，将被跳过！
3 atoms
2 atom types

0.0 100.0 xlo xhi
0.0 100.0 ylo yhi
0.0 100.0 zlo zhi

Masses

1 12.0
2 16.0

Atoms

1 2 1.0 10.0 10.0
2 1 2.0 10.0 10.0
3 1 -0.1 10.0 10.0

该分子由一个碳原子和两个氧原子沿 X 轴线性排列组成，碳位于 x = 1.0 Å，氧原子分别位于 x = -0.1 Å 和 x = 2.0 Å。这个最小配置使我们能够清楚地观察 LAMMPS 如何处理本地原子与幽灵原子、对相互作用以及周期性边界条件——这些是理解接口中数据流的关键要素。

图 1. 本示例中使用的 CO2 分子示意图，突出显示了原子沿 x 轴相对于单元边界的位置。

以下 LAMMPS 输入脚本 sample.in 加载位置文件，加载我们的势，并运行一个 MD 步骤：

units           metal
atom_style      atomic
atom_modify     map array
boundary p p p
read_data   sample.pos

# 加载 ML-IAP 模型
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

使用 pair_style mliap，unified 关键字指示 LAMMPS 使用 Python 接口，后跟要加载的模型文件，最后是值 0，它会影响邻居列表（有关详细信息，请参阅文档）。我们指定我们感兴趣的元素子集（C 和 O）以及它们与位置文件中索引的映射。

我们使用以下命令在单个 GPU 上运行带 Kokkos 的 LAMMPS：

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

查看输出的相关部分：

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. ]

将此与我们的类进行比较，我们理解了数据结构：

• 原子分为两类：物理原子，范围从 [0, nlocal)，需要更新；和幽灵原子，范围从 [nlocal, ntotal)，它们是周期性副本或来自另一个处理器/GPU 的物理原子。
• 有三个本地原子 (nlocal) 和三个幽灵原子，总共六个 (ntotal)。
• iatoms 列出了本地原子的索引（前 nlocal 个）。
• elems 列出了所有原子的原子种类。
• npairs 是截断范围内涉及至少一个本地原子的对的数量。
• 请注意，对 (i, j) 和 (j, i) 是分开报告的：这里我们有四个对，对应于两个 (C, O) 键的两个方向。
• 我们可以通过 pair_i、pair_j 和 rij 访问每对的 (i, j) 索引和相对位移。
• 请注意，大于或等于 nlocal 的索引指的是幽灵原子（例如，3 是 0 的周期性幽灵，5 是 2 的周期性幽灵）。

要实现一个简单的势，想象一个模型，它在本地原子上计算一些“节点特征”，并根据邻居特征和键长更新能量，这是 MLIPs 的典型设置。为简单起见，我们将所有特征设为 1，径向函数设为原子距离（在实际实现中，这些将是我们的模型计算的更复杂的对象）。

这个势的（部分）实现可能是：

def compute_forces(self, data):
    # 将邻居对的相对位移向量转换为 PyTorch 张量
    rij = torch.as_tensor(data.rij)
    # 设置 requires_grad=True 以便进行自动微分
    rij.requires_grad_()
    # 初始化所有原子（包括幽灵原子）的特征张量
    features = torch.zeros((data.ntotal, 1)).cuda()
    # 仅为本地原子设置特征值（这里简化为1.0）
    for i in range(data.nlocal):
        features[i, 0] = 1.0
    # 初始化本地原子的能量张量
    Ei = torch.zeros(data.nlocal).cuda()

    # 消息传递步骤 (j -> i)，遍历所有邻居对
    for r, i, j in zip(rij, torch.as_tensor(data.pair_i), torch.as_tensor(data.pair_j)):
        # 根据邻居j的特征和键长更新原子i的能量
        Ei[i] += features[j, 0] * torch.norm(r)

    # 计算总能量
    Etotal = torch.sum(Ei)
    # 使用自动微分计算每个键上的力
    Fij = torch.autograd.grad(Etotal, rij)[0]

    # 将计算出的总能量和力复制回 LAMMPS
    data.energy = Etotal
    data.update_pair_forces_gpu(Fij)

请注意此代码中的一些细节：

• 在 ML-IAP 和 LAMMPS 之间传递的所有量都是 fp64 并且在设备上，尽管在模型中允许类型转换和主机<->设备移动。
• 总能量在我们的模型中直接更新。
• 我们通过 autograd 函数计算相对于 rij 的力，即力是“按键”而不是“按原子”计算的。直接计算原子力的势必须重铸为此形式。
• 我们通过 update_pair_forces_gpu 函数更新 LAMMPS 中的力，传递形状为 (npairs, 3) 的 fp64 GPU 张量。

我们在我们的小 CO2 分子上测试这段代码。作为参考，我们手动计算只有两个对和两个距离（1 和 1.1）的能量和力，总能量应该是 4.2，每个对贡献 1 的力，导致碳原子上产生大小为 2 的相反方向的力，氧原子上的力为 0。

然而，运行这段代码我们得到的总能量是 2（在标准输出中看到），并且在一个碳原子和氧原子上力为 2（在 output.xyz 中看到）。很明显，问题在于幽灵原子的存在；它们的特征没有在真实原子上更新，因此只有 (0, 1) 键对总能量和力有贡献。

在这个简化的案例中，我们可以将幽灵初始化为 1，但在实际场景中，这可能是环境的一个更复杂的函数。在这种情况下，我们使用 ML-IAP-Kokkos 内置的消息传递钩子来将幽灵特征更新为其相应真实对应物的值。这处理了来自周期性边界条件的幽灵粒子和驻留在完全不同 GPU 上的幽灵粒子。

提供了两个例程来帮助这个初始化过程：forward_exchange，它从它们的物理对应物设置幽灵原子；和 reverse_exchange，它将对幽灵原子的贡献求和回它们的物理对应物。

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

为了解决上述问题，我们需要在计算前同步幽灵原子的特征。这可以通过 forward_exchange 函数实现。

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()

    # 更新特征
    new_features = torch.empty_like(features)
    # 使用 forward_exchange 同步幽灵原子的特征
    data.forward_exchange(features, new_features, 1)
    features = new_features

    # 本地消息传递步骤
    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)

    # 计算总能量和力
    Etotal = torch.sum(Ei)
    Fij = torch.autograd.grad(Etotal, rij)[0]

    # 复制总能量和力
    data.energy = Etotal
    data.update_pair_forces_gpu(Fij)

再次运行此代码，我们现在获得了正确的能量和力。然而，在多个 GPU 上，力计算将无法正常工作，因为梯度不会在不同设备之间计算。

为了解决这个问题，我们需要定义 LAMMPS_MP.forward() 的梯度并将其注册到 PyTorch。forward_exchange 的梯度通过 reverse_exchange 提供。调用 reverse_exchange 后，梯度被限制在真实原子上，并在所有幽灵原子上清零。

这是一个完整的消息传递实现的简单示例：

import torch

class LAMMPS_MP(torch.autograd.Function):
    """一个自定义的 PyTorch Autograd 函数，用于包装 LAMMPS 的消息传递。"""
    @staticmethod
    def forward(ctx, *args):
        """前向传播：将物理原子的特征交换到幽灵原子。"""
        feats, data = args  # 解包输入
        ctx.vec_len = feats.shape[-1] # 保存特征向量的长度
        ctx.data = data # 保存 data 对象以备反向传播使用
        out = torch.empty_like(feats)
        # 调用 LAMMPS 的前向交换函数
        data.forward_exchange(feats, out, ctx.vec_len)
        return out

    @staticmethod
    def backward(ctx, *grad_outputs):
        """反向传播：将幽灵原子上的梯度贡献累加回物理原子。"""
        (grad,) = grad_outputs  # 解包梯度
        gout = torch.empty_like(grad)
        # 调用 LAMMPS 的反向交换函数
        ctx.data.reverse_exchange(grad, gout, ctx.vec_len)
        return gout, None

有了这个新类，我们插件的最终版本变为：

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()

    # 使用自定义的 Autograd 函数更新特征
    features = LAMMPS_MP.apply(features, data)

    # 本地消息传递步骤
    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)

    # 计算总能量和力
    Etotal = torch.sum(Ei)
    Fij = torch.autograd.grad(Etotal, rij)[0]

    # 复制总能量和力
    data.energy = Etotal
    data.update_pair_forces_gpu(Fij)

我们现在可以访问支持 autograd 的正确的多 GPU 消息传递。这个相同的结构应该独立于消息传递 MLIP 模型工作。

在任何消息传递之前调用 LAMMPS_MP 可确保特征保持同步并且模型正常工作。然而，可能存在优化，其中可以跳过某些通信步骤。由于 LAMMPS 在调用 ML-IAP 之前收集幽灵，并且之后将幽灵上的力传播到其他等级，因此多层 MLIP 中的第一个前向交换和最后一个反向交换通常可以跳过。

如前所述，力计算需要计算能量相对于键的梯度。此外，将所有项传递给 LAMMPS 允许 LAMMPS 汇总贡献以计算应力；插件不需要计算应力。

使用 HIPPYNN 模型评估性能

为了展示消息传递钩子带来的性能提升，我们在 1 到 512 个 NVIDIA H100 GPU 上对运行 HIPPYNN (Chigaev 2023) 模型的 LAMMPS 进行了基准测试，模型由 MLIP 中 h=1 到 h=4 个交互层组成，分别测试了有和没有通信钩子的情况。我们使用一个简单的铝原子晶格作为性能测试，每个 GPU 保持大约 203 个原子。

模型在论文 Automated discovery of a robust interatomic potential for aluminum 的数据集上进行训练，遵循 HIPPYNN 中的多交互层训练脚本示例。为了模拟的正确性，其中 r 是一个 MLIAP 层的交互半径，我们在启用钩子时使用长度为 r+epsilon 的幽灵截断，在禁用钩子时在邻居列表中启用幽灵-幽灵对，使用 hr+epsilon。

图 2. 顶部图表显示了每 GPU 每秒原子步数与 GPU 数量的关系。底部是使用通信钩子带来的加速比图。

上图所示的使用通信钩子带来的加速直接归因于幽灵原子的减少。对于一层、两层、三层和四层幽灵原子，真实原子占总原子（真实加幽灵）的百分比分别为 54%、38%、26% 和 18%。启用钩子后，我们可以对所有情况使用一个半径的幽灵，因此真实原子占 54%。这导致两层、三层和四层的总原子数分别减少 1.4 倍、2.1 倍和 3 倍（对于一层，钩子不起作用），直接对应于观察到的加速。

使用 MACE 集成比较结果

作为进一步的演示，我们使用流行的 MLIP MACE 对 LAMMPS 模拟进行了基准测试，比较了 LAMMPS 中原始的 MACE 配对样式（利用 C++ 中的 libtorch 加载模型权重）与使用 ML-IAP-Kokkos 接口的基于 cuEquivariance 的 MACE 插件。cuEquivariance 是一个用于加速几何神经网络的 NVIDIA Python 库。

图 3. 不同 LAMMPS 插件的模拟时间随系统大小变化的图表。

对于此基准测试，我们使用 MACE-OFF23 势的“中等”版本，在单个具有 80GB 内存的 A100 GPU 上模拟越来越大的水盒子，直到 GPU 内存耗尽，模型精度为 fp32 和 fp64。图中显示了原始 MACE 配对样式的速度以供比较。新插件速度更快，内存效率更高。这种加速归因于使用 cuEquivariance 的模型加速和 ML-IAP-Kokkos 接口中的消息传递改进。通过 ML-IAP-Kokkos 接口，cuEquivariance 加速的 MLIPs 可以无缝地在 LAMMPS 中使用。

结论

LAMMPS ML-IAP-Kokkos 接口是目前使用 MLIPs 进行多 GPU、多节点 MD 模拟的最重要工具。它能够高效地模拟极大系统，弥合了现代基于 ML 的力场与高性能计算基础设施之间的差距。

通过 ML-IAP-Kokkos 接口将基于 PyTorch 的 MLIPs 与 LAMMPS 集成，开发人员可以实现可扩展和高效的 MD 模拟。本教程概述了设置和执行模拟所需的步骤，并强调了 LAMMPS 中 GPU 加速和消息传递功能的优势。基准测试结果进一步强调了该接口在大规模模拟中的潜力，使其成为该领域研究人员和开发人员的宝贵工具。

准备好加速您的模拟了吗？立即尝试本教程，并与社区分享您的成果。