deepmodeling / deepmodeling/deepmd-kit
[BUG] PT-expt se_e2_a type_one_side=False fails symbolic tracing in compiled training
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
### Bug summary
PT-expt compiled training fails for `se_e2_a` with `type_one_side=False` while tracing the first model forward. The two-sided implementation computes atom-type-dependent slice boundaries by converting a tensor reduction to a Python integer; symbolic tracing cannot concretize that value.
The failure reproduces with `disp_avg=False` and occurs before display-metric accumulation. It was found during extra validation for #5809, not as a failure of an existing CI unit test.
### DeePMD-kit version and environment
- Source checkout: `67837d7ff8b971f1503b8b8d7efcf210d91688c4`, based on upstream/master `28b7d068801716765ab8119257f814596e49a10c`. The failing descriptor code is unchanged from that upstream commit.
- Editable CUDA installation built from source; installed version metadata: `3.2.1.dev13+g5580ebd3f.d20260911`.
- Backend: PyTorch Exportable (`pt_expt`).
- Python `3.13.15`, PyTorch `2.13.0+cu130`, CUDA runtime `13.0`.
- GPU: NVIDIA RTX PRO 6000 Blackwell.
### Steps to reproduce
Run from the repository root in an installed DeePMD-kit environment. The script uses only the repository's public water example and writes training output in a temporary directory.
```bash
CUDA_VISIBLE_DEVICES=0 NUM_WORKERS=0 python reproduce_se_a_compile.py
```
`reproduce_se_a_compile.py`:
```python
"""Reproduce the two-sided SeA symbolic tracing failure from a public example."""
import json
import os
import tempfile
from pathlib import Path
from deepmd.pt_expt.entrypoints.main import get_trainer
from deepmd.utils.argcheck import normalize
from deepmd.utils.compat import update_deepmd_input
repo = Path.cwd()
config = json.loads((repo / "examples/water/se_e2_a/input_torch.json").read_text())
config["model"]["descriptor"].update(
type_one_side=False, sel=[4, 8], neuron=[4, 8], axis_neuron=4,
rcut=3.0, rcut_smth=0.5,
)
config["model"]["fitting_net"]["neuron"] = [8]
config["model"]["data_stat_nbatch"] = 1
config["training"]["training_data"] = {
"systems": [str(repo / "examples/water/data/data_0")],
"batch_size": 1,
}
config["training"].pop("validation_data", None)
config["training"].update(numb_steps=1, enable_compile=True, disp_avg=False)
config = normalize(update_deepmd_input(config, warning=False))
with tempfile.TemporaryDirectory(prefix="deepmd-se-a-compile-") as work:
os.chdir(work)
get_trainer(config).run()
```
This exact self-contained reproducer was run and reaches the error below.
### Observed error
```text
deepmd/pt_expt/train/training.py: _trace_and_compile -> make_fx
deepmd/dpmodel/descriptor/se_e2_a.py:596
offset += int(xp.sum(xp.astype(atype_loc == ti, xp.int32)))
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode:
Could not extract specialized integer from data-dependent expression u0
(unhinted: u0). (Size-like symbols: none)
```
The count depends on the values of `atype_loc`, not merely its shape. The subsequent slice boundaries therefore require a data-dependent Python integer during symbolic tracing.
### Expected behavior and controls
- Compiled training should handle the two-sided descriptor without extracting these tensor values into Python integers.
- Eager training with `type_one_side=False` succeeds.
- Compiled single-task and multi-task training with `type_one_side=True` succeeds in the same environment.
- Disabling `disp_avg` does not change the failure.
The ordinary descriptor-consistency suite includes `type_one_side=False`, but the checked PT-expt compiled-training fixtures use `type_one_side=True`. A regression for `type_one_side=False` together with `enable_compile=True` is needed. The failing combination has not yet been retested on the CI-pinned PyTorch versions; this report establishes the behavior on the environment above.
Contributor guide
Assessment
This issue has not been assessed yet.