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.