NVIDIA ALCHEMI工具包：加速AI驱动化学与材料科学模拟

用户576110558132
0 阅读11分钟

加速AI驱动的化学与材料科学模拟：NVIDIA ALCHEMI 工具包

机器学习原子间势能（MLIP）正在变革计算化学和材料科学领域。MLIP使得原子模拟能够将计算成本高昂的量子化学的保真度与AI的扩展能力相结合。

然而，在此交叉领域工作的开发者面临一个持续的挑战：缺乏一个强大且Python化的、用于GPU加速原子模拟的工具箱。对于运行大量并行、GPU加速模拟等用例，当前软件生态系统中要么缺少稳健且支持良好的工具，要么这些工具分散在多个开源软件中。

过去几年中，用于运行MLIP原子模拟的可用软件一直以CPU为中心。核心操作，如近邻识别、色散校正、长程相互作用及其相关的梯度计算，传统上仅支持CPU计算，这常常难以满足当代研究所需的速度。在混合工作流（模型在PyTorch中进行GPU加速，但模拟工具是串行且基于CPU的）中，中小型原子系统的高通量模拟很快就会因GPU使用效率低下而形成瓶颈。

尽管开发者多年来一直尝试直接在PyTorch中实现这些操作，但PyTorch的通用设计在处理原子模拟所需的专用空间和力计算操作时，性能有所不足。PyTorch能力与原子模拟需求之间的这种根本性不匹配引发了一个重要问题：需要什么来弥合这一差距？

NVIDIA ALCHEMI（化学与材料创新AI实验室）为化学和材料科学领域的开发者和研究人员提供了领域专用的工具包，以及针对NVIDIA加速计算平台优化的NVIDIA NIM微服务。它是一套高性能、批处理且GPU加速的工具集合，专门用于在机器学习框架层面实现化学和材料科学研究中的原子模拟。

NVIDIA ALCHEMI通过三个集成层次提供能力：

  • ALCHEMI 工具包-操作层 (Toolkit-Ops)：一个用于AI赋能的原子模拟任务的GPU加速、批处理通用操作仓库，例如近邻列表构建、DFT-D3色散校正和长程静电相互作用。
  • ALCHEMI 工具包 (Toolkit)：一套GPU加速的模拟构建模块，包括几何优化器、积分器和数据结构，以支持利用AI的大规模、批处理模拟。
  • ALCHEMI NIM 微服务：一个可扩展的、面向化学和材料科学的云就绪领域特定微服务层，支持在NVIDIA加速平台上的部署和编排。

本文介绍NVIDIA ALCHEMI Toolkit-Ops，即ALCHEMI的加速批处理通用操作层。ALCHEMI Toolkit-Ops使用NVIDIA Warp来加速和批处理AI驱动的原子建模中的常见操作。这些操作通过模块化的、PyTorch可访问的API公开，能够与现有及未来的原子模拟包快速迭代和集成。

图1展示了此初始版本ALCHEMI Toolkit-Ops中包含的用于原子模拟的加速批处理通用操作。此测试版包括两种近邻列表（朴素算法和单元格算法）、DFT-D3色散校正以及长程库仑（Ewald和Particle Mesh Ewald）函数。

图1. NVIDIA ALCHEMI Toolkit-Ops是一个专门为支持MLIP和分子动力学引擎而开发的、用于GPU加速批处理操作（单个GPU，多个系统）的模块仓库。

图2展示了ALCHEMI Toolkit-Ops中的加速核函数与流行的核加速模型（如MACE（cuEquivariance）和TensorNet（Warp））相比的性能，以实现完全并行化的性能和可扩展性。蓝色的MLIP基线允许与高级特性（如近邻列表、DFT-D3色散校正和显式静电计算（Ewald和Particle-Mesh Ewald (PME)））进行比较。测试系统由使用Packmol填充到各种晶胞中、尺寸不断增大的氨团簇组成。计时结果在NVIDIA H100 80 GB GPU上运行20次取平均值。由于D3的长程性质，DFT-D3计算不包括6Å的结果。

图2. 基准测试显示了ALCHEMI Toolkit-Ops近邻列表（包括朴素的O(N²)和单元格列表O(N)实现）、DFT-D3校正以及两种版本的静电相互作用的速度。所有方法都与流行的核加速MLIP的计算成本进行了比较。左侧面板概述了在固定原子数和可变系统尺寸 x [批大小] 情况下的批处理扩展性，而右侧面板展示了单个系统尺寸增长时的耗时情况。

ALCHEMI Toolkit-Ops 生态系统集成

ALCHEMI Toolkit-Ops旨在与更广泛的、基于PyTorch的原子模拟生态系统无缝集成。我们很高兴地宣布，正在与化学和材料科学领域的领先开源工具进行集成：TorchSim、MatGL和AIMNet Central。

TorchSim

TorchSim，一个下一代开源原子模拟引擎，正在采用ALCHEMI Toolkit-Ops内核来驱动其GPU加速工作流。TorchSim是一个为MLIP时代专门构建的、PyTorch原生的模拟引擎，能够在单个GPU上同时对数千个系统进行批处理分子动力学和结构弛豫。TorchSim将利用我们优化的近邻列表来驱动高通量批处理操作，同时不牺牲灵活性或性能。

MatGL

MatGL（材料图库）是一个开源框架，用于为无机、分子和混合材料系统构建基于图的机器学习原子间势能和基础势能。通过集成ALCHEMI Toolkit-Ops，MatGL显著加速了基于图的长程相互作用的处理，使得大规模原子模拟不仅速度更快、计算效率更高，而且不损失精度。

AIMNet Central

AIMNet Central是AIMNet2的仓库，AIMNet2是一个通用MLIP，能够高保真地模拟中性、带电、有机和元素有机系统。AIMNet Central正在利用ALCHEMI Toolkit-Ops来进一步提升其灵活的长程相互作用模型的性能。通过使用NVIDIA加速的DFT-D3和近邻列表内核，AIMNet2可以在不牺牲精度的情况下，为大型和周期性系统提供更快的原子模拟。

如何开始使用 ALCHEMI Toolkit-Ops

开始使用ALCHEMI Toolkit-Ops很简单，设计上注重易用性。

系统和包要求

  • Python 3.11+
  • 操作系统：Linux（主要）、Windows (WSL2)、macOS
  • NVIDIA GPU（推荐A100或更新），CUDA计算能力 ≥ 8.0
  • CUDA Toolkit 12+, NVIDIA 驱动 570.xx.xx+

安装

要安装ALCHEMI Toolkit-Ops，请使用以下代码片段：

# 通过pip wheel安装
pip install nvalchemi-toolkit-ops

# 确保它可以被导入
python -c "import nvalchemiops; print(nvalchemiops.__version__)"

有关其他安装说明，请参阅ALCHEMI Toolkit-Ops文档。探索GitHub仓库中的examples目录并在您自己的硬件上运行它们以测试加速效果。

典型的故障排除提示：

  • 验证CUDA安装和设备可用性：nvidia-smi, nvcc --version
  • 确保兼容的Python版本：python --version
  • 根据需要升级依赖项：pip list | grep torch 和 pip list | grep warp

功能亮点

本节深入探讨ALCHEMI Toolkit-Ops的三个初始功能：高性能近邻列表、DFT-D3色散校正和长程静电相互作用。

近邻列表

近邻列表构建是原子模拟的基石，它使得能够使用局域或半局域MLIP计算能量和力。ALCHEMI Toolkit-Ops在PyTorch中提供了最先进的GPU性能，对于许多中小型原子系统或单个大型原子系统的批处理，实现了每秒数百万原子的性能扩展。

能力

  • 支持批处理的O(N)（单元格列表）和O(N²)（朴素）算法
  • 支持任意晶胞尺寸和部分周期性的三斜晶胞的周期边界
  • 支持端到端计算图编译
  • 与PyTorch的直接API兼容性

API示例

import torch
from nvalchemiops.neighborlist import neighbor_list

# 水分子
water_positions = torch.tensor([
   [0.0, 0.0, 0.0],      # O
   [0.96, 0.0, 0.0],     # H
   [-0.24, 0.93, 0.0],   # H
], device="cuda", dtype=torch.float32)
# 氨分子 (NH3)
ammonia_positions = torch.tensor([
   [0.0, 0.0, 0.0],      # N
   [1.01, 0.0, 0.0],     # H
   [-0.34, 0.95, 0.0],   # H
   [-0.34, -0.48, 0.82], # H
], device="cuda", dtype=torch.float32)

# 为批处理拼接位置
positions = torch.cat([water_positions, ammonia_positions], dim=0)

# 创建批次索引（0代表水，1代表氨）
batch_idx = torch.cat([
   torch.zeros(3, dtype=torch.int32, device="cuda"),   # 水
   torch.ones(4, dtype=torch.int32, device="cuda"),    # 氨
])

# 为每个分子定义晶胞（足够大以容纳分子，不使用PBC）
cells = torch.stack([
   torch.eye(3, device="cuda") * 10.0,  # 水晶胞
   torch.eye(3, device="cuda") * 10.0,  # 氨晶胞
])

# 非周期性分子情况
pbc = torch.tensor([
   [False, False, False],  # 水
   [False, False, False],  # 氨
], device="cuda")

# 截断距离，单位为埃
cutoff = 4.0

# 计算近邻列表；这里明确请求批处理单元格列表算法
neighbor_matrix, num_neighbors, shift_matrix = neighbor_list(
   positions, cutoff, cell=cells, pbc=pbc, batch_idx=batch_idx, method="batch_cell_list"
)

print(f"近邻矩阵: {neighbor_matrix.cpu()}")  # [7, num_neighbors.max()]
print(f"每个原子的近邻数: {num_neighbors.cpu()}")  # [7,]
print(f"周期位移: {shift_matrix.cpu()}")

DFT-D3 色散校正

真实的分子建模必须充分考虑范德华相互作用，而标准DFT泛函并未系统地考虑这些。DFT-D3使用经验性的成对校正，从而显著改进了常见DFT泛函的结合能、晶格结构、构象分析和吸附研究。

能力

  • Becke-Johnson (BJ) 有理阻尼变体
  • 支持批处理和周期性计算
  • 支持在截断距离处平滑处理
  • 联合计算能量、力和应力张量

API示例

from nvalchemiops.interactions.dispersion import dftd3

batch_ptr = torch.tensor([0, 3, 7], dtype=torch.int32, device="cuda")
atomic_numbers = torch.tensor(
    [6, 1, 1, 7, 1, 1, 1], dtype=torch.int32, device="cuda"
)

# 对于此代码片段，假设 d3_params 已加载为：
# d3_params = D3Parameters(rcov=..., r4r2=..., c6ab=..., cn_ref=...)
# 用户可以参阅文档以获取DFT-D3参数来源并了解预期的数据结构
d3_params = ...

# 调用 DFT-D3 函数接口
energy, forces, coordination_numbers = dftd3(
    positions=positions,
    numbers=atomic_numbers,
    a1=0.3981, a2=4.4211, s8=0.7875,  # PBE 参数
    neighbor_matrix=neighbor_matrix,
    neighbor_matrix_shifts=shift_matrix,
    batch_idx=batch_idx,
    d3_params=d3_params
)

print(f"能量: {energy.cpu()}")  # [2,]
print(f"力: {forces.cpu()}")  # [7, 3]

局限性 当前的实现仅计算二体项（C6和C8）。不包括三体Axilrod-Teller-Muto (ATM/C9)贡献。这通常会导致对色散能的一定程度高估。

长程静电相互作用

对于涉及离子/带电物质和极性系统的模拟，精确建模静电相互作用至关重要。目前，MLIP最常见的方法是在短程模型内学习库仑相互作用。对长程库仑效应的系统性低估会导致结合能、溶剂化结构和界面现象的精度损失。

ALCHEMI Toolkit-Ops提供了完全GPU加速的Ewald求和方法——包括标准Ewald和particle mesh Ewald——使得能够在PyTorch中进行GPU加速、高效且精确的长程静电处理。

对于大型周期性系统，基于Ewald的方法将静电相互作用分为短程和长程部分，每个部分都在最适合性能的域中计算。ALCHEMI Toolkit-Ops提供了一种双截断策略，与朴素的全部成对方法相比，显著减少了冗余的近邻查询和内存开销，使得在现代GPU上对带电系统进行高通量模拟变得切实可行。用户可以根据特定的性能和精度需求，为较小系统选择标准Ewald，或为大型周期性系统选择PME。

能力

  • Ewald求和方法
  • 使用B样条的 Particle Mesh Ewald (PME)
  • 支持批处理和周期性系统
  • GPU优化计算，利用cuFFT加速倒易空间评估
  • PyTorch集成为端到端可微工作流提供原生张量支持

API示例

from nvalchemiops.interactions.electrostatics import particle_mesh_ewald

# 此处每个原子的电荷是随机生成的
atomic_charges = torch.randn(
    positions.size(0),  dtype=torch.float32, device="cuda"
)

# 使用 particle mesh ewald 计算能量和力
energy, forces = particle_mesh_ewald(
    positions,
    atomic_charges,
    cells,
    alpha=0.3,  # 调整 Ewald 分裂参数
    batch_idx=batch_idx,
    neighbor_matrix=neighbor_matrix,
    neighbor_matrix_shifts=shift_matrix,
    compute_forces=True
)

print(f"能量: {energy.cpu()}")  # [2]
print(f"力: {forces.cpu()}")  # [7, 3]

深入了解 ALCHEMI Toolkit-Ops

ALCHEMI Toolkit-Ops通过提供在NVIDIA GPU上高性能、易访问的原子建模工具来赋能社区。要加速您的化学和材料科学模拟，请访问NVIDIA/nvalchemi-toolkit-ops GitHub仓库和NVIDIA ALCHEMI Toolkit-Ops文档。您还可以浏览示例集。此ALCHEMI Toolkit-Ops测试版专注于高效近邻列表、色散校正和长程静电学。敬请关注未来版本的新功能和性能优化。

致谢

我们感谢Shyue Ping Ong教授、Olexandr Isayev教授以及TorchSim委员会成员Abhijeet Gangan、Orion Archer Cohen、Will Engler和Ben Blaiszik与我们合作，将NVIDIA ALCHEMI Toolkit-Ops集成到他们的开源项目中。我们也感谢来自某机构的Wen Jie Ong、Piero Altoe和Kibibi Moseley对准备本文的帮助。FINISHED

avatar
文章
阅读
粉丝