Examples#

End-to-end examples demonstrating ALCHEMI Toolkit functionality.

Basic Examples#

These examples introduce the core nvalchemi-toolkit workflow to users familiar with tools like ASE. Each script focuses on one feature set and runs end-to-end on a single GPU in under 60 seconds.

01 — Data Structures: AtomicData and Batch API. 02 — Geometry Optimization: FIRE optimizer with NeighborListHook and ConvergenceHook. 03 — ASE Integration: Loading ASE structures, FreezeAtomsHook on a surface system. 04 — NVE MD: Microcanonical dynamics, WrapPeriodicHook, EnergyDriftMonitorHook. 05 — NVT MD: Langevin thermostat, thermalization, LoggingHook to CSV.

AtomicData and Batch: Graph-structured molecular data

FIRE Geometry Optimization with Lennard-Jones Argon

ASE Integration: Real Molecular Structures in nvalchemi-toolkit

Microcanonical (NVE) Molecular Dynamics

Canonical (NVT) Molecular Dynamics with Langevin Thermostat

Intermediate Examples#

These examples assume familiarity with the basic tier and introduce the storage layer, performance monitoring, and more complex pipeline patterns. Training examples are labeled explicitly in this list; if the collection grows further, split training workflows into a dedicated examples section so inference, dynamics, and training entry points remain easy to scan.

01 — Multi-Stage Pipeline: FusedStage composition, LoggingHook CSV output, step-budget migration, fused hooks for global status monitoring.

02 — Trajectory I/O: Writing trajectories to Zarr, reading back with DataLoader, round-trip validation.

03 — NPT MD: Pressure-controlled dynamics with the MTK barostat, LJ stress computation, cell fluctuation monitoring.

04 — Inflight Batching: SizeAwareSampler, Mode 2 FusedStage run (batch=None), system_id tracking, ConvergedSnapshotHook collecting results.

05 — Safety and Monitoring: NaNDetectorHook, MaxForceClampHook, EnergyDriftMonitorHook, StageTimingHook — defensive MD patterns.

06 — DDP MLP Training: DDPHook with a simple MLP, dummy AtomicData, single-node torchrun launch, and auto/gloo/nccl backend selection.

07 — Rich Training Reporting: Live Rich dashboard driven by synthetic training losses, validation metrics, progress counters, and learning-rate scheduler values.

Multi-Stage Dynamics Pipelines with FusedStage

Writing and Replaying Trajectories with Zarr

NPT Barostat Validation: Expansion and Contraction

Processing Large Datasets with Inflight Batching

Defensive MD: Safety Hooks and Performance Monitoring

Distributed Training: DDPHook with a Dummy MLP

Rich Training Reporting

Advanced Examples#

These examples are for users who want to extend the nvalchemi-toolkit framework. They require understanding of the intermediate tier.

01 — Biased Potential: BiasedPotentialHook for harmonic COM restraints and umbrella sampling patterns.

02 — Custom Hook: Implementing the Hook protocol with a full radial distribution function accumulator.

03 — Custom Convergence: ConvergenceHook with multiple criteria and custom_op for arbitrary convergence logic.

04 — MACE NVT: Using a real MACE MLIP for NVT dynamics; automatic neighbor list wiring via ModelConfig; LJ fallback for CI.

05 — Custom Integrator: Subclassing BaseDynamics to implement a velocity-rescaling thermostat; the pre_update/post_update contract; _init_state for stateful integrators.

07 — Composable Model Composition: Combining LJ + Ewald models with the + operator; PipelineModelWrapper for dependent pipelines.

08 — AIMNet2 + Ewald Pipeline: Composing AIMNet2 with Ewald electrostatics and DFTD3 dispersion in a multi-group pipeline.

09 — UMA NVE/NVT: Driving the fairchem UMA foundation model through NVE / NVT dynamics with energy-drift tracking; OMat crystals and OMol molecules via task selection on UMAWrapper.from_checkpoint.

10 — MACE Training: Training a ScaleShiftMACE model with the ALCHEMI training stack; Zarr dataloading, scheduled Huber losses, EMA, checkpointing, validation, and distributed launch patterns.

Biased Sampling with BiasedPotentialHook

Writing a Custom Hook: Radial Distribution Function

Multi-Criteria Convergence with Custom Operators

NVT MD with MACE (with LJ Fallback)

Building a Custom Integrator by Subclassing BaseDynamics

Long-Range Electrostatics with Ewald Summation

Additive Model Composition (LJ + Ewald)

AIMNet2 + Ewald Pipeline Geometry Optimization

UMA Foundation Model: NVE / NVT / NPT Molecular Dynamics

MACE Training with ALCHEMI Training Utilities

Distributed Examples#

These examples cover the two multi-GPU paths in NVAlchemi:

Pipeline parallelism (examples 01–02) — map ranks to dynamics stages with DistributedPipeline.
Domain decomposition (examples 03–05) — shard one system across ranks with DomainParallel, including the “bring your own model” arc.

All require multiple GPUs and must be launched with torchrun.

Warning

These examples are not executed during the Sphinx documentation build. To run them, use torchrun as shown in each example.

Pipeline Architecture Overview#

A DistributedPipeline maps GPU ranks to dynamics stages. Systems flow between stages via fixed-size NCCL communication buffers:

Key concepts:

Upstream ranks (prior_rank=None): hold a SizeAwareSampler and push graduated (converged) systems to the next rank.
Downstream ranks (next_rank=None): receive systems from the prior rank and write results to a sink.
BufferConfig: must be set to a fixed size on all ranks; NCCL requires identical message sizes every communication step.
torchrun --nproc_per_node=N launches one process per GPU; each process runs only the stage assigned to its rank.

Running the Examples#

01 — Parallel FIRE → Langevin (4 GPUs required):

torchrun --nproc_per_node=4 examples/distributed/01_distributed_pipeline.py

# CPU/debug mode (set backend="gloo" in the script first):
torchrun --nproc_per_node=4 --master_port=29500 examples/distributed/01_distributed_pipeline.py

02 — Monitoring with LoggingHook, StageTimingHook, and ZarrData (4 GPUs required):

torchrun --nproc_per_node=4 examples/distributed/02_distributed_monitoring.py

After running example 02, per-rank CSV logs and Zarr trajectory stores are written to the working directory. Rank 0 also prints a collated summary.

Example Descriptions#

01 — Distributed Pipeline: Two independent FIRE → NVTLangevin sub-pipelines running on 4 GPUs. Demonstrates DistributedPipeline wiring, BufferConfig, and HostMemory sinks.
02 — Distributed Monitoring: Same topology as example 01, augmented with per-rank LoggingHook and StageTimingHook for observability, and ZarrData sinks for persistent trajectory storage. Shows post-run log collation on rank 0.

Domain-Decomposition Examples#

These shard a single system across ranks with DomainParallel (halo exchange + force consolidation handled by the framework).

# 03 — MACE NVT Langevin MD, trajectory written to xyz from rank 0
torchrun --nproc_per_node=2 examples/distributed/03_mace_nvt_distributed.py

# 04 / 05 — bring-your-own model, validated against a single-process reference
torchrun --nproc_per_node=2 examples/distributed/04_byo_pytorch_mpnn.py
torchrun --nproc_per_node=2 examples/distributed/05_byo_graph_transformer.py

# 06 — MACE NPT (barostat) MD, evolving-cell trajectory written from rank 0
torchrun --nproc_per_node=2 examples/distributed/06_mace_npt_distributed.py

# 07 — 2-D-parallel dynamics: FIRE → NVT, each stage domain-decomposed
torchrun --nproc_per_node=4 examples/distributed/07_fire_nvt_dd.py

03 — MACE NVT Distributed: End-to-end distributed MD with a stock MACEWrapper: a short NVTLangevin trajectory under DomainParallel, with per-step neighbour-list rebuild and xyz snapshot logging from rank 0. No distributed-aware code at the user layer.
04 — BYO PyTorch MPNN: The full bring-your-own arc for a plain-PyTorch Behler-Parrinello potential: architecture → wrapper → run → trace_and_validate against a single-process reference → MLIPSpec.save/load. An MPNN-halo model whose forward is scatter-aggregations + autograd needs no distributed code.
05 — BYO Graph Transformer (Warp kernel): The same arc when the model embeds a performance-critical Warp kernel that is opaque to ShardTensor dispatch. Shows declaring the kernel’s distribution semantics once via OpAdapter.
06 — MACE NPT Distributed: The constant-pressure sibling of example 03: a NPT trajectory (Nosé–Hoover thermostat + isotropic barostat) under DomainParallel, with the cell relaxing toward equilibrium. The barostat/thermostat couple to global quantities (total kinetic energy, degrees of freedom, pressure tensor); the framework’s dynamics coordinator all-reduces them and broadcasts the replicated cell + barostat state each step, so the only user change from example 03 is requesting stress and swapping in NPT.
07 — 2-D-parallel dynamics: FIRE → NVT, each stage domain-decomposed: The 2-D generalization of example 01: a FIRE relaxation → NVTLangevin MD pipeline where each stage is itself domain-decomposed. A (pipeline, domain) DeviceMesh gives each stage a whole domain sub-mesh row; a stage is just DomainParallel(dynamics) handed to DistributedPipeline(stages, mesh=mesh). DomainParallel overrides the pipeline’s communication seam so the group lead performs the cross-stage hand-off (over the pipeline axis) while the group scatters/gathers to its sub-mesh — no distributed-aware code in the model or the integrators. Keep the per-step domain dimension intra-node (NVLink) and let the rare-hand-off pipeline dimension span nodes (IB); 4 GPUs = 2 stages × 2 domain. FIRE’s velocity mixing couples to global power/norm scalars (v·f / v·v / f·f), which the dynamics coordinator all-reduces within each stage’s domain group.

Benchmarks#

Performance + force-equivalence benchmarks for the domain-decomposition path live in benchmark/distributed/ (two config-driven runners covering LJ, Ewald, PME, MACE, AIMNet2, and UMA). See benchmark/distributed/README.md.