deepmodeling / deepmodeling/deepmd-kit
[Feature Request] Support decoupled observer-model inference for model deviation in deepmd/kk
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
### Summary
Add model-deviation support to the Kokkos-accelerated LAMMPS `pair_style deepmd/kk` for compatible edge/graph `.pt2` models by separating the **driver model** from the additional **observer models**.
Only model 0 should run every MD step and supply the energy, force, and virial used to advance the trajectory. Models 1 through N-1 only need inference on `out_freq` steps to estimate committee uncertainty. This preserves normal model-deviation semantics while keeping the Kokkos device-resident path for trajectory integration.
Currently, the same configuration is rejected during initialization:
```text
ERROR: pair style deepmd/kk does not support model deviation.
```
### Detailed Description
#### Current behavior
The regular LAMMPS `pair_style deepmd` supports an ensemble of models:
```lammps
pair_style deepmd \
graph.000.pt2 graph.001.pt2 graph.002.pt2 graph.003.pt2 \
out_freq 100 out_file model_devi.out
pair_coeff * * H C N O Cl
```
Its effective execution model is already decoupled:
- on ordinary MD steps, evaluate only model 0;
- on `out_freq` steps, evaluate all models;
- use model 0 output for dynamics;
- use the ensemble outputs only to calculate deviation statistics.
Running the equivalent input with:
```bash
lmp -k on g 1 -sf kk -in in.lammps
```
selects `pair_style deepmd/kk` and fails in `PairDeepMDKokkos::init_style()` because `source/lmp/pair_deepmd_kokkos.cpp` explicitly rejects `numb_models != 1`:
```cpp
if (numb_models != 1) {
error->all(
FLERR,
"pair style deepmd/kk does not support model deviation."
);
}
```
This restriction was introduced together with `deepmd/kk` in #5758, whose description states that the Kokkos path requires one model.
#### Why the driver and observers can be decoupled
For a committee of models M0, M1, ..., MN-1:
- **M0 is the driver**: it runs on every step and its force advances the trajectory.
- **M1 ... MN-1 are observers**: they run only when `step % out_freq == 0`.
- Observer outputs never modify the trajectory; they only contribute to force, energy, and virial deviation statistics.
Therefore, supporting model deviation does not require all models to participate in MD integration or to retain all model workspaces simultaneously.
A possible Kokkos execution flow is:
```text
Every MD step:
build the device graph once
evaluate M0
scatter M0 outputs and advance dynamics
On model-deviation steps:
reuse the same device graph
evaluate M1, update online statistics
evaluate M2, update online statistics
...
evaluate MN-1, update online statistics
write model_devi.out
```
Sequential observer inference would avoid scaling temporary inference workspace with the committee size. Online Welford accumulation could avoid retaining an `N_models x N_atoms x 3` force tensor. The persistent state can remain O(N_atoms):
- model-0 force used by dynamics;
- one observer scratch force/virial buffer;
- running mean and M2 accumulators for deviation.
For MPI/domain-decomposed execution, each observer's ghost force contributions should be reverse-communicated to owner atoms before atom-wise force-deviation statistics are updated.
#### Requested behavior
Allow a multi-model command such as:
```lammps
pair_style deepmd \
graph.000.pt2 graph.001.pt2 graph.002.pt2 graph.003.pt2 \
out_freq 100 out_file model_devi.out
```
to run with:
```bash
lmp -k on g 1 -sf kk -in in.lammps
```
while preserving the regular `pair_style deepmd` semantics:
- model 0 drives the trajectory;
- observer models run only on model-deviation output steps;
- `model_devi.out` reports compatible force/energy/virial statistics;
- `atomic`, `relative`, and `relative_v` work where applicable;
- single-model `deepmd/kk` behavior and performance remain unchanged.
#### Possible implementation direction
`PairDeepMDKokkos` already builds the device graph and owns device-resident energy, force, and atomic-virial buffers. A possible implementation could:
1. expose device-edge/canonical-graph inference for individual models held by `DeepPotModelDevi`, or provide a suitable iterator/evaluation API;
2. validate that every committee model supports the same device graph schema, type map, cutoff, parameter dimensions, edge-vector precision, and communication contract;
3. build the Kokkos graph once per timestep and reuse it across the driver and observers;
4. evaluate model 0 on every step;
5. evaluate models 1 through N-1 sequentially only on `out_freq` steps;
6. reverse-communicate observer ghost forces/virials before owner-atom statistics are accumulated;
7. calculate model-deviation statistics online where possible;
8. reuse the regular `pair_style deepmd` output format and option semantics.
#### Acceptance criteria
- Two or more compatible graph/edge `.pt2` models initialize under `pair_style deepmd/kk`.
- Model 0 alone drives the trajectory.
- Observer models execute only at `out_freq` steps.
- `model_devi.out` agrees with regular `pair_style deepmd` for a small deterministic system.
- Energy, force, global virial, and atom-wise force deviation are covered.
- `atomic`, `relative`, and `relative_v` are supported or clearly diagnosed.
- Serial and at least two-rank MPI/domain-decomposition tests pass.
- Observer evaluation does not require inference workspace proportional to the committee size.
- Existing single-model `deepmd/kk` behavior and performance remain unchanged.
- Automated Kokkos regression coverage uses at least two `.pt2` models.
This is relevant to active-learning workflows because model-deviation exploration conventionally uses four independently trained models. I can provide a four-model compressed DPA4C `.pt2` input and help validate the implementation on an NVIDIA H20.
### Further Information, Files, and Links
Related work:
- `deepmd/kk` introduction and explicit one-model limitation: #5758
- DPA4C graph/compact `.pt2` support: #5972
- DP-GEN request for `.pt2` model-deviation deployment: deepmodeling/dpgen#1925
- DP-GEN implementation of `.pt2` export and artifact forwarding: deepmodeling/dpgen#1926
Observed environment:
- DeePMD-kit 3.2.0 built from a clean upstream source checkout
- compressed PyTorch-exportable `.pt2` DPA4C models
- clean upstream LAMMPS `stable_22Jul2025_update2` source build
- NVIDIA H20
- four independently trained committee models
The explicit `deepmd/kk` multi-model rejection is also present in current upstream DeePMD-kit source (`source/lmp/pair_deepmd_kokkos.cpp`) and is not specific to a locally modified LAMMPS build.
Contributor guide
Assessment
This issue has not been assessed yet.