[Feature Request] Support decoupled observer-model inference for model deviation in deepmd/kk #6000

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:

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:

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:

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:

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:

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:

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:

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: perf(dpa1): optimize graph CUDA inference and deployment #5758
• DPA4C graph/compact .pt2 support: feat(pt-expt): add compact descriptor DPA4C 🎉🎉🎉 #5972
• DP-GEN request for .pt2 model-deviation deployment: [Feature Request] Support PyTorch-exportable (.pt2) models for Kokkos-accelerated LAMMPS model deviation dpgen#1925
• DP-GEN implementation of .pt2 export and artifact forwarding: feat: support DPA4 and DPA4C model formats 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.