deepmodeling / deepmodeling/deepmd-kit

feat(pt_expt): decouple DPA1 graph normalization and statistics from dense sel

Open
#5,824 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

The pt_expt DPA1/`se_atten_v2` carry-all graph lower no longer uses `sel` to cap or pad the runtime neighbor graph. However, `sel` is still part of the descriptor's mathematical and statistical state:

- `self.nnei = sum(sel)` is used to normalize graph aggregation;
- descriptor mean/std buffers are shaped as `(ntypes, nnei, 4)`;
- input statistics build a dense `sel`-sized neighbor list and include padded zeros;
- omitting `sel` therefore still changes model construction and statistics.

Graph-native DPA1 is uncapped at runtime but is not truly `sel`-independent.

## Proposed design

Separate the three concepts currently represented by `sel`:

1. `neighbor_capacity`: required only by dense neighbor-list execution;
2. `aggregation_normalization`: serialized model semantics used by DPA1 aggregation;
3. `input_stat_policy`: defines how descriptor statistics are accumulated.

Use a versioned compatibility mapping:

- legacy configs/checkpoints map `sum(sel)` to a fixed aggregation normalizer and retain legacy padded-stat semantics;
- graph-native configs may use `neighbor_capacity=None`;
- a graph-native fixed or data-derived normalizer is computed once and serialized;
- graph-native statistics store the canonical per-type vector, such as `(ntypes, 1, 4)`, rather than duplicating it over `nnei`.

Implement statistics as streaming reductions over graph edges. Legacy padded statistics should be reproducible from sums, squared sums, real-edge counts, and the legacy slot count without materializing `(nframes, nloc, sel, 4)` tensors.

Dynamic raw-degree normalization should not become the default because it changes DPA1 semantics and may break smoothness at the cutoff. Any smooth-degree policy should be explicit and versioned separately.

## Compatibility requirements

- Existing configs and checkpoints containing `sel` must preserve their current normalization and statistics.
- Serialization/deserialization must distinguish legacy and graph-native semantics.
- PT dense execution must retain its current capacity-based behavior.
- Continuing a legacy checkpoint in pt_expt graph mode must not silently reinterpret its normalization.

## Performance requirements

- Graph-native statistics must not allocate memory proportional to `nloc * sel`.
- Runtime graph execution must remain carry-all and dynamic in the number of edges.
- The normalizer must be compile-friendly and must not introduce per-atom Python control flow.
- No regression is allowed in compiled graph throughput or peak memory.

## Acceptance tests

- A legacy `sel=416` DPA1 model preserves energy, force, virial, statistics, and checkpoint round-trip results.
- A new graph-native `se_atten_v2` config can omit neighbor capacity while using every neighbor within `rcut`.
- A system whose degree exceeds a legacy capacity is not truncated on the graph path.
- Streaming legacy-stat accumulation matches the current padded implementation.
- Graph-native HDF5 statistics round-trip correctly with a 118-element type map.
- `se_atten_v2`, `attn_layer=0`, `silu`, and `float32` pass eager/compiled forward, loss, and parameter-gradient comparisons with varying atom and edge counts.

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.