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

Activity

1. iduygnay

   added

   bug

   on Aug 14, 2026
2. dosubot commented on Aug 14, 2026

   

   dosubotBot

   on Aug 14, 2026 – with DosuBot

   More actions

   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.

   
3. njzjz commented on Aug 14, 2026

   

   njzjz

   on Aug 14, 2026 – with ChatGPT Codex Connector

   Member

   More actions

   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
4. OutisLi commented on Aug 17, 2026

   

   OutisLi

   on Aug 17, 2026

   Last edited by OutisLi

   Collaborator

   More actions

   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.