From e5520f665883a03748c43ccd59b566fdd8ce6b1a Mon Sep 17 00:00:00 2001 From: Nikita Fedik Date: Fri, 31 Jul 2026 01:53:04 -0400 Subject: [PATCH] docs: update README for 0.2.0 Bring the README up to date for the 0.2.0 release. Add the subsystems that landed during this cycle, training (nvalchemi.training) and spatial domain decomposition (nvalchemi.distributed), along with the UMA wrapper and MultiDataset composition. Restructure the list to lead with capabilities rather than architecture, with batched GPU simulation first. Split model listings into interatomic potentials and interaction terms, since Ewald and PME read charges an MLIP predicts. Clarify the | operator as stage-parallel so it reads distinctly from the new intra-system spatial decomposition. Add a Roadmap section following the format used in nvalchemi-toolkit-ops, and replace the inline breaking-changes paragraph with a pointer to CHANGELOG.md, which already carried all five items. Add training and domain-decomposition example snippets, and fix three snippets that do not run against the current API: - LoggingHook(frequency=100) raised TypeError; backend is required - DemoDynamics.run(batch) raised AttributeError on batch.forces; forces, energy and velocities must be pre-allocated on AtomicData - DemoDynamics(model=model, dt=0.5) raised TypeError; n_steps is required Snippets 1-3 were verified by execution in a uv sync --extra cu13 environment. Snippets 4-6 need two ranks; their API calls were confirmed but they were not run end to end. Signed-off-by: Nikita Fedik --- README.md | 223 +++++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 197 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 3bcd4499..83cb9c18 100644 --- a/README.md +++ b/README.md @@ -11,32 +11,91 @@ 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. +machine-learned interatomic potentials (MLIPs), batched molecular dynamics, +composable multi-stage simulation pipelines, and model training and fine-tuning: +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`](https://github.com/NVIDIA/nvalchemi-toolkit-ops) - for GPU-optimized neighbor lists, dispersion, and electrostatics via - NVIDIA `warp-lang` -- Agents as first-class citizens; includes a core skills library that - teaches agents how to use `nvalchemi` efficiently in agentic workflows +- **Batched GPU simulation**: + - many systems batched together in one pass on a single GPU + - molecular dynamics: NVE, NVT with Langevin or Nosé-Hoover chain + thermostats, NPT and NPH + - geometry relaxation: FIRE and FIRE2, with optional cell relaxation + - `+` fuses stages so a relax-then-MD workflow costs one forward pass per + step (`FusedStage`) + - `SizeAwareSampler` swaps in new samples as others converge, keeping the + GPU full on long runs + +- **Training and fine-tuning** ([guide](docs/userguide/training.md)): + - composable energy, force and stress losses with weight schedules + - validation and restartable checkpoints + - EMA and DDP hooks + +- **Multi-GPU scaling** ([guide](docs/userguide/distributed.md)): + - spatial domain decomposition splits one large system across ranks via + `DomainParallel` and `ShardTensor`, with automatic halo exchange + - `|` distributes pipeline stages across ranks (`DistributedPipeline`) + +- **Interatomic potentials**: + - wrappers for MACE, AIMNet2 and UMA + - your own model via `BaseModelMixin` + +- **Interaction terms**: + - Ewald and particle mesh Ewald electrostatics, which read the per-atom + charges an MLIP such as AIMNet2 predicts + - DFT-D3(BJ) dispersion + - Lennard-Jones, usable on its own as a lightweight potential + +- **Data at scale**: + - `AtomicData` and `Batch`, Pydantic-backed and GPU-resident + - Zarr serialization for trajectories and datasets + - `MultiDataset` mixes several datasets with balanced sampling; + `InMemoryDataset` keeps a hot dataset resident + - per-sample and per-batch transforms, CUDA-stream prefetching + +- **Extensible by design**: + - one `BaseModelMixin` wrapper, with capabilities declared through + `ModelConfig`, works unchanged in dynamics, training and multi-GPU runs + - custom integrators and optimizers by subclassing `BaseDynamics` + - custom loss functions by subclassing `BaseLossFunction`, and custom + training steps through `training_fn` + - custom storage backends by subclassing `Reader` + - your own hooks at nine per-step insertion points, for logging, safety + checks, enhanced sampling, profiling and convergence detection + - stack interaction terms onto a potential with `PipelineModelWrapper`, + wiring outputs between steps so AIMNet2 supplies the charges Ewald reads + +- **Agent skills**: task-specific API guides under `.claude/skills/` (see [Using with AI coding agents](#using-with-ai-coding-agents)) +Built on [`nvalchemi-toolkit-ops`](https://github.com/NVIDIA/nvalchemi-toolkit-ops) +for GPU-optimized neighbor lists and interaction kernels via NVIDIA `warp-lang`. + +Upgrading from 0.1.0? See [CHANGELOG.md](CHANGELOG.md) for breaking changes and +migration snippets. + +### Roadmap + +Features planned for upcoming releases: + +- **Generative models**: model-agnostic abstraction of generative models for + the ALCHEMI Toolkit simulation pipeline +- **Crystal structure prediction (CSP) primitives**: composable, batched + building blocks for molecular CSP workflows +- **Enhanced sampling**: GPU-resident collective variables and biasing methods +- **Model distillation**: pipeline for distilling large, accurate potentials + into compact models for fast production inference +- **LoRA adapters**: parameter-efficient fine-tuning that maintains many + specialized variants of one base potential without duplicating its weights +- **Hessians and phonons**: analytical second derivatives through automatic + differentiation for vibrational and thermodynamic property prediction +- **Domain decomposition optimization**: continued performance improvement of + spatial domain decomposition +- **Kernel improvements** at the + [`nvalchemi-toolkit-ops`](https://github.com/NVIDIA/nvalchemi-toolkit-ops) + level + ### Using with AI coding agents The repository ships agent-facing guidance at two levels: @@ -88,15 +147,34 @@ print(outputs["forces"].shape) # [7, 3] — one force vector per atom Geometry optimization with convergence detection ```python +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(frequency=100), NaNDetectorHook()], + hooks=[ + LoggingHook(backend="csv", log_path="run.csv", frequency=100), + NaNDetectorHook(), + ], ) with dynamics: result = dynamics.run(batch) @@ -110,8 +188,8 @@ with dynamics: ```python from nvalchemi.dynamics import DemoDynamics -optimizer = DemoDynamics(model=model, dt=0.5) -md = DemoDynamics(model=model, dt=1.0) +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 @@ -128,8 +206,8 @@ with fused: # 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) +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 @@ -139,6 +217,96 @@ with pipeline: +
+Train a model with validation and checkpointing + +```python +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`](examples/advanced/10_mace_training.py). + +
+ +
+Split one large system across GPUs (domain decomposition) + +```python +# Launch with: torchrun --nproc_per_node=2 my_dd_run.py +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh + +from nvalchemi.distributed import DomainConfig, DomainParallel +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.models.mace import MACEWrapper + +dist.init_process_group(backend="nccl") +device = torch.device(f"cuda:{dist.get_rank()}") +torch.cuda.set_device(device) +mesh = DeviceMesh( + "cuda", list(range(dist.get_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) +dynamics = DomainParallel(dynamics=integrator, config=domain_cfg, n_steps=200) + +owned = dynamics.partition(batch if dist.get_rank() == 0 else None) +dynamics.run(owned) +dynamics.close() +``` + +For complete runnable scripts, see +[`examples/distributed/`](examples/distributed/) and the +[distributed guide](docs/userguide/distributed.md). + +
+ ## Installation The quickest way to install: @@ -184,6 +352,9 @@ pip install \ 'nvalchemi-toolkit[cu12,mace]' # MACE model support, CUDA 12 ``` +The `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](docs/userguide/about/install.md) for detailed setup instructions.