deepmodeling / deepmodeling/deepmd-kit
refactor(pt_expt): make graph-native preprocessing and export independent of get_sel
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Background
Several pt_expt orchestration paths still treat `get_sel()` as a universal model contract even when the selected execution is a carry-all graph:
- the training entrypoint unconditionally runs `BaseModel.update_sel()`;
- graph compile sample size is derived from `sum(model.get_sel())`;
- export metadata, tracing, dynamic-shape construction, and inference wrappers read `get_sel()`;
- model interfaces require a dense capacity even for graph-native models.
For large legacy values such as `sel=416`, compile tracing also creates an unnecessarily large synthetic graph despite runtime edge count being dynamic.
This issue covers orchestration and capability discovery. DPA1 normalization/statistics semantics belong to the companion model-semantics issue.
## Proposed design
Introduce a backend-neutral `NeighborContract`, owned by dpmodel rather than pt_expt, for example with:
- `representation="graph"`;
- `requires_capacity=False`;
- `capacity=None`.
The contract should be exposed through model/descriptor plugins and be usable before model construction when preprocessing configuration.
Use it to:
- skip neighbor-capacity discovery and `update_sel()` for graph-native models;
- retain the legacy update path for dense models;
- replace graph-path `get_sel()` queries in training, compile, export, and inference;
- serialize a versioned neighbor contract in exported metadata;
- provide legacy adapters for PT and old checkpoints.
The entrypoint should call one model-level preparation hook that recursively handles descriptor/model composition. Avoid backend- or descriptor-name conditionals in the trainer.
Graph compile/export should use a small deterministic synthetic topology that exercises dynamic atom and edge counts, multiple edges per center, and isolated centers. Its size must not be derived from a physical neighbor capacity.
## Dependency direction
The shared contract belongs under `deepmd/dpmodel` or another backend-neutral module. Making dpmodel or PT import it from `deepmd/pt_expt` would invert the current dependency direction.
Torch-specific tracing, compilation, and export adapters should remain under pt_expt.
## Compatibility requirements
- Dense PT and legacy pt_expt models continue to expose and consume their existing `sel`.
- Existing exported metadata remains readable.
- New metadata represents graph models without inventing a dummy capacity.
- Composite models obtain one consistent contract or report an actionable incompatibility.
- Input-stat computation may still run when needed; only neighbor-capacity discovery must be skipped.
## Performance requirements
- Synthetic trace size and compile latency must be independent of legacy `sel`.
- Graph preprocessing remains O(E) and does not construct dense padded neighbor lists.
- Dynamic atom and edge counts must not cause recompilation for every batch shape.
- The capability layer must not add graph breaks or hot-path Python dispatch.
## Acceptance tests
- A graph-native DPA1 training config without `sel` reaches model construction and training without calling `update_sel()`.
- Dense models retain current automatic-selection behavior.
- Graph compile and export contain no capacity-dependent sample sizing.
- Compiled/eager energy, force, virial, loss, and parameter gradients agree across varying atom and edge counts.
- Export/import and inference work for both legacy-capacity and graph-native metadata.
- A 118-type `se_atten_v2`, `attn_layer=0`, `silu`, `float32` model covers the highest type index and HDF5 stat loading.
- Increasing a legacy `sel` value does not increase graph trace size or compile setup memory.
Refs #5755.
Contributor guide
Assessment
This issue has not been assessed yet.