v0.2.0

Second public-beta release of NVIDIA ALCHEMI Toolkit. The headline is domain
decomposition — multi-GPU inference and dynamics for machine-learned
interatomic potentials — alongside first-class training validation, restartable
checkpoints, and composable atomic datapipes.

Domain Decomposition

Multi-GPU inference and dynamics over a spatial partition of one system, with
energy, forces and stress matching a single-GPU reference to fp32 rounding.

• DomainParallel — wraps any dynamics stage to run it across a device mesh.
  A decomposed run is the same script as a single-GPU one plus this wrapper.
• Two strategies — a spatial halo (each rank owns a region and exchanges
  a ghost shell with its neighbours) and graph parallel (atoms split by a
  balanced index range, node features all-gathered per message-passing layer).
  Halo suits locality-bound models; graph parallel suits dense neighbour
  matrices and small boxes.
• MLIPSpec / distribution_spec — a model declares what its outputs mean
  under decomposition (per-node vs per-graph, how each reduces across ranks) and
  the framework owns the mechanism. Wrappers stay distribution-agnostic
  everywhere else.
• Bring-your-own model — OpAdapter, MethodAdapter, PythonAdapter and
  JitAdapter let an external model's opaque kernels participate in the halo
  without editing that model.
• DistributedModel / DistributedPipelineModel — single models and
  composed pipelines both decompose; each sub-model gets its own right-sized
  ghost region over one shared owned partition.
• torch.compile support — fixed-shape graph padding (GraphPadder,
  DenseBatchPadder) holds compiled graphs stable across MD steps, so a
  decomposed run reaches a recompile-free steady state.
• trace_and_validate — spawns a real multi-rank run of a candidate spec and
  reports where a distributed result diverges from single-process, including
  halo-completeness and partition-degeneracy diagnostics.
• pin_fp32 — pins full-precision fp32 for runs that must match a reference;
  TF32 makes distributed and single-process results diverge well beyond fp32
  rounding.

Shipping specs: Ewald, PME, MACE, AIMNet2 and UMA.

Training

• First-class validation — set a ValidationConfig on
  strategy.validation_config and validation runs at the configured step or
  epoch cadence plus a final pass, with the latest summary on
  strategy.last_validation. Mechanics live in a public, context-managed
  ValidationLoop usable standalone. A new AFTER_VALIDATION hook stage fires
  immediately after each pass.
• Restartable checkpoints — strategy checkpoints carry models, optimizers,
  schedulers, runtime counters and restart-safe device placement, with a
  periodic hook for step- or epoch-based saves.
• Checkpointable hooks — hooks such as EMA save restart state alongside the
  strategy, so a resumed run keeps averaged weights instead of restarting them.
• EMA on sharded models — an EMAHook._build_averaged_model seam lets a
  caller that owns model sharding supply a pre-built AveragedModel, enabling
  EMA on fully_shard (FSDP2) / DTensor models.
• Metric-driven schedulers — ReduceLROnPlateau and friends step on a
  validation metric.

Data Layer

• Atomic datapipes — PhysicsNeMo-compatible pipes with MultiDataset
  composition, multidataset-aware sampling policies, and fused batch loading
  that preserves the Zarr reader's coalesced I/O path.
• InMemoryDataset — a fully materialized in-memory dataset for
  small-to-medium corpora, avoiding per-epoch Zarr reads.
• User-specified transforms — Dataset accepts a transforms= kwarg.

Models

• UMAWrapper — UMA (fairchem-core) as a first-class wrapper, including under
  domain decomposition.
• Pipeline neighbor-list adaptation policy — PipelineModelWrapper states
  how each sub-model's neighbour list is derived, so composed models with
  different cutoffs compose predictably.

Dynamics

• NVE energy conservation across neighbour rebuilds — rebuilds no longer
  perturb the conserved quantity.
• NPT/NPH under decomposition — halo geometry now tracks a barostat-deformed
  cell, and the cross-rank thermostat state stays synchronised.

Developer & Agent Experience

• 12 Claude Code agent skills (.claude/skills/), including a new
  nvalchemi-distributed guide covering strategy choice, authoring a
  distribution_spec for a bring-your-own model, and writing a custom
  integrator that stays correct across ranks.
• Expanded user guide and API reference for 0.2.0, with an install matrix.

Breaking Changes

• EwaldModelWrapper and PMEModelWrapper default to hybrid_forces=False.
  The analytic direct-output path does not produce consistent gradients and is
  rejected under domain decomposition. Forces and stress now come from autograd
  over the energy; pass hybrid_forces=True explicitly for the previous path.
• stress outputs standardised on tensile-positive Cauchy stress. Callers
  relying on the previous sign convention must flip it.
• Hook context split into HookContext, DynamicsContext and their training
  counterpart.
• EvaluateHook removed in favour of first-class validation on
  TrainingStrategy.
• Dataset-level explicit batch reads now use load_batches(...); the raw
  read_many(...) API remains on readers.

Requirements

• Python 3.11–3.13
• PyTorch >= 2.8
• nvalchemi-toolkit-ops[torch] >= 0.4.1
• Optional: [mace], [aimnet], [uma], [ase], [pymatgen],
  [tensorboard], [cu12], [cu13] extras

Domain decomposition requires a multi-GPU host; single-GPU workflows are
unchanged.