deepmodeling / deepmodeling/deepmd-kit

[pt_expt] NeighborGraph compiled training is energy-only; generalize the trace to non-energy fittings

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

## 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

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.