nvidia_configs: add deepmd env hosting the built-in DPA pretrained models

[OwenPriceSkelly]

What

Adds sample_model_configurations/nvidia_configs/deepmd.py, a DeePMD-kit environment, plus tests/sample_configs/test_deepmd_env.py. Several groups we are onboarding (Delta especially) run DeePMD-kit and fine-tune DPA3; this closes the "DeePMD isn't in the catalog" gap.

Ten canonical ids, all mapped to names in deepmd-kit's built-in pretrained registry (deepmd/pretrained/registry.py, the same names dp pretrained download accepts, each with HF + mirror URLs and a sha256):

id	upstream	notes
dpa3-omol-large	DPA3-Omol-Large	single-task, OMol25, charge/spin via fparam
dpa4-{nano,mini,neo,air,plus}-omat24	DPA4-<Size>-OMat24-v20260805	single-task, DPA4/SeZM, OMat24
dpa-3.3-1m, dpa-3.2-5m, dpa-3.1-3m, dpa-2.4-7m	DPA-3.3-1M …	multitask OpenLAM, head= required
dpa:custom	—	any .pt checkpoint / frozen .pth via weights=

Design

• Deps: deepmd-kit[torch]>=3.2.0 (released 2026-08-19; first release with the DPA4 loaders and the full registry). The [torch] extra pins torch==2.11.0.*; torch is listed explicitly so it routes through the cu128 index, since PyPI's 2.11 wheels bundle CUDA 13. uv lock --script resolves universally (83 packages: torch 2.11.0+cu128, deepmd-kit 3.2.0, e3nn, vesin, nvalchemi-toolkit-ops 0.4.1, mpich).
• Fetch: setup() calls deepmd's resolve_model_path(name, cache_dir=$XDG_CACHE_HOME/deepmd/pretrained/models) and hands DP() the local path. A cache hit only re-hashes the file (read-only), so serve time never writes to the shared install and needs no network.
• Heads: the multitask ids raise without head= (same rule as uma task / mace-mh-1 head from remove bad uma checkpoint and make task heads mandatory to specify #219), even though deepmd itself would silently pick the Default alias (OMat24) on 3.2-5M / 3.3-1M. HEADS lists valid branches per id; VERIFY_KWARGS selects OMat24 (covers verify's H2O) for the four multitask ids and for dpa:custom.
• Custom family is dpa, not deepmd: the smoke-test weights= leg finds its baseline by id prefix (_family_of), and every shipped id starts with dpa. The first declared id (dpa3-omol-large) is single-task so that leg needs no head.
• Device: deepmd's PyTorch backend has no device argument; it reads DEVICE=cpu or cuda:{LOCAL_RANK} at first import. setup() sets LOCAL_RANK for cuda:N (correct under an existing CUDA_VISIBLE_DEVICES mask, unlike overwriting it) and DEVICE=cpu for CPU, before importing deepmd. Works because the worker imports no torch before setup().
• Charge/spin: two upstream mechanisms exist — fparam=[charge, spin] (DPA3-Omol-Large, DPA-3.2-5M) and a charge_spin embedding (DPA-3.3-1M, DPA4-OMol). A thin DP subclass fills whichever the loaded model uses from atoms.info["charge"] / ["spin"], matching the other OMol-trained envs; explicit fparam / charge_spin keys win; absent both, the model's trained-in neutral-singlet default applies.

Not verified yet

Not built or verified on any cluster. Delta is the natural first target (most of the DeePMD groups are there). Expect the first verify to surface something, per docs/environments.md.

Deliberately left out

• DPA4-OMol25 sizes, DPA4C (needs the pt-expt backend / .pt2 export), DPA4-Beta, and DPA-3.1-3M-FT (the Matbench-Discovery .pth): on HF but not in deepmd's registry, so they would need a second fetch path.
• No amd_configs / aurora_configs counterpart (ROCm/XPU untested).
• Licenses: DPA-2/3 checkpoints are CC-BY-4.0; the DPA4 checkpoints are CC-BY-NC-4.0.

Tests

uv run pytest tests/sample_configs — 152 passed (13 new: registry resolution into the shared cache, mandatory head, device env contract, custom-path loading, charge/spin translation for both mechanisms, VERIFY_KWARGS ↔ HEADS consistency).

🤖 Generated with Claude Code