deepmodeling / deepmodeling/deepmd-kit

[BUG] Multi-frame DeepPot.eval produces NaNs and corrupted outputs for frozen native-spin DPA4 graph models

Open
#5,969 3 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

### Bug summary

### Issue body

## Summary

Multi-frame inference with a frozen native-spin DPA4 graph-form `.pt2` model produces NaNs and extremely large corrupted values.

Single-frame inference with the same model and the same input structures works correctly.

The affected public API is:

```python
from deepmd.infer import DeepPot

dp = DeepPot("dpa4_step128000.pt2")

energy, force, virial, force_mag, mask_mag = dp.eval(
coords, # (nframes, natoms, 3)
cells, # (nframes, 9)
atom_types, # (natoms,)
spin=spins, # (nframes, natoms, 3)
)

The issue appears when multiple frames are passed to DeepPot.eval() in a single call.

## Environment

- DeepMD-kit: 3.2.0b1.dev211+gf573ca045
- Backend: PyTorch
- PyTorch: 2.11.0+cu126
- Model type: DPA4 native-spin model
- Frozen model format: graph-form .pt2
- Number of atoms per frame: 54

The model was frozen using dp --pt freeze.

## Observed behavior

The input coordinates, cells, atom types, and spins contain no NaNs or infinite values.

I set DP_INFER_BATCH_SIZE=4096, which is large enough to keep the tested frames in a single internal inference batch.

The results depend on the number of frames passed in one call:

Frames in one eval() call Observed result
━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
1 All outputs are finite and correct
─────────────────────────── ─────────────────────────────────────────────────────────────────────────────
2 The second frame contains NaNs in energy, force, virial, and magnetic force
─────────────────────────── ─────────────────────────────────────────────────────────────────────────────
3 Two frames contain NaNs
─────────────────────────── ─────────────────────────────────────────────────────────────────────────────
4 All outputs happened to be finite in this test
─────────────────────────── ─────────────────────────────────────────────────────────────────────────────
8 NaNs and extremely large finite values appear
─────────────────────────── ─────────────────────────────────────────────────────────────────────────────
16 NaNs and extremely large finite values appear

For two frames, the NaN counts were:

Output Number of NaNs
━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━━━
Energy 1
──────────────── ────────────────
Force 162
──────────────── ────────────────
Virial 9
──────────────── ────────────────
Magnetic force 162
──────────────── ────────────────
Magnetic mask 0

Here, 162 equals 54 atoms × 3, and 9 corresponds to one full virial tensor. Therefore, one complete frame is corrupted.

For eight frames, I observed:

- 2 NaN energy values
- 801 NaN force components
- 45 NaN virial components
- 324 NaN magnetic-force components
- finite but unphysical energy values up to approximately 2.7e9 eV
- finite but unphysical force-like values of approximately 7.7e7

The magnetic mask remains finite, while the model-computed outputs are corrupted.

The behavior is shape-dependent: a four-frame call happened to be finite, so the problem does not occur monotonically for every batch size.

## Expected behavior

Calling DeepPot.eval() with multiple frames should give the same result as evaluating those frames individually and concatenating the results:

batched_result = dp.eval(
coords,
cells,
atom_types,
spin=spins,
)

should be numerically equivalent to:

single_frame_results = []

for i in range(len(coords)):
result = dp.eval(
coords[i:i + 1],
cells[i:i + 1],
atom_types,
spin=spins[i:i + 1],
)
single_frame_results.append(result)

## Reliable workaround

Explicitly evaluating one frame per call works correctly:

results = []

for i in range(len(coords)):
result = dp.eval(
coords[i:i + 1],
cells[i:i + 1],
atom_types,
spin=spins[i:i + 1],
)
results.append(result)

I tested 2,118 frames using this workaround. All frames completed successfully with finite outputs.

## Relevant inference path

Based on Python introspection, the multi-frame call follows this path:

DeepPot.eval
→ deepmd.pt_expt.infer.deep_eval.DeepEval.eval
→ _eval_model_spin
→ _eval_model_graph_spin
→ batched NeighborGraph construction
→ AOTICompiledModel

The .pt2 runner is:

torch.export.pt2_archive._package.AOTICompiledModel

The graph-form native-spin inference path constructs flattened batched graph inputs, including:

atype
n_node
n_local
edge_index
edge_vec
edge_mask
destination_order
destination_row_ptr
source_order
source_row_ptr
spin
fparam
aparam
charge_spin

Therefore, the issue appears to be localized to multi-frame inference through the graph-form native-spin .pt2 AOTI path.

It may involve either:

1. node/edge offsets or row pointers in the batched NeighborGraph, or
2. dynamic-shape handling or shape specialization in the AOT-compiled .pt2 model.

I have not yet isolated whether the corruption originates in graph construction or in the compiled AOTI model.

## Scope

This does not appear to indicate corrupted model parameters:

- every tested frame works when evaluated individually;
- the same input frame can fail in a multi-frame call and succeed in a single-frame call;
- all 2,118 frames were successfully evaluated using single-frame calls.

The issue has currently only been confirmed for the following combination:

- development version of DeepMD-kit;
- DPA4 native-spin model;
- graph-form frozen .pt2 model;
- DeepPot.eval() with more than one frame per call.

I have not confirmed whether the same issue affects non-spin DPA4 models, non-graph input formats, or unfrozen/raw checkpoints.

### DeePMD-kit Version

v3.2.0b1.dev211+gf573ca045

### Backend and its version

PyTorch 2.11.0+cu126

### How did you download the software?

Offline packages

### Input Files, Running Commands, Error Log, etc.

上面有

### Steps to Reproduce

上面有

### Further Information, Files, and Links

_No response_

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.