deepmodeling / deepmodeling/deepmd-kit

refactor(pt_expt): make graph-native preprocessing and export independent of get_sel

Open
#5,825 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.