NVIDIA ALCHEMI Toolkit#
GPU-first Python framework for AI-driven atomic simulations.
NVIDIA ALCHEMI Toolkit gives you a unified, composable API for machine-learned interatomic potentials — from single-GPU prototyping to distributed high-throughput production. Wrap any MLIP, assemble multi-stage simulation pipelines with Python operators, and let inflight batching keep your hardware fully utilized.
Who Is This For?#
Run batched geometry/cell optimization, molecular dynamics (NVE/NVT/NPT), or multi-stage relaxation-equilibration-production workflows — all from a single Python script.
Plug your own potential into the framework with BaseModelMixin, compose
it with existing force fields, and generate training data through
GPU-buffered trajectory capture.
Scale from one GPU to an entire node with DistributedPipeline. Inflight
batching and size-aware sampling handle load balancing automatically.
Highlights#
Bring your own model — wrap MACE, AIMNet2, or any PyTorch MLIP in a few lines with a standardized
ModelCardinterface.Compose, don’t configure — fuse stages on one GPU with
+, distribute across GPUs with|, and inject behavior at nine hook points per step.GPU-native data —
AtomicDataandBatchare Pydantic-validated,jaxtyping-annotated graph structures that live on-device.Inflight batching — converged samples are replaced mid-run so the GPU never idles during high-throughput screening.
Zarr-backed I/O — write trajectories with zero-copy GPU buffering; reload through a CUDA-stream-prefetching
DataLoader.