Skip to content

v0.6.0

Latest
Latest

Choose a tag to compare

View all tags
@yury-lysogorskiy yury-lysogorskiy released this 07 Jul 10:06
0.6.0
4110538
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

This is effectively the 0.6.0 release. Three big updates headline it.

1. Uncertainty Quantification (UQ)

New tensorpotential/uq/ package: GMM-based per-atom uncertainty (gamma) via Mahalanobis distance on a Johnson–Lindenstrauss random-projection ("basis-RP") feature of the model's invariant B-basis.

  • Two new ComputeFunctions: full HAL-style uncertainty forces (dsigma/dr, virial_sigma) and a cheaper gamma-only variant for active-learning monitoring.
  • New grace_uq CLI: build (artifact + SavedModel export), info (summary + health diagnostics), predict (E/F/S + per-atom gamma), select (candidate selection by extrapolation/diversity with per-element stratification).
  • Build pipeline (uq/cli/build/): out-of-core, multi-worker/multi-GPU artifact builder with a resumable 3-step pipeline — clustering (clusters.py) → covariance (workers.py) → threshold calibration (thresholds.py) — orchestrated by master.py, with multi-source data resolution (data_resolve.py), weighted/filtered builds, elbow-based cluster selection, streaming ingest (streaming.py), and crash-safe .npz writes.
  • docs/gracemaker/uq.md: full user guide (theory, CLI, weighted builds, filtering, Python API).

2. LAMMPS Kokkos support

New Kokkos weight-export tooling

  • Export GRACE-1L/2L/3L weights to .npz for the LAMMPS Kokkos pair styles via grace_utils -p <potential> export_kokkos (backed by scripts/_kokkos_export.py), with compile-time cap validation.
  • Optionally bakes the UQ random-projection matrix into the exported weights, so per-atom uncertainty is available in LAMMPS-Kokkos runs.

3. GRACE foundation models

Hosted on HuggingFace: AMS-ICAMS-RUB/grace-foundation-models.

  • UQ-enabled: the released foundation models now ship with the baked-in UQ random-projection matrix, so per-atom uncertainty (gamma) is available out of the box.
  • Default precision changed to fp32: a bare model name (e.g. GRACE-2L-OMAT-large-ft-E) now resolves to the fp32 saved model/checkpoint. Full fp64 builds are published non-default under the -fp64 suffix; -mx models remain natively mixed precision.
  • New 3-layer (3L) foundation models (v4/uqv6 release): GRACE-3L-OMAT-large and GRACE-3L-OMAT-large-ft-AM, natively fp32 (no -fp64 variant). LAMMPS-Kokkos: pair_style grace/3l/kk (mixed) or grace/3l/kk/fp32.

4. Misc

  • Fit dashboard (grace_dashboard): Flask app that scans fit output folders and shows training curves, Pareto scatter plots, and metrics tables.
  • grace_models CLI overhaul: list / info <name> / download <name> [--kokkos], with a 3-state on-disk glyph (cached / available-not-downloaded / not-offered), introspected artifact flags, and download hints.
  • Intra-epoch (mid-epoch) checkpointing: resume from mid-epoch via an intra_epoch_save flag with ff_skip + epoch-rewind logic;
  • metadata_utils.py: read the model metadata block from model.yaml and resolve parameter dtype.

Performance updates

  • fp32 default for foundation models — roughly halves memory and ~2× throughput vs fp64.
  • Optional dense neighbour-list format (dense_nbr / compute_dense): per-atom-uniform bond reshape + matmul instead of segment_sum/scatter aggregation, ~1.3–1.7× faster per structure; dual compute/compute_dense SavedModel export lets the calculator pick the engine (mode="uniform" vs "diverse").
  • NCCL multi-GPU (distribute.py): forces NCCL and avoids MirroredStrategy's slow CPU-hop scalar reduction.
  • Tight dense padding (DensePaddingManager): ski-rental shape promotion minimises wasted padded FLOPs and recompiles in the dense calculator.

Supporting changes

Distributed / multi-GPU training

  • cli/distribute.py: single source of truth for MultiWorkerMirroredStrategy (NCCL forced), auto-synthesizing single-worker TF_CONFIG to avoid MirroredStrategy's slow CPU-hop scalar reduction.
  • examples/slurm_multinode/submit.sbatch: multi-node SLURM launcher building per-node TF_CONFIG for MWMS over InfiniBand/NCCL.

Note: This release also carries substantial rework of core modules not itemized above (notably functions/couplings.py, instructions/compute.py, calculator/asecalculator.py, foundation_models.py, cli/train.py, tpmodel.py, tensorpot.py, presets.py, loss.py) plus matching test and doc updates.

Assets 2
Loading