fix(pt_expt): trace graph lower on the model's device; fix CUDA test device placement

The graph-lower compile/export path traces the model on CPU (make_fx keeps real
parameters, and the export path moves the model to CPU to dodge a CUDA
autograd-stream limitation). _trace_and_compile_graph built its synthetic trace
inputs on the global DEVICE instead, so on a GPU host it mixed a CPU model with
CUDA inputs and failed make_fx tracing. Build the synthetic inputs on the model's
own device. (No effect on CUDA training, where the model already lives on DEVICE.)

The graph-export tests then run the exported program, which _trace_and_export
moves to env.DEVICE via move_to_device_pass, but fed it CPU inputs and compared
against a CPU eager model. Place the graph inputs (and the eager reference) on
env.DEVICE; the with-comm host control tensors stay on CPU by design. Also read
the CUDA repinit.mean stats buffer through .detach().cpu() before np.asarray, and
relax the escape-hatch force check to CUDA scatter-add reduction-order noise
(energy stays bit-exact). All CPU-only, so the CPU CI never exercised these paths.

Author: Han Wang

deepmd/pt_expt/train/training.py

@@ -691,12 +691,20 @@ def _trace_and_compile_graph(
     # sel-derived estimate overflows whenever the real degree exceeds sel),
     # then prime-padded to stay distinct from nf and N.  ``+ 2`` keeps at
     # least two masked padding rows so the padded-tail branch is traced.
+    # Trace on the MODEL's device, not the global ``DEVICE``: make_fx keeps the
+    # real model parameters (``_allow_non_fake_inputs``), so the synthetic trace
+    # inputs must live where the model does. A CUDA training run keeps the model
+    # on ``DEVICE`` (these match), but callers that trace a CPU-placed model
+    # (e.g. the graph .pt2/export path, which moves the model to CPU to dodge a
+    # CUDA autograd-stream limitation) would otherwise mix a CPU model with
+    # CUDA inputs and fail only on a GPU host.
+    _trace_device = next(model.parameters()).device
     e_real = count_synthetic_graph_edges(
         model,
         nframes=trace_nf,
         nloc=nloc_trace,
         dtype=GLOBAL_PT_FLOAT_PRECISION,
-        device=DEVICE,
+        device=_trace_device,
     )
     e_max = _next_safe_prime(e_real + 2, _forbidden | {trace_nf, trace_N})
     sample = build_synthetic_graph_inputs(
@@ -705,7 +713,7 @@ def _trace_and_compile_graph(
         nframes=trace_nf,
         nloc=nloc_trace,
         dtype=GLOBAL_PT_FLOAT_PRECISION,
-        device=DEVICE,
+        device=_trace_device,
         want_fparam=fparam is not None,
         want_aparam=aparam is not None,
         want_charge_spin=charge_spin is not None,

source/tests/pt_expt/model/test_dpa2_graph_lower.py

@@ -326,12 +326,19 @@ def test_disable_graph_lower_escape_hatch(self) -> None:
             box,
             neighbor_graph_method="legacy",
         )
-        # with the hatch on, default (None) == legacy (dense), bit-identical
+        # With the hatch on, default (None) takes the same dense path as legacy,
+        # so the energy is bit-identical. The force goes through a scatter/index_add
+        # in the backward, which is non-deterministic on CUDA across two separate
+        # passes, so it matches only to reduction-order noise (~1e-18), far below
+        # any real route divergence (>=1e-8, cf. test_binding_sel_diverges).
         torch.testing.assert_close(
             default_after["energy_redu"], legacy["energy_redu"], rtol=0, atol=0
         )
         torch.testing.assert_close(
-            default_after["energy_derv_r"], legacy["energy_derv_r"], rtol=0, atol=0
+            default_after["energy_derv_r"],
+            legacy["energy_derv_r"],
+            rtol=1e-10,
+            atol=1e-12,
         )
 
     def test_binding_sel_diverges(self) -> None:
@@ -745,7 +752,7 @@ def _build(set_davg_zero: bool) -> EnergyModel:
             "box": self.cell.clone(),
         }
         model.atomic_model.descriptor.compute_input_stats([sample])
-        davg = np.asarray(model.atomic_model.descriptor.repinit.mean)
+        davg = np.asarray(model.atomic_model.descriptor.repinit.mean.detach().cpu())
         assert np.abs(davg).max() > 0, "computed statistics must be nonzero"
 
         # 1. the gate: nonzero-mean configs stay on the legacy route.

source/tests/pt_expt/model/test_graph_export.py

@@ -160,6 +160,9 @@ def test_graph_export_aparam_flat_node_axis():
         lower_kind="graph",
     )
     loaded = exported.module()
+    # ``_trace_and_export`` moves the exported program to ``env.DEVICE``; run the
+    # eager reference there too so both sides use matching-device inputs.
+    model.to(env.DEVICE)
 
     # nf=1, N=1 (single node: zero real edges, guard rows only) and a
     # multi-frame nf=3, N=15 system: both must pass the input guards and
@@ -171,7 +174,7 @@ def test_graph_export_aparam_flat_node_axis():
             nframes=nframes,
             nloc=nloc,
             dtype=torch.float64,
-            device=torch.device("cpu"),
+            device=env.DEVICE,
         )
         (
             a2,
@@ -191,9 +194,9 @@ def test_graph_export_aparam_flat_node_axis():
         # single-rank runtime: every node is owned
         nl2 = nn2.clone()
         # distinct per-row values so the sensitivity check below is real
-        ap2 = torch.linspace(0.1, 0.9, ap2.numel(), dtype=torch.float64).reshape(
-            ap2.shape
-        )
+        ap2 = torch.linspace(
+            0.1, 0.9, ap2.numel(), dtype=torch.float64, device=ap2.device
+        ).reshape(ap2.shape)
         out = loaded(a2, nn2, nl2, ei2, ev2, em2, do2, drp2, so2, srp2, fp2, ap2, cs2)
         ref = model.forward_common_lower_graph(
             a2,

source/tests/pt_expt/model/test_graph_export_with_comm.py

@@ -16,6 +16,9 @@
 
 import pytest
 
+from deepmd.pt_expt.utils.env import (
+    DEVICE,
+)
 from deepmd.pt_expt.utils.serialization import (
     deserialize_to_file,
 )
@@ -239,15 +242,27 @@ def test_graph_with_comm_n_local_is_separate_device_input(
     atype = np.array([[i % 2 for i in range(n_total)]])
     graph = build_neighbor_graph(coord, atype, None, rcut, canonicalize=True)
 
-    atype_t = torch.tensor(atype.reshape(-1), dtype=torch.int64)
-    n_node_t = torch.as_tensor(np.asarray(graph.n_node), dtype=torch.int64)
-    ei = torch.as_tensor(np.asarray(graph.edge_index), dtype=torch.int64)
-    ev = torch.as_tensor(np.asarray(graph.edge_vec), dtype=torch.float64)
-    em = torch.as_tensor(np.asarray(graph.edge_mask), dtype=torch.bool)
-    do_t = torch.as_tensor(np.asarray(graph.destination_order), dtype=torch.int64)
-    drp_t = torch.as_tensor(np.asarray(graph.destination_row_ptr), dtype=torch.int64)
-    so_t = torch.as_tensor(np.asarray(graph.source_order), dtype=torch.int64)
-    srp_t = torch.as_tensor(np.asarray(graph.source_row_ptr), dtype=torch.int64)
+    # Graph inputs live on the device the exported program was moved to
+    # (env.DEVICE); the 8 comm tensors below stay on CPU (host control metadata).
+    atype_t = torch.tensor(atype.reshape(-1), dtype=torch.int64, device=DEVICE)
+    n_node_t = torch.as_tensor(
+        np.asarray(graph.n_node), dtype=torch.int64, device=DEVICE
+    )
+    ei = torch.as_tensor(np.asarray(graph.edge_index), dtype=torch.int64, device=DEVICE)
+    ev = torch.as_tensor(np.asarray(graph.edge_vec), dtype=torch.float64, device=DEVICE)
+    em = torch.as_tensor(np.asarray(graph.edge_mask), dtype=torch.bool, device=DEVICE)
+    do_t = torch.as_tensor(
+        np.asarray(graph.destination_order), dtype=torch.int64, device=DEVICE
+    )
+    drp_t = torch.as_tensor(
+        np.asarray(graph.destination_row_ptr), dtype=torch.int64, device=DEVICE
+    )
+    so_t = torch.as_tensor(
+        np.asarray(graph.source_order), dtype=torch.int64, device=DEVICE
+    )
+    srp_t = torch.as_tensor(
+        np.asarray(graph.source_row_ptr), dtype=torch.int64, device=DEVICE
+    )
 
     sendlist_indices = np.ascontiguousarray(
         np.arange(nghost, dtype=np.int32)
@@ -265,7 +280,7 @@ def test_graph_with_comm_n_local_is_separate_device_input(
     )
 
     def run(n_local_val: int) -> torch.Tensor:
-        n_local_t = torch.tensor([n_local_val], dtype=torch.int64)
+        n_local_t = torch.tensor([n_local_val], dtype=torch.int64, device=DEVICE)
         out = loaded(
             atype_t,
             n_node_t,