deepmodeling / deepmodeling/deepmd-kit
feat(pt_expt): decouple DPA1 graph normalization and statistics from dense sel
- 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
Assessment
This issue has not been assessed yet.