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

[iduygnay]

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:

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_

[dosubot]

Hi @iduygnay! I'm Dosu and I'm helping the DeepModeling team.

This looks like a real bug in the multi-frame Python inference path for native-spin DPA4 graph .pt2 models. The corruption pattern you describe — frame 0 correct, subsequent frames producing NaNs, non-monotonic dependence on batch size — matches a class of per-frame offset bugs that have already been found and fixed on two other backends.

Closely related bugs already fixed

• #5660 / PR #5851 — The TensorFlow DeepSpin C++ path had exactly this symptom: extension buffers and output copyback only handled frame zero, so force, force_mag, atomic energy, and virial for frames beyond 0 were garbage. The fix added explicit frame * natoms * stride offsets throughout. [1]
• PR #5852 — The dpmodel DeepEval adapter had the same issue: the flat spin tensor was not being reshaped and sliced per frame under AutoBatchSize, so multi-frame spin evaluations silently used only frame 0's spin for every frame. [2]

Why the pt_expt graph-spin path is affected

The native-spin graph route was introduced in PR #5884. Its validation covered dpmodel parity, C++ gtests (single-frame), and LAMMPS single-rank runs, but did not include multi-frame Python DeepEval regressions for the graph-spin path [3]. The NeighborGraph batches multiple frames by stacking nodes/edges on a flat axis with n_node, destination_row_ptr, and source_row_ptr encoding per-frame offsets [4]. If those offsets are applied incorrectly to the spin tensor or edge indices when nframes > 1, the compiled AOTI kernel indexes into the wrong nodes — producing exactly the NaN/garbage pattern you see.

Additionally, #5965 identifies that the graph-native spin path is still missing forward_ragged / call_common_ragged seams and does not propagate the padding mask, which are related gaps in the same code area [5].

Most likely root cause

Based on the pattern and the analogous fixes in #5851 and #5852, the probable culprit is that the spin tensor fed to the AOTI runner has shape (nf * nloc, 3) but per-frame node-offset bookkeeping in n_node / row pointers isn't applied consistently to it — so frame N's descriptor reads frame 0's spin values, or edge src/dst indices from frame N point into frame 0's node range.

Debugging suggestions

1. Dump the NeighborGraph fields (n_node, destination_row_ptr, source_row_ptr, spin) for a 2-frame call and verify the second frame's slice starts at offset nloc, not 0.
2. Check whether the spin tensor passed to the AOTI runner has correct frame stacking (nf * nloc, 3) or is inadvertently (nloc, 3) from only frame 0.
3. The fix pattern from PR #5852 — reshaping and validating spin to (numb_test, natoms, 3) before AutoBatchSize slicing — is directly applicable here [2].

Your workaround of evaluating one frame at a time is correct. The fix belongs in the multi-frame NeighborGraph construction or spin-tensor slicing logic in deepmd/pt_expt/infer/deep_eval.py.

To reply, just mention @dosu.

---

Docs are dead. Just use Dosu.

[njzjz]

Thanks for the detailed report. I checked the exact f573ca045 code path and current master, and I think we can narrow this down further.

I do not think the earlier “missing per-frame offset / wrong spin slicing” diagnosis matches the pt_expt graph-spin code:

• neighbor_graph_from_ijs() explicitly converts frame-local indices to the flattened graph with i + nframe_id * nloc / j + nframe_id * nloc, and the CSR views are built over the full nf * nloc node axis.
• _eval_model_graph_spin() reshapes atype and spin onto that same flat node axis; in particular spin is passed as (nframes * natoms, 3). aparam, when present, is flattened the same way.
• _build_graph_dynamic_shapes(..., is_native_spin=True) marks nframes, the flat node count N, and edge count E as dynamic, and makes spin share the same dynamic N symbol as atype.
• More importantly, test_graph_lower_exportable_torch_export for native-spin DPA4 already builds its default sample with nframes=2, exports it with dynamic nf/N/E, and checks the exported program against eager execution. That path is finite and agrees with eager execution.

The coverage hole appears one layer later: AOTInductor packaging/execution. The end-to-end test_deep_eval_graph_spin_parity freezes a graph .pt2, loads it through DeepPot, and runs the AOTI package, but its _SPIN_EVAL_COORDS, cell, and spins are all hard-coded with leading dimension 1. That test is also skipped in CI because the AOTI compile is slow. I checked current master; it is still single-frame.

So the first layer that is both consistent with this report and not covered by the existing multi-frame tests is:

torch.export ExportedProgram (passes multi-frame) → aoti_compile_and_package / CUDA AOTInductor → AOTICompiledModel (only tested with one frame)

My current diagnosis is therefore a CUDA AOTInductor dynamic-shape miscompile/runtime bug in the frozen graph artifact, rather than a missing frame offset in DeepEval or NeighborGraph. The non-monotonic shape dependence (2/3/8/16 bad while 4 happened to work) also fits generated-kernel shape sensitivity better than a deterministic frame-offset error.

One additional clue is that energy itself is corrupted, together with force/virial/force_mag. That makes a bug confined to the final force/virial CSR scatter unlikely; the corruption must occur in the energy/descriptor path or an earlier compiled intermediate as well.

I would not claim the exact generated kernel yet without a reproducer. The clean isolation test is to run the same 1/2/3/4/8-frame graph inputs through three stages: (1) eager forward_common_lower_graph, (2) the pre-AOTI exported module, and (3) the packaged .pt2. If (1) and (2) remain finite while (3) fails, that pins the bug to Inductor codegen/runtime and gives us a small target for inspecting the generated kernel/configuration. This should also become a regression test: batched .pt2 results must match concatenated single-frame calls for native-spin DPA4.

If possible, could you attach the affected .pt2 (or a minimal checkpoint/config that reproduces it)? The current report is enough to localize the layer, but the model would let us identify the exact AOTI kernel/operation instead of guessing from static inspection.

For now, evaluating one frame per call is a valid workaround.

---

Agent: ChatGPT
Agent version: not exposed in this runtime
Model: GPT-5.6 Sol

[OutisLi]

I could not reproduce this on current master using the existing local DeepMD/PyTorch environment and CUDA hardware.

Environment used

• DeepMD-kit source: ed691aab147d9d7686d296e30f46902c08e9fb68 (current upstream/master when tested)
• PyTorch: 2.13.0+cu130
• CUDA toolkit used for the custom-op rebuild: 13.3
• GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition, 96 GB; the tests were bound to physical GPU 7
• DP_INFER_BATCH_SIZE=4096
• DeepPot(..., auto_batch_size=False), so every requested frame count stayed in one AOTI call

The installed libdeepmd_op_pt.so was older than the checked-out Python/CUDA source and initially failed during export because its edge_force_virial schema still contained an edge_spin_gradient argument. I rebuilt only the existing build-tree deepmd_op_pt target from current source and explicitly loaded that build-tree library. The verified schema was:

deepmd::edge_force_virial(
    Tensor edge_gradient,
    Tensor edge_vec,
    Tensor edge_index,
    Tensor edge_mask,
    Tensor destination_order,
    Tensor destination_row_ptr,
    Tensor source_order,
    Tensor source_row_ptr,
    Tensor n_node_per_frame,
    SymInt node_capacity,
    bool want_atom_virial
) -> (Tensor force, Tensor atom_virial, Tensor virial)

The stale-library failure is unrelated to this issue and was not used as evidence below.

Reproduction procedure

I built a jittered native-spin DPA4 model so that the model has non-trivial spin dependence and force_mag is not vacuously zero. The model used fp32 descriptor and fitting parameters with a default-scale architecture:

channels=64
n_radial=16
lmax=3
mmax=1
n_blocks=3
fitting=[64, 64]
precision=float32

I exported a graph-form CUDA AOTInductor .pt2 through deepmd.pt_expt.utils.serialization.deserialize_to_file(..., lower_kind="graph"), loaded it through the public DeepPot API, and compared each batched call against concatenated single-frame calls.

The test system contained 54 atoms per frame, matching the report. Every frame used distinct coordinates, cells, and spins. I tested nframes = 1, 2, 3, 4, 8.

All outputs were finite for every frame count:

output        maximum batched-vs-single difference    NaN count
energy        1.669e-6                                0
force         8.008e-6                                0
virial        5.216e-5                                0
force_mag     1.129e-7                                0
mask_mag      exact                                   0

These are ordinary finite fp32 reduction-order differences, not the reported NaNs or 1e7-1e9 corruption. No frame was wholly or partially corrupted.

Related paths checked

The previously fixed frame-offset bugs in #5660/#5851 and #5852 do not match this current path: the graph builder applies frame offsets, and spin is flattened onto the same N = nframes * natoms node axis as atype. #5965 concerns ragged training/padding loss and is also separate.

There is also a provenance inconsistency worth resolving. At both report commit f573ca045 and current master, the standard PT SeZM exporter used by dp --pt freeze reports lower_input_kind="edge_vec", whereas the issue says the artifact entered _eval_model_graph_spin, which requires lower_input_kind="graph". Please paste the artifact's model/extra/metadata.json or clarify whether it was actually frozen through the pt_expt graph exporter or a locally modified path.

Current assessment

I cannot confirm a DeepMD frame-offset or graph-construction bug with the available information. The report's non-monotonic shape dependence remains consistent with an environment-specific CUDA AOTInductor wrong-code/runtime problem. Since the same fp32 architecture and frame/atom shapes are stable on the RTX PRO 6000 Blackwell test host, a reporter-side GPU/driver/hardware issue is also plausible. This is a hypothesis, not a confirmed hardware diagnosis; the exact production model may still trigger generated code that this synthetic model does not.

The remaining major differences are:

• reporter: PyTorch 2.11.0+cu126; local test: 2.13.0+cu130
• reporter GPU model/driver/ECC state: not provided
• reporter's portable checkpoint, complete model config, and minimal failing input frames: not provided

AOTInductor .pt2 packages are tied to their build/runtime toolchain and target device, so moving the reported .pt2 to a different host is not a reliable reproduction method. To reproduce by rebuilding the artifact in a controlled environment, please provide:

portable training checkpoint
complete model/training configuration
minimal failing coordinates, cells, atom types, and spins
model/extra/metadata.json copied as text
GPU model and driver version
nvidia-smi -q
ECC/error counters and any Xid messages

The hardware/environment discriminator is to rebuild .pt2 packages independently from the same checkpoint and config on the two environments, then run the same input frames. The codegen discriminator in the reporter's own environment remains comparing the same graph inputs through eager, the pre-AOT exported module, and the locally packaged .pt2.