pet-kokkos models
Every model published in lab-cosmo/upet (the PET-MAD / UPET family), converted for pet-kokkos, a C++/Kokkos implementation of PET that runs without libtorch on CPUs and on NVIDIA and AMD GPUs.
The weights are unchanged from upstream. Each model is two files:
<name>.json holds the hyperparameters, species map and composition energies,
and <name>.safetensors holds the weights. They were produced with
tools/convert_pet.py, and each JSON's source_model names the checkpoint it
came from.
Usage
# everything (~12 GB)
hf download peterspackman/pet-kokkos-models --local-dir ~/.local/share/pet/models
# one model
hf download peterspackman/pet-kokkos-models --include "pet-mad-s.*" --local-dir models
pet-eval pet-mad-s structure.xyz --forces
pet-kokkos looks for models by name in $PET_MODEL_DIR, ./models, ., and
~/.local/share/pet/models.
Models
| model | source checkpoint |
|---|---|
| pet-mad-xs | pet-mad-xs-v1.6.0 |
| pet-mad-s | pet-mad-s-v1.6.0 |
| pet-mad-m | pet-mad-m-v1.6.0 |
| pet-mols-s | pet-mols-s-v1.1.0 |
| pet-omad-xs | pet-omad-xs-v1.0.0 |
| pet-omad-s | pet-omad-s-v1.0.0 |
| pet-omad-l | pet-omad-l-v0.1.0 |
| pet-omat-xs | pet-omat-xs-v1.0.0 |
| pet-omat-s | pet-omat-s-v1.0.0 |
| pet-omat-m | pet-omat-m-v1.0.0 |
| pet-omat-l | pet-omat-l-v1.0.0 |
| pet-omat-xl | pet-omat-xl-v1.0.0 |
| pet-omatpes-l | pet-omatpes-l-v0.1.0 |
| pet-oam-l | pet-oam-l-v0.1.0 |
| pet-oam-xl | pet-oam-xl-v1.0.0 |
| pet-omol-s | pet-omol-s-v1.0.0 |
| pet-omol-m | pet-omol-m-v1.0.0 |
| pet-omol-l | pet-omol-l-v1.0.0 |
| pet-spice-s | pet-spice-s-v0.2.0 |
| pet-spice-l | pet-spice-l-v0.2.0 |
For what each model was trained on and the level of theory behind it, see the upstream model card. Please cite the upstream work when you use these models.
Verification
The files in this repository, downloaded as published, were evaluated with
pet-kokkos (HIP, on an MI250X) and compared against metatomic's evaluation of
the original checkpoint. The test structures were a water molecule, diamond, and
rattled diamond (tools/test_all_models.py). The cond models were also checked
at charge +1 and spin 2. dE/E is the relative energy deviation, and dF (eV/Å)
and dStress (eV/ų) are the worst absolute deviations.
model architecture dE/E dF dStress
-------------------------------------------------------------------------------------
PASS pet-mols-s resi PostLN G2 A2 d256 1.1e-07 7.6e-06 2.3e-07
PASS pet-omad-xs feed PreLN G2 A1 d128 grid 1.4e-07 9.3e-06 5.2e-07
PASS pet-omat-xs feed PreLN G2 A1 d128 grid 1.3e-07 1.1e-05 1.1e-06
PASS pet-mad-xs feed PreLN G2 A1 d128 solver 1.9e-07 9.8e-06 4.3e-07
PASS pet-spice-s feed PreLN G3 A1 d192 5.2e-08 9.7e-06 5.2e-07
PASS pet-omad-s feed PreLN G3 A1 d256 grid 8.8e-08 8.3e-06 3.4e-07
PASS pet-omat-s feed PreLN G3 A1 d256 grid 5.5e-08 1.2e-05 4.2e-07
PASS pet-mad-s feed PreLN G3 A1 d256 solver 1.2e-07 9.1e-06 3.9e-07
PASS pet-omol-s feed PreLN G3 A1 d256 grid cond 1.1e-07 9.0e-06 3.6e-07
PASS pet-omol-m feed PreLN G4 A2 d384 grid cond 2.0e-07 7.4e-06 7.1e-07
PASS pet-omat-m feed PreLN G3 A2 d384 grid 1.3e-07 8.5e-06 4.1e-07
PASS pet-mad-m feed PreLN G3 A2 d384 solver 1.4e-07 8.0e-06 2.3e-07
PASS pet-spice-l feed PreLN G3 A2 d512 1.2e-07 1.8e-05 6.2e-07
PASS pet-omad-l feed PreLN G3 A2 d512 9.2e-08 8.7e-06 3.6e-07
PASS pet-omatpes-l feed PreLN G3 A2 d512 9.0e-08 9.1e-06 3.7e-07
PASS pet-oam-l feed PreLN G3 A2 d512 5.8e-08 7.9e-06 3.8e-07
PASS pet-omat-l feed PreLN G4 A2 d512 grid 1.8e-07 9.5e-06 5.0e-07
PASS pet-omol-l feed PreLN G4 A2 d512 grid cond 2.0e-07 7.6e-06 3.4e-07
PASS pet-oam-xl feed PreLN G5 A3 d640 grid 2.2e-07 1.1e-05 9.5e-07
PASS pet-omat-xl feed PreLN G5 A3 d640 grid 1.4e-08 1.3e-05 6.0e-07
20 pass, 0 mismatch, 0 unsupported, 0 error (of 20)
License
BSD-3-Clause, the same as the upstream checkpoints.