deepmodeling / deepmodeling/deepmd-kit
[BUG] pt-expt DPA1 graph lower fails with f_use_norm on NPY data
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Bug summary
PT-experimental training fails when all of the following are used:
- a graph-eligible DPA1 descriptor (`attn_layer: 0` is sufficient),
- a legacy NPY training system,
- energy loss with `f_use_norm: true`.
The graph lower returns force with shape `(nf, nloc, 3)`, while `DeepmdDataSystem` supplies NPY force labels with shape `(nf, nloc * 3)`. Most force-loss paths work because `EnergyLoss.call` first flattens both tensors into `diff_f`. The `f_use_norm` paths instead subtract the original tensors before reshaping, which raises a shape mismatch.
This reproduces with eager PT-experimental execution; `torch.compile` is not involved.
## DeePMD-kit Version
`3.2.0b1.dev110+g8399520c5.d20260704`, source checkout commit `0210cdb0edd5`.
## Backend and its version
PyTorch Experimental (`pt_expt`), PyTorch `2.11.0+cu128`, Python `3.13.13`.
## How did you download the software?
Built from source.
## Input Files, Running Commands, Error Log, etc.
Save the following as `input.json` at the repository root. It uses the bundled `examples/water/data/data_0` NPY system.
```json
{
"model": {
"type_map": ["O", "H"],
"descriptor": {
"type": "dpa1",
"sel": 12,
"rcut_smth": 0.5,
"rcut": 3.0,
"neuron": [8, 16],
"axis_neuron": 4,
"attn_layer": 0,
"precision": "float64",
"seed": 1
},
"fitting_net": {
"neuron": [16, 16],
"precision": "float64",
"seed": 1
},
"data_stat_nbatch": 1
},
"learning_rate": {
"type": "exp",
"start_lr": 0.001,
"stop_lr": 1e-6
},
"loss": {
"type": "ener",
"loss_func": "mae",
"f_use_norm": true,
"start_pref_e": 0,
"limit_pref_e": 0,
"start_pref_f": 1,
"limit_pref_f": 1,
"start_pref_v": 0,
"limit_pref_v": 0
},
"training": {
"training_data": {
"systems": ["examples/water/data/data_0"],
"batch_size": 1
},
"numb_steps": 1,
"disp_freq": 1,
"save_freq": 100
}
}
```
Run:
```bash
OMP_NUM_THREADS=1 \
DP_INTER_OP_PARALLELISM_THREADS=0 \
DP_INTRA_OP_PARALLELISM_THREADS=0 \
dp --pt-expt train input.json --skip-neighbor-stat
```
Observed shapes immediately before loss evaluation:
```text
force label: (1, 576)
graph prediction: (1, 192, 3)
```
Error:
```text
Traceback (most recent call last):
...
File "deepmd/dpmodel/loss/ener.py", line 445, in call
diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3))
RuntimeError: The size of tensor a (576) must match the size of tensor b (3) at non-singleton dimension 2
```
Expected behavior: `f_use_norm` accepts force labels from every supported data backend and the training step completes.
## Steps to Reproduce
1. Check out commit `0210cdb0edd5` or a revision containing the PT-experimental DPA1 graph lower.
2. Save the input above as `input.json` in the repository root.
3. Run the command above.
4. The first eager training step fails before the optimizer step.
## Further Information, Files, and Links
The mismatch is localized to force-difference construction:
- `deepmd/dpmodel/loss/ener.py` first creates layout-independent `diff_f` by flattening `force_hat` and `force`.
- The `f_use_norm` branches later recompute `force_hat - force` using the incompatible original layouts (masked/unmasked MAE and Huber paths).
- `deepmd/pt_expt/train/training.py::_CompiledModel._forward_graph` and the eager graph path expose force as `(nf, nloc, 3)`.
- `deepmd/dpmodel/utils/batch.py::normalize_batch` preserves the legacy flat NPY label shape.
- LMDB does not reproduce because `LmdbDataReader` normalizes each frame's force to `(nloc, 3)` before collation.
A structural fix is to derive every norm-based force difference from the already flattened `diff_f`, reshaping it to `(_nf, _nloc, 3)` where required. A regression test should cover both eager and compiled PT-experimental graph training on an NPY system with `f_use_norm: true`.
Contributor guide
Assessment
This issue has not been assessed yet.