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

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.