[BUG] PT-expt se_e2_a type_one_side=False fails symbolic tracing in compiled training

[OutisLi]

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.

CUDA_VISIBLE_DEVICES=0 NUM_WORKERS=0 python reproduce_se_a_compile.py

reproduce_se_a_compile.py:

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

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.

[njzjz-bot]

Thanks for the minimal reproducer. This looks like a real pt_expt compile-path bug rather than anything related to disp_avg.

I checked the current implementation on master: the two-sided se_e2_a path still derives per-type slice boundaries from atype_loc with a tensor reduction followed by Python int(...), then uses those values as Python slice bounds. Because the counts are data-dependent, make_fx/symbolic tracing cannot specialize them safely. Your eager-mode and type_one_side=True controls are therefore consistent with the implementation.

The fix should keep this selection/grouping tensor-native instead of materializing data-dependent boundaries as Python integers, and we should add a regression covering pt_expt + enable_compile=True + type_one_side=False so this combination is exercised in CI.

Until that is fixed, the practical workarounds are to disable compiled training (enable_compile=False) or use type_one_side=True when that is acceptable for the model.

Agent: ChatGPT
Model: GPT-5.6 Sol