NVIDIA/nvalchemi-toolkit

NVIDIA/nvalchemi-toolkit

Description: ALCHEMI Toolkit is a developer toolkit for accelerating training and inference for AI in chemistry and material science.

Language: Python

License: Apache-2.0

Stars: 92

Forks: 20

Open issues: 16

Created: 2026-02-26T21:56:09Z

Pushed: 2026-06-10T00:27:27Z

Default branch: main

Fork: no

Archived: no

README:

NVIDIA ALCHEMI Toolkit

 

High-performance deep-learning framework for atomic simulations

NVIDIA ALCHEMI Toolkit is a GPU-first Python framework for building, running, and deploying AI-driven atomic simulation workflows. It provides a unified interface for machine-learned interatomic potentials (MLIPs), batched molecular dynamics, and composable multi-stage simulation pipelines: all designed for high throughput on NVIDIA GPUs.

Key Features

• Bring your own model — wrap any MLIP (MACE, AIMNet2, or your own) with

a standard BaseModelMixin that handles input/output adaptation, capability negotiation, and runtime control via ModelConfig

• Graph-structured data — AtomicData and Batch provide Pydantic-backed,

GPU-resident graph representations with built-in serialization to Zarr

• Composable dynamics — subclass BaseDynamics for custom integrators;

compose stages with + (single-GPU FusedStage) or | (multi-GPU DistributedPipeline)

• Pluggable hook system — nine insertion points per step for logging,

safety checks, enhanced sampling, profiling, and convergence detection

• Inflight batching — SizeAwareSampler replaces graduated samples

on the fly, maximizing GPU utilization across long-running pipelines

• High-performance primitives — built on

`nvalchemi-toolkit-ops` for GPU-optimized neighbor lists, dispersion, and electrostatics via NVIDIA warp-lang

• Agents as first-class citizens; includes core SKILLS.md library that

teach agents how to use nvalchemi efficiently in agentic workflows. Simply copy the .claude/skills folder contents to your project repository or home directory depending on use case and agent platform (e.g. Claude Code, Cursor, OpenCode).

Example Snippets

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 atom

Geometry optimization with convergence detection

from nvalchemi.dynamics import DemoDynamics, ConvergenceHook
from nvalchemi.dynamics.hooks import LoggingHook, NaNDetectorHook

dynamics = DemoDynamics(
model=model,
n_steps=10_000,
dt=0.5,
convergence_hook=ConvergenceHook.from_fmax(0.05),
hooks=[LoggingHook(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, dt=0.5)
md = DemoDynamics(model=model, 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, dt=0.5)
md = DemoDynamics(model=model, dt=1.0)

# | distributes stages: one dynamics per GPU rank
pipeline = optimizer | md
with pipeline:
pipeline.run()

Installation

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 cu13

cu13 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 12

See the [Installation Guide](docs/userguide/about/install.md) for detailed setup instructions.

Contributions & Disclaimers

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.

License

Apache 2.0 — see [LICENSE](LICENSE) for details.