NVIDIA ALCHEMI Toolkit is a GPU-first Python framework for AI atomic simulations. Run batched molecular dynamics (MD) and relaxation, train or fine-tune machine-learned interatomic potentials (MLIPs), and scale the same composable workflows from one GPU to multi-GPU systems.
- Batched GPU simulation: run many MD or geometry relaxation jobs in one model pass; inflight batching keeps the GPU occupied as systems finish.
- Bring your own model: use MACE, AIMNet2 or UMA, or wrap another MLIP with
BaseModelMixin; compute energies, forces and stress, plus local Hessians and Hessian-vector products for supported models; add Ewald/PME electrostatics or DFT-D3(BJ) dispersion. - Training, fine-tuning and distillation: combine energy, force and stress objectives with validation, restartable checkpoints, exponential moving averages and distributed data-parallel training; adapt models with LoRA or distill them into compact students.
- Generative workflows: use a model-agnostic API for batched structure generation, optional conditioning and hooks; compose generators with relaxation or MD.
- Molecular crystal structure prediction (CSP) primitives: generate random rigid-molecule packings, expand crystal symmetry, and compare structures for de-duplication.
- Energy landscape exploration: find minimum-energy paths with batched regular or climbing-image NEB; explore configurations with umbrella sampling, metadynamics, adaptive biasing force and replica exchange.
- Multi-GPU scaling: split one large atomic system across GPUs with spatial domain decomposition.
- Data at scale: use GPU-resident
AtomicDataandBatch, Zarr storage, balanced dataset mixing, transforms and CUDA-stream prefetching. - Composable dynamics and hooks: create custom integrators, chain simulation
stages with
+or|, compile fused steps, and attach logging, safety, sampling, profiling or convergence logic throughout each step and on system admission or graduation. - Agent-ready guidance: an agent skill and
AGENTS.mdteach coding agents the toolkit APIs and repository conventions.
Built on nvalchemi-toolkit-ops
for GPU-optimized neighbor lists and interaction kernels via NVIDIA warp-lang.
Skills — one agent skill, skills/nvalchemi-toolkit/, that scopes a request and
routes to task notes for data, models, dynamics, training, generative workflows, and
extensions. Agents that read AGENTS.md find it inside a clone; to use it elsewhere,
copy the folder into your agent's skills directory (see the
Agent skills user guide).
AGENTS.md — repository-wide guidance (setup, conventions, gotchas). Agents that follow the AGENTS.md convention (e.g. Codex, Cursor, OpenCode) load it natively. Claude Code auto-loads CLAUDE.md instead: to get the same guidance there, add a CLAUDE.md to your clone containing the single line @AGENTS.md (an import), or symlink it (ln -s AGENTS.md CLAUDE.md).
Build atomic data and run a batched forward pass
import torch
from nvalchemi.data import AtomicData, Batch
from nvalchemi.models.demo import DemoModel, DemoModelWrapper
# Create two molecules
mol_a = AtomicData(
positions=torch.randn(4, 3),
atomic_numbers=torch.tensor([6, 6, 1, 1], dtype=torch.long),
)
mol_b = AtomicData(
positions=torch.randn(3, 3),
atomic_numbers=torch.tensor([8, 1, 1], dtype=torch.long),
)
# Batch for GPU-efficient inference
batch = Batch.from_data_list([mol_a, mol_b])
# Wrap a model and run
model = DemoModelWrapper(DemoModel())
outputs = model(batch)
print(outputs["energy"].shape) # [2, 1] — one energy per system
print(outputs["forces"].shape) # [7, 3] — one force vector per atomGeometry optimization with convergence detection
import torch
from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics import DemoDynamics, ConvergenceHook
from nvalchemi.dynamics.hooks import LoggingHook, NaNDetectorHook
# Dynamics reads and writes these per-step buffers, so allocate them up front.
def system(n_atoms: int, z: int) -> AtomicData:
return AtomicData(
positions=torch.randn(n_atoms, 3),
atomic_numbers=torch.full((n_atoms,), z, dtype=torch.long),
forces=torch.zeros(n_atoms, 3),
energy=torch.zeros(1, 1),
velocities=torch.zeros(n_atoms, 3),
)
batch = Batch.from_data_list([system(4, 6), system(3, 8)])
dynamics = DemoDynamics(
model=model,
n_steps=10_000,
dt=0.5,
convergence_hook=ConvergenceHook.from_fmax(0.05),
hooks=[
LoggingHook(backend="csv", log_path="run.csv", frequency=100),
NaNDetectorHook(),
],
)
with dynamics:
result = dynamics.run(batch)Multi-stage pipeline: relax then MD (single GPU)
from nvalchemi.dynamics import DemoDynamics
optimizer = DemoDynamics(model=model, n_steps=500, dt=0.5)
md = DemoDynamics(model=model, n_steps=1_000, dt=1.0)
# + fuses stages: one forward pass, masked updates per sub-stage
fused = optimizer + md
with fused:
fused.run(batch)Distributed pipeline across GPUs
# Launch with: torchrun --nproc_per_node=2 my_pipeline.py
from nvalchemi.dynamics import DemoDynamics
optimizer = DemoDynamics(model=model, n_steps=500, dt=0.5)
md = DemoDynamics(model=model, n_steps=1_000, dt=1.0)
# | distributes stages: one dynamics per GPU rank
pipeline = optimizer | md
with pipeline:
pipeline.run()Train a model with validation
import torch
from nvalchemi.training import (
EnergyMSELoss,
ForceMSELoss,
OptimizerConfig,
TrainingStrategy,
ValidationConfig,
default_training_fn,
)
# Assumes `model` is a BaseModelMixin wrapper and `train_loader` /
# `val_loader` are nvalchemi DataLoaders (see the data pipeline guide).
device = torch.device("cuda")
# Compose an objective: weighted energy + force terms
loss_fn = 1.0 * EnergyMSELoss() + 10.0 * ForceMSELoss()
strategy = TrainingStrategy(
models=model,
optimizer_configs=OptimizerConfig(
optimizer_cls=torch.optim.AdamW,
optimizer_kwargs={"lr": 1e-3},
),
num_steps=10_000,
training_fn=default_training_fn,
loss_fn=loss_fn,
devices=[device],
validation_config=ValidationConfig(
validation_data=val_loader,
validation_fn=default_training_fn,
loss_fn=loss_fn,
every_n_steps=500,
),
)
strategy.run(train_loader)
print(strategy.last_validation)For a complete runnable script, see
examples/advanced/10_mace_training.py.
Split one large system across GPUs (domain decomposition)
# Launch with: torchrun --nproc_per_node=2 my_dd_run.py
import torch
from nvalchemi.distributed import DistributedManager, DomainConfig, DomainParallel
from nvalchemi.dynamics import NVTLangevin
from nvalchemi.models.mace import MACEWrapper
DistributedManager.initialize()
dm = DistributedManager()
device = torch.device(dm.device)
mesh = dm.initialize_mesh(
mesh_shape=(dm.world_size,), mesh_dim_names=("domain",)
)
# The wrapper and the integrator are the same objects you would use on a
# single GPU; `batch` is the full system, built on rank 0 only.
wrapper = MACEWrapper.from_checkpoint("medium-0b2", device=device).eval()
integrator = NVTLangevin(
model=wrapper, dt=0.5, temperature=300.0, friction=0.01, n_steps=200
)
# One DomainConfig + one wrap is the entire user-facing addition. Atoms are
# partitioned spatially; halo exchange and cross-rank reductions are automatic.
domain_cfg = DomainConfig(cutoff=float(wrapper.cutoff), skin=0.5, mesh=mesh)
with DomainParallel(
dynamics=integrator, config=domain_cfg, n_steps=200
) as dynamics:
owned = dynamics.partition(batch if dm.rank == 0 else None)
dynamics.run(owned)
DistributedManager.cleanup()For complete runnable scripts, see
examples/distributed/ and the
distributed guide.
The quickest way to install:
pip install \
--extra-index-url https://download.pytorch.org/whl/cu130 \
--extra-index-url https://pypi.nvidia.com \
'nvalchemi-toolkit[cu13]'For development:
git clone https://github.com/NVIDIA/nvalchemi-toolkit.git
cd nvalchemi-toolkit
uv sync --extra cu13cu13 is the default development CUDA variant. For CUDA 12 environments, run
uv sync --extra cu12 instead and pass the same extra to uv run, for example
uv run --extra cu12 pytest test/. The Makefile does this automatically:
make test CUDA_EXTRA=cu12. CUDA-aligned optional extras follow the same
pattern, for example uv sync --extra cu12 --extra mace or
make test CUDA_EXTRA=cu12 OPTIONAL_EXTRAS=mace. To include documentation
dependencies, add --group docs. Avoid uv sync --all-extras, because the
CUDA variants are mutually exclusive.
Optional extras:
pip install \
--extra-index-url https://download.pytorch.org/whl/cu126 \
--extra-index-url https://pypi.nvidia.com \
'nvalchemi-toolkit[cu12]' # Specify CUDA 12 version
pip install \
--extra-index-url https://download.pytorch.org/whl/cu130 \
--extra-index-url https://pypi.nvidia.com \
'nvalchemi-toolkit[cu13,mace]' # MACE model support, CUDA 13
pip install \
--extra-index-url https://download.pytorch.org/whl/cu126 \
--extra-index-url https://pypi.nvidia.com \
'nvalchemi-toolkit[cu12,mace]' # MACE model support, CUDA 12The uma extra is mutually exclusive with mace and the CUDA extras
(incompatible e3nn / torch pins) and resolves into its own environment.
See the Installation Guide for detailed setup instructions.
Features planned for upcoming releases:
- In-flight NEB batching: admit new reaction paths as completed paths leave the batch
- Multipole electrostatics: support multipole-aware potentials, including MACE-POLAR, in Toolkit workflows
- Domain decomposition performance: improve overlap of computation and communication across GPUs
- OPES enhanced sampling: add the on-the-fly probability enhanced sampling method
NVIDIA ALCHEMI Toolkit is in public beta. During this phase, the API is subject to change. Feature requests, bug reports, and general feedback are welcome via GitHub Issues.
Apache 2.0 — see LICENSE for details.