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

Start at DeepPot.eval and follow the mentioned path through deepmd.pt_expt.infer.deep_eval.DeepEval.eval, _eval_model_spin, and _eval_model_graph_spin. Reproduce batched versus single-frame calls and inspect the flattened NeighborGraph inputs and AOTICompiledModel boundary. Done means multi-frame outputs are finite and numerically equivalent to concatenated single-frame results.

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_