deepmodeling / deepmodeling/deepmd-kit

[pt_expt] Graph-lower .pt2 export/inference is energy-only; non-energy models lack deployment consumers

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

Graph-lower `.pt2` export is restricted to energy-output models. This is not an artificial gate: for non-energy models there is currently no deployment stack that could consume a graph artifact. This issue records the missing consumers so the restriction can be lifted per output type as demand appears.

## Current state

- The graph export entry points are methods on `EnergyModel` (`forward_lower_graph_exportable` / `forward_lower_graph_exportable_with_comm` in `deepmd/pt_expt/model/ener_model.py`), and the export gate in `deepmd/pt_expt/utils/serialization.py` keys on `model_uses_graph_lower`, which requires `"energy"` in the atomic output def.
- The Python `DeepEval` graph fast path (`deepmd/pt_expt/infer/deep_eval.py`) translates energy output keys only.
- C++: `DeepPotPTExpt` is the only graph runner. There is no C++ DeepDOS in any backend (dos inference is Python-only), and `DeepTensor` (dipole/polar) has no pt_expt support even for the DENSE `.pt2` schema.
- LAMMPS consumes energy models only.

An exported graph artifact for a dos/dipole/polar model would therefore be a file nothing can load; the gate refuses to produce it.

## What lifting the restriction requires (per output type)

1. **dos / property (Python-only inference):** generalize the exportable graph lower off `EnergyModel` (output-def-driven key translation, shared with the training generalization in the companion issue) and extend the `DeepEval` graph branch to non-energy output keys. No C++ work needed.
2. **dipole / polar:** first requires dense `.pt2` support in `DeepTensor` for the pt_expt backend, then the graph schema on top. Substantially larger; only worth doing with a concrete use case.

## Suggested order

Do the training-side generalization first (companion issue), then 1, then 2 on demand.

## Context

Gate introduced on #5779; companion training-side issue: #5805.

Contributor guide

Open the contributing guide

Research direction

Start with deepmd/pt_expt/model/ener_model.py, deepmd/pt_expt/utils/serialization.py, and deepmd/pt_expt/infer/deep_eval.py, then read companion issue #5805 and the existing DeepPotPTExpt runner. Define the output type to address first; done means its graph export has a deployment consumer, with non-energy keys supported in DeepEval where applicable. Dipole and polar additionally require DeepTensor dense .pt2 support.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.