Related Resources
模拟/建模/设计

使用 NVIDIA ALCHEMI 工具包为化学和材料科学构建自定义原子模拟工作流

2026年 4月 14日
作者：, , , 和

几十年来，计算化学一直面临准确性与速度之间的拔河。密度泛函理论 (DFT) 等 Ab initio 方法提供高保真度，但计算成本高昂，将研究人员限制在只有几百个原子的系统上。相反，传统的力场速度很快，但通常缺乏复杂的键断裂或过渡态分析所需的化学准确性。

机器学习原子间相互作用势 (MLIP) 已成为以经典速度提供量子准确性的桥梁。但是，软件生态系统是一个新的瓶颈。虽然 MLIP 模型本身在 GPU 上运行，但周围的仿真基础设施通常依赖于以 CPU 为中心的传统代码。

NVIDIA ALCHEMI (化学和材料创新 AI 实验室) 通过利用 AI 加速化学和材料发现，帮助应对这些挑战。我们之前宣布了 ALCHEMI 产品组合的两个组成部分：

  • ALCHEMI NIM 微服务: 适用于化学和材料科学领域 AI 加速批量原子模拟的可扩展、云就绪型微服务
  • ALCHEMI 工具包 – Ops：一组基础 GPU 内核，旨在加速模拟背后的计算，例如近邻列表、色散校正和静电

今天，我们将推出 NVIDIA ALCHEMI Toolkit，这是一组 GPU 加速的模拟构建块，整合并扩展了 ALCHEMI Toolkit-Ops。ALCHEMI 工具包旨在管理加速化学和材质领域特定内核与深度学习模型之间的数据流。ALCHEMI Toolkit 不仅能扩展单个模型和内核，还能为研究人员和开发者提供模块化的 PyTorch 原生结构，以构建自定义仿真工作流。

图 1 显示了此初始版本的 ALCHEMI Toolkit 中支持的 ALCHEMI 架构堆栈和产品功能，包括 Toolkit-Ops 中的扩展功能。此版本包括几何松和分子动力学功能，以及用于组合多个仿真工作流的支持工作流基础设施。

Graphic of ALCHEMI architectural stack (left) with list of product features supported in the initial release of ALCHEMI Toolkit (right).
图 1. NVIDIA ALCHEMI 工具套件是 GPU 加速模拟基础模组的集合，可借助 AI 实现大规模批量模拟

ALCHEMI 工具包如何推进数字化学的发展？

ALCHEMI 工具包不仅仅是脚本的集合。它旨在使研究人员和开发者能够轻松构建自定义的高性能原子模拟工作流程。

扩展 ALCHEMI 工具包 – Ops

ALCHEMI Toolkit 利用 Toolkit-Ops 的功能来处理模拟的基础计算。上一个版本包含几个关键运算：

  • 近邻列表结构
  • DFT-D3 色散校正
  • 远距离静电相互作用

此版本扩展了常见操作的范围，包括:

  • 批处理动态内核
  • JAX 支持 (适用于 v0.2.0 版本功能)

与原子模拟生态系统集成

ALCHEMI 工具包旨在与更广泛的原子模拟生态系统无缝集成。我们很高兴地宣布，我们将与化学和材料科学领域的领先平台进行以下集成。

轨道

Orbital 开发了先进的 AI 基础模型，用于加速发现适用于数据中心和可持续材料的新型冷却系统。Orbital 已将 ALCHEMI Toolkit 集成到其新的 OrbMolv2 模型中，以大幅缩短推理所需的时间。新模型将利用 ALCHEMI Toolkit 组件，例如用于定期库伦相互作用的 PME 静电和用于批量恒定压力分子动力学的 MTK 积分器。现有的 Orb 模型已经利用 Toolkit-Ops 进行 GPU 加速的图形构建，在 TorchSim 支持下，可为大型系统提供约 1.7 倍的加速，为批量较小的系统提供约 33 倍的加速。

材质图形库 (MatGL)

MatGL 是一个开源框架，适用于基于图形的先进 MLIP。ALCHEMI 工具包与 MatGL TensorNet 模型集成，可显著加速材料模拟和属性预测工作流程。通过利用 ALCHEMI Toolkit GPU 原生内核和批处理基础设施，MatGL 用户可以在大规模模拟中实现更高的计算效率和更低的内存消耗。

Matlantis

Matlantis 通过将通用 MLIP 与高性能云计算相结合，实现快速的材料发现。Matlantis 正在积极探索 ALCHEMI 工具包，并确定其可组合动力学如何为工业材料模拟客户提供最大价值。这基于其经过验证的 ALCHEMI Toolkit-Ops 集成 (包括 Warp 优化的近邻列表构建和 DFT-D3 色散校正) ，可显著减少原子交互的计算开销，速度提升高达 10 倍。

此外，通过评估 ALCHEMI Toolkit 中的特定组件，这项合作有可能使 Matlantis 能够超越单一结构优化，实现数百万分子配置的高吞吐量、并行松。最终，这种集成旨在进一步推动小规模研究和工业规模的材料设计，以无与伦比的 GPU 效率加速化学评估。

如何开始使用 ALCHEMI 工具包

本节将为您介绍如何开始使用 ALCHEMI 工具包，该工具包简单易用。

系统和软件包要求

  • Python* 3.11、* 3.14
  • PyTorch@ 2.8
  • CUDA 工具包 12+，NVIDIA 驱动 470.57.02+
  • 操作系统：Linux (主要) 、macOS
  • NVIDIA GPU ( RTX 20xx 或更高版本) ，CUDA 计算能力 = 7.0
  • 最低 4 GB RAM (对于大型系统，建议使用 16 GB)

安装

使用以下代码安装 ALCHEMI Toolkit：

# Install Atomic Simulation Environment (ASE, used in the examples below)
uv pip install ase
# Using pip
pip install nvalchemi-toolkit
# Using uv
uv venv --seed --python 3.12
uv pip install nvalchemi-toolkit
# Install from source
git clone https://github.com/NVIDIA/nvalchemi-toolkit.git
cd nvalchemi-toolkit
uv sync --all-extras
# Add nvalchemi as a project dependency
uv add nvalchemi-toolkit

有关更多信息，请参阅 NVIDIA/nvalchemi-toolkit GitHub 存储库和 ALCHEMI 工具包文档。

用于构建端到端工作流的 ALCHEMI Toolkit 的主要特性

本节将深入探讨 ALCHEMI Toolkit 的四个核心功能：可定制的批量仿真工作流、构建您自己的动态类、模型包装器和高级数据管理。这些功能为研究人员和开发者提供了创建定制端到端工作流所需的工具和灵活性，从而更大限度地提高 NVIDIA GPU 的效率和性能。

可定制的批量模拟工作流程

NVIDIA ALCHEMI 工具包的显著特点是 GPU 原生批量动态引擎。没有哪种 MLIP 模型能够完美适用于每种化学环境，尤其是在处理非局部、远程相互作用时。

ALCHEMI 工具包使研究人员能够将模块化化学和材料科学领域特定的内核和模型整合到定制的仿真工作流中。该架构支持开发专门的计算工作流，并运行具有数百万并发原子交互的虚拟实验室，而不会像传统软件堆栈那样延迟。

功能

  • 将 MLIP 与基于物理性质的校正相结合的可组合计算器
  • 高性能封装器 ( MACE、TensorNet、AIMNet2)

API 示例

以下示例构建数据、设置 MLIP 并配置 FIRE2 几何优化，然后将其用作 velocity Verlet (微规范) 动力学的起点：

from ase import Atoms
from nvalchemi.data import AtomicData, AtomicBatch
from nvalchemi.dynamics import ConvergenceHook
from nvalchemi.dynamics.optimizers import FIRE2
from nvalchemi.dynamics.integrator import VelocityVerlet
# setup some batch of atomic structures
atomic_data = [AtomicData.from_atoms(Atoms(...), device="cuda") for _ in range(16)]
batch = Batch.from_data_list(atomic_data)
# setup your MLIP and dynamics classes
mlip = ...
# optimizer convergence depends on the force norm and max values
conv_criteria = ConvergenceHook(
    criteria=[
        {"key": "forces", "threshold": 0.05, "reduce_op": "norm"},
        {"key": "forces", "threshold": 0.1, "reduce_op": "max"}
    ]
)
optimizer = FIRE2(
    mlip,
    convergence_hook=conv_criteria,
    n_steps=200
)
velverlet = VelocityVerlet(mlip, n_steps=1000)

您可以通过以下两种方式之一运行和扩展仿真流程：在单个 GPU 上，或在多个 CPU 和 GPU 上。

在单个 GPU 上运行和扩展工作流：FusedStage 类通过将两个或多个动态对象“添加”在一起而形成。这有助于在 torch.compile 中封装端到端工作流，并共享 CUDA 流上下文。

fused = optimizer + velverlet
# context manager handles compilation and CUDA stream
with fused:
    # runs 200 steps of optimization and 1000 steps of MD
    fused.run(batch)

通过这种方法，您可以轻松构建仿真工作流，在批量内的样本立即收时运行顺序步骤，并充分利用 GPU。

在多个 CPU 和 GPU 上运行和扩展工作流：第二种方法是在多个 CPU/ GPU 上分配工作流。然后，在两个动态类上使用管道运算符将 FIRE2 优化分配到一个 GPU 上，并在另一个 GPU 上分配 velocity Verlet 集成。

pipeline = optimizer | velverlet
# equivalent to manual allocation with explicit producer/consumer
# optimizer.next_rank = 1, velverlet.prior_rank = 0
# DistributedPipeline({0: optimizer, 1: velverlet})
with pipeline:
    pipeline.run(batch)

虽然为了说明目的，我们特意简化了这个示例，但这种抽象允许用户将其工作流扩展到节点上的多个 GPU，并扩展到多个节点，以获得任意规模的数据集和秩数。

以下示例配置了 8 个 GPU 以运行几何优化，并通过管道传输结果在另外 8 个 GPU 上运行 Langevin 动力学：

from torch import distributed as dist
from torch.utils.data.distributed import DistributedSampler
from nvalchemi.data.datapipes import Dataset, DataLoader
# set up distributed; torchrun --nproc-per-node 8 --nnodes 2 ...
dist.initialize_process_group()
# set up data and distributed sampler
dataset = Dataset(...)
data_sampler = DistributedSampler(
   dataset,
   num_replicas=dist.get_world_size(),
   rank=dist.get_rank()
)
loader = DataLoader(
   dataset,
   batch_size=128,
   sampler=sampler,
   use_stream=True
)
# configure your pipeline; 8 ranks do optimization, 8 do langevin dynamics
optimizers = [FIRE2(mlip, ..., next_rank=index + 8) for index in range(8)]
dynamics = [Langevin(mlip, ..., prior_rank=index) for index in range(8)]
pipeline = DistributedPipeline(
    {index: stage for index, stage in enumerate(optimizers + dynamics)}
)
with pipeline:
    for batch in loader:
        pipeline.run(batch)

构建您自己的动态类

ALCHEMI 工具包提供模块化架构，用于从头开始构建和自定义动态类。这种方法使社区能够将新的采样方法或热力学集成集成到 ALCHEMI 环境中，同时保持对底层内核的直接访问。这将动态转换为完全可定制的环境，用户可以在其中从头开始构建专门的动态类。

功能

  • 专用的 GPU 优先轨迹分析工具
  • 集成和可定制的动态内核 ( Velocity Verlet、NPT、Langevin 自动调温器)
  • FIRE 和 FIRE2 优化器

API 示例

from enum import Enum
import torch
from nvalchemi.data import Batch
from nvalchemi.dynamics.base import BaseDynamics, DynamicsStage
from nvalchemi.hooks import Hook, HookContext
class MySimulatedAnnealer(Hook):
       def __init__(
           self,
           t_start: float,
           t_end: float,
           cooldown_steps: int,
           frequency: int,
           stage: DynamicsStage
       ) -> None:
           # this hook will fire off every `frequency` MD steps,
           # bringing the temperature from `t_start` to `t_end`
           self.frequency = frequency
           self.t_start = t_start
           self.t_end = t_end
           self.cooldown_steps = cooldown_steps
           self.stage = DynamicsStage.BEFORE_STEP
           self.decay = (t_end / t_start) ** (1.0 / cooldown_steps)
           self._current_temp = t_start
       def __call__(self, ctx: HookContext, stage: Enum) -> None:
           # access the calling dynamics class through `HookContext`
           dynamics = ctx.workflow
           dynamics.target_temperature = max(
               dynamics.target_temperature * self.decay,
               self.t_end
           )
class VelocityVerlet(BaseDynamics)
    __needs_keys__: {"energies", "forces", "masses", "velocities"}
    __provides_keys__: {"positions"}
       def __init__(
           self,
           model: BaseModelMixin,
           n_steps: int,
           dt: float = 1.0,  # timestep
           target_temperature: float = 300.0,  # initial temperature
           tau: float = 10.0,  # coupling constant
           hooks: list[Hook] | None = None,
           convergence_hook: ConvergenceHook | dict | None = None,
           **kwargs,
       ):
           super().__init__(model=model, n_steps=n_steps, hooks=hooks, convergence_hook=convergence_hook)
           self.dt = dt
           self.target_temperature = target_temperature
           self.tau = tau
           self._prev_accelerations = None
       
       def pre_update(self, batch: Batch) -> None:
           # perform the first half of velocity Verlet
           with torch.no_grad():
                accelerations = batch.forces / batch.masses
                self._prev_accelerations = accelerations.clone()
                batch.positions.add_(
                    batch.velocities * dt + 0.5 * accelerations * dt**2.0
                )
       def post_update(self, batch: Batch) -> None:
           # perform second half of velocity Verlet, with thermostat
           # temperature update
           with torch.no_grad():
               new_accelerations = batch.forces / batch.masses
               batch.velocities.add_(0.5 * (self._prev_accelerations + new_accelerations) * self.dt)
               ke_per_atom = 0.5 * batch.masses * (batch.velocities**2).sum(dim=-1, keepdim=True)
               # get the total kinetic energy per system
               total_ke = scatter_add_(...)
               current_temp = 2.0 * total_ke / (batch.num_atoms * 3.0)
               ratio = self.target_temperature / current_temp
               lam = torch.sqrt(
                  torch.tensor(1.0 + (self.dt / self.tau) * (ratio - 1.0))
               ).clamp(min=0.8, max=1.2)  # clamp for stability
               batch.velocities.mul_(lam)
# configure the new dynamics class
my_velverlet = VelocityVerlet(
    ...,
    hooks=[
        MySimulatedAnnealer(t_start=900.0, t_end=300.0, cooldown_steps=10, frequency=100, stage=DynamicsStage.BEFORE_STEP)
    ],
)

模型包装器

借助 ALCHEMI 工具包，您可以使用自己的预训练模型和加速物理组件。它为将您自己的模型导入管道提供了必要的基础架构，确保专有架构或特定领域的架构能够利用 GPU 原生编排。这抽象化了不同模型类型的复杂性，提供了从独立模型到生产就绪型高吞吐量仿真的标准化路径。

功能

  • MLIP 支持 ( MACE、TensorNet、AIMNet2)
  • 可组合计算器
  • 标准化模型配置

API 示例

from beartype import beartype
from super_mlip import BestMLIPModel
from nvalchemi._typing import ModelOutputs
from nvalchemi.models.base import BaseModelMixin, ModelConfig, NeighborConfig
class BestMLIPWrapper(nn.Module, BaseModelMixin):
    def __init__(self, model: BestMLIPModel, **kwargs):
          super().__init__(**kwargs)
          # ModelConfig declares model capabilities (which are frozen)
          # and runtime control (mutable) for the rest of the framework
          self.model_config = ModelConfig(
                     outputs=frozenset({"energy", "forces", "hessians"}),
                     # this is actually the default value
                     required_inputs=frozenset({"positions", "atomic_numbers"})
                     autograd_outputs=frozenset({"forces"}),
                     neighbor_config=NeighborConfig(cutoff=5.0, format="coo")
)
    def adapt_input(self, data: Batch, **kwargs) -> dict[str, Any]:
        # adapts the nvalchemi data structure to what is
        # expected by the model
        model_inputs = super().adapt_input(data, **kwargs)
        # dict structure expected by BestMLIPModel
        model_inputs["atom_numbers"] = data.atomic_numbers
        model_inputs["coords"] = data.positions
        return model_inputs
   def adapt_output(self, model_output: any, data: Batch) -> ModelOutputs:
       # adapt the model outputs from the model's forward pass to
       # format expected by nvalchemi
       output = super().adapt_output(model_output, data)
       energies = model_output["energies"]
       output["energies"] = energies
       # check model config for expected outputs
       if "forces" in self.model_config.active_outputs:
           output["forces"] = model_output["forces"]
       return output
    # beartype decorator is optional, but will runtime type check arguments
    @beartype
    def forward(self, data: Batch, **kwargs) -> ModelOutputs:
        model_inputs = self.adapt_input(data, **kwargs)
        # calls BestMLIPModel's forward definition based on MRO
        model_outputs = super().forward(**model_inputs)
        return self.adapt_output(model_outputs, data)

高级数据管理

传统上，在 CPU 和 GPU 之间移动数据的“内存税”是 AI 驱动的发现过程中的一个重大瓶颈。ALCHEMI Toolkit 可充当科学数据的专用编排器，提供构建自定义提取流程所需的基础架构，以便将信息从标准研究文件移动到优化的 GPU 张量中。

这支持扩展发现，使工业规模的模拟能够通过熟悉的界面进行访问。通过标准化原子信息的表示和加载方式，ALCHEMI Toolkit 可确保数据始终驻留在设备上，这意味着整个模拟都保留在 GPU 上，从而实现批量模拟，以优化 GPU 利用率并消除通信开销。

功能

  • 高性能数据加载器
  • ASE 和 Pymatgen 接口
  • AtomicData 和批处理对象

API 示例

from nvalchemi import AtomicData, Batch
from nvalchemi import data
from ase.build import slab
atoms = slab(...)
# Create AtomicData object from ase.Atoms object
data = AtomicData.from_atoms(atoms, device="cuda")
>>> data ...
data.node_properties
data.system_properties
# Create a Batch object from a list of AtomicData
batch = Batch.from_data_list([data, data, data])
batch.num_graphs
batch.get_data(0)
# get first three samples
batch[:2]
batch[mask]
batch["energies"] -> ...
batch.from_atoms([ase.Atoms,...])
# Create a dataset from ase.Atoms
writer = data.AtomicDataZarrWriter("atom_dataset.zarr")
# writer will amortize overhead by writing batches of data;
# this is equivalent to writing individual samples but efficiently
writer.write(batch)
# Read the data from zarr
reader = data.AtomicDataZarrReader("atom_dataset.zarr")
# Dataset treats device natively; individual samples # are placed on GPU and it accelerates preprocessing transforms;
# num_workers sets the number of threads used for async prefetching
dataset = data.Dataset(reader, device = "cuda", num_workers=4)
dataloader = data.DataLoader(dataset, batch_size=16)
for batch in dataloader:
    # do something with batch

开始使用 ALCHEMI 工具包构建分子工作流程

ALCHEMI 工具包为研究人员和开发者提供构建端到端 GPU 原生分子工作流所需的低级基元和高级抽象。将关键瓶颈 (例如近邻列表构建、结构松和集成步骤) 转移到 PyTorch 生态系统中，可以消除主机到设备的内存传输开销，而这种开销传统上限制了 MLIP 驱动的模拟。

无论您是合成混合 ML 或物理势，还是扩展批量分子动力学，ALCHEMI Toolkit 都可提供必要的 API Hook，以便在不牺牲性能的情况下管理复杂的张量状态。

要加速化学和材料科学模拟，并探索构建您自己的自定义工作流，请访问 NVIDIA/nvalchemi-toolkit GitHub 资源库和 ALCHEMI Toolkit 文档。在继续扩展受支持的操作和架构库的同时，我们鼓励您克隆存储库，探索所提供的 Jupyter Notebook，并开始将这些 GPU 加速的工作流集成到您自己的发现流程中。

致谢

我们在此感谢 Orbital 的 James Gin、Tim Duignan、Vaidas imkus、MatGL 的 Shyue Ping Ong 教授、Matlantis 的 Susumu Rhino、Ryuhei Okuno 和 Jethro Tan 与我们合作，将 NVIDIA ALCHEMI Toolkit 应用到他们的平台中。我们还要感谢 Nikita Fedik、Roman Zubadyuk、Atul Thakur 和 Logan Ward 对此博文的贡献。

标签

开发工具与技巧 | 模拟/建模/设计 | ALCHEMI | PyTorch | 精选 | 计算化学/ 材料科学

关于作者

Erica Tsai 是 NVIDIA ALCHEMI 的高级产品营销经理。经过训练，她是一名无机化学家，拥有普林斯顿大学的学士学位和麻省理工学院的博士学位。加入 NVIDIA 之前，Erica 曾在麦肯锡公司领导科学 AI 计划。如今，她专注于弥合深度科学、技术执行和商业价值之间的差距。
Justin S.Smith 是 NVIDIA 化学和材料科学 AI 部门的高级开发者关系经理。他是一位计算化学家，于 2018 年获得佛罗里达大学博士学位，从事原子模拟 AI 研究。随后，他成为洛斯阿拉莫斯国家实验室的一名工作人员科学家，专注于 ML 在反应化学和材料科学中的应用。
Wen Jie Ong 是 NVIDIA ALCHEMI 的高级产品经理。他受训成为有机和聚合物化学家，并在麻省理工学院获得博士学位，在麻省理工学院发现了一种新型动态共价化学。他还是一位科学家兼企业家，拥有将新型聚合物技术商业化的经验。最近，他担任 McKinsey 的副合作伙伴，为化学品和技术公司提供增长和创新方面的建议。
Dallas Foster 是 NVIDIA HPC 和 AI 部门的高级深度学习软件工程师。他在俄勒冈州立大学获得数学博士学位，并曾在洛斯阿拉莫斯国家实验室、美国国家大气研究中心和麻省理工学院工作。作为 NVIDIA PhysicsNeMo 团队的成员，他特别关注深度学习在天气预报和分子动力学方面的应用和部署。
Kelvin Lee 是 NVIDIA 的高级深度学习软件工程师。他在澳大利亚新南威尔士大学获得物理化学博士学位。加入 NVIDIA 之前，他曾在天体物理学中心担任学术研究职务 | 此前，哈佛大学史密森分校和麻省理工学院曾在英特尔实验室担任研究科学家。他的工作重点是加速计算和开发者生产力强制化学和材料科学、光谱学和天体物理学。