deepmodeling / deepmodeling/deepmd-kit
[pt_expt] NeighborGraph compiled training is energy-only; generalize the trace to non-energy fittings
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
The NeighborGraph route's compiled-training trace is hardcoded to the energy output, so non-energy models (dos / property / dipole / polar) are excluded from the graph default-flip even when their descriptor supports the graph lower.
The graph forward itself is already output-agnostic: `forward_common_lower_graph` (deepmd/pt_expt/model/make_model.py) routes any fitting through `fit_output_to_model_output_graph`, reducing every reducible output and differentiating every `r_differentiable` output. The eager graph route works for non-energy fittings today via explicit `neighbor_graph_method=` opt-in (exercised by the dos graph test).
## Where the energy-specificity lives
All of it is in the compiled-training trace, `_trace_and_compile_graph` in `deepmd/pt_expt/train/training.py`:
- `model.do_grad_r("energy")` / `model.do_grad_c("energy")` — the output-def key is a literal.
- `_translate_energy_keys` (from `deepmd/pt_expt/model/ener_model.py`) — maps internal keys (`energy`, `energy_redu`, `energy_derv_r`, ...) to the public keys the loss consumes (`atom_energy`, `energy`, `force`, `virial`). The public mapping is per model type (dos: `dos` -> `atom_dos`, etc.).
Because of this, `model_uses_graph_lower` (`deepmd/pt_expt/model/graph_lower.py`) and the `_resolve_graph_method` default-flip (`deepmd/pt_expt/model/make_model.py`) gate the graph DEFAULT to energy-output models: a non-energy model default-flipped onto the graph route would train fine in eager but raise `KeyError('energy')` at its first compiled batch — and gating only the compiled side would silently diverge eager (graph) from compiled (dense).
## Proposed fix
1. Derive the differentiability flags from the model's `atomic_output_def()` instead of the `"energy"` literal.
2. Replace `_translate_energy_keys` in the trace with an output-def-driven translation (via `get_reduce_name`/`get_deriv_name`) or a per-model-class `translate_lower_keys` hook that `EnergyModel` implements with the current mapping.
3. Add a compiled-vs-eager parity test for graph training of a non-energy model (e.g. dpa1/dpa2 descriptor + dos fitting).
4. Drop the energy condition from `model_uses_graph_lower` / `_resolve_graph_method` for the TRAINING default (the export gate is a separate concern, tracked in #5806).
## Context
Gate introduced on #5779 (see the rationale comments at the two gate sites). The model type should be decoupled from how the descriptor consumes neighbors (graph vs dense nlist); this issue removes the remaining coupling on the training side.
Contributor guide
Assessment
This issue has not been assessed yet.