feat(dpmodel,pt_expt,cc): graph-native DPA-2 — first message-passing model on the NeighborGraph lower (PR-G-dpa2)

.github/workflows/test_cuda.yml

@@ -17,6 +17,9 @@ jobs:
   test_cuda:
     name: Test Python and C++ on CUDA
     runs-on: gpu
+    # The full CUDA suite (serial pytest ~3h43m + C++ + LAMMPS) exceeds the
+    # default 360-min job limit; the self-hosted GPU runner has no hard cap.
+    timeout-minutes: 480
     # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
     # container:
     #   image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -381,6 +381,7 @@
         fparam: Array | None = None,
         aparam: Array | None = None,
         charge_spin: Array | None = None,
+        comm_dict: dict | None = None,
     ) -> dict:
         """Graph analogue of :meth:`forward_common_atomic` on the flat node axis.
 
@@ -406,6 +407,11 @@
         charge_spin
             charge/spin conditioning. Unused by the dpa1 graph path; accepted so
             the interface stays stable for charge/spin-conditioned descriptors.
+        comm_dict
+            MPI communication metadata forwarded to :meth:`forward_atomic_graph`
+            (and, from there, the descriptor's ``call_graph``). ``None`` for
+            non-parallel inference (default). Mirrors :meth:`forward_common_atomic`'s
+            ``comm_dict`` on the dense route.
 
         Returns
         -------
@@ -420,6 +426,7 @@
             fparam=fparam,
             aparam=aparam,
             charge_spin=charge_spin,
+            comm_dict=comm_dict,
         )
         return self._finalize_atomic_ret(ret_dict, output_mask, atype)
 
@@ -551,16 +558,16 @@
             out = ret_dict[kk]
             # explicit trailing product (NOT -1): a zero-atom forward (nloc==0)
             # has size 0, and numpy cannot infer -1 for a size-0 array.
             trail = math.prod(out.shape[len(lead) :])
             flat = xp.reshape(out, (*lead, trail))
             flat = xp.where(atom_mask[..., None], flat, xp.zeros_like(flat))
             ret_dict[kk] = xp.reshape(flat, out.shape)
         ret_dict["mask"] = xp.astype(atom_mask, xp.int32)
         return ret_dict
 
     def call(
         self,
         extended_coord: Array,
         extended_atype: Array,
         nlist: Array,
         mapping: Array | None = None,

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -301,6 +301,7 @@ def forward_atomic_graph(
         fparam: Array | None = None,
         aparam: Array | None = None,
         charge_spin: Array | None = None,
+        comm_dict: dict | None = None,
     ) -> dict[str, Array]:
         """Graph analogue of :meth:`forward_atomic` on the flat node axis.
 
@@ -322,6 +323,11 @@ def forward_atomic_graph(
         charge_spin
             charge/spin conditioning. Unused by the dpa1 graph path; accepted so
             the interface stays stable for charge/spin-conditioned descriptors.
+        comm_dict
+            MPI communication metadata forwarded to the descriptor's
+            ``call_graph`` (the message-passing part). ``None`` for
+            non-parallel inference (default). Mirrors :meth:`forward_atomic`'s
+            ``comm_dict`` on the dense route.
 
         Returns
         -------
@@ -338,14 +344,16 @@ def forward_atomic_graph(
         xp = array_api_compat.array_namespace(graph.edge_vec)
         type_embedding = self.descriptor.type_embedding.call()
         gg, rot_mat = self.descriptor.call_graph(
-            graph, atype, type_embedding=type_embedding
+            graph, atype, type_embedding=type_embedding, comm_dict=comm_dict
         )
         fparam_node = None
         if fparam is not None:
-            frame_id = frame_id_from_n_node(
-                graph.n_node,
-                n_total=atype.shape[0],
-            )
+            # Pass the STATIC flat node count (``atype.shape[0] == N``) so the
+            # helper does not fall back to ``int(sum(n_node))``: that int() on a
+            # traced tensor breaks make_fx / torch.export
+            # (``GuardOnDataDependentSymNode``) for the graph .pt2 export and
+            # compiled-training paths when ``numb_fparam > 0``.
+            frame_id = frame_id_from_n_node(graph.n_node, n_total=atype.shape[0])
             fparam_node = xp.take(fparam, frame_id, axis=0)  # (N, ndf)
         aparam_node = aparam
         if aparam is not None and graph.n_local is not None and aparam.ndim == 3:

deepmd/dpmodel/descriptor/dpa1.py

@@ -459,6 +459,23 @@
             )
         return self.se_atten.tebd_input_mode in ("concat", "strip")
 
+    def uses_compact_edge_pairs(self) -> bool:
+        """Returns whether the graph lower traces compact edge pairs.
+
+        The transformer attention lower (``attn_layer > 0``) enumerates
+        neighbor pairs via the compact ``center_edge_pairs`` realization
+        (unbacked-SymInt ``nonzero``/``repeat`` sizes, ``pairs.py``);
+        the factorizable lower (``attn_layer == 0``) traces with backed
+        symbols only.  ``check_graph_trace_torch_version`` keys its
+        torch >= 2.6 requirement on this capability.
+
+        Returns
+        -------
+        bool
+            Whether tracing :meth:`call_graph` runs ``center_edge_pairs``.
+        """
+        return self.se_atten.attn_layer > 0
+
     def disable_graph_lower(self) -> None:
         """Force the legacy dense lower for this descriptor.
 
@@ -540,16 +557,16 @@
         )
         remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map)
         obj = self.se_atten
         obj.ntypes = len(type_map)
         self.type_map = type_map
         self.type_embedding.change_type_map(type_map=type_map)
         obj.reinit_exclude(map_pair_exclude_types(obj.exclude_types, remap_index))
         if has_new_type:
             # the avg and std of new types need to be updated
             extend_descrpt_stat(
                 obj,
                 type_map,
                 des_with_stat=model_with_new_type_stat.se_atten
                 if model_with_new_type_stat is not None
                 else None,
             )
@@ -781,6 +798,7 @@
         atype: Array,
         type_embedding: Array | None = None,
         static_nnei: int | None = None,
+        comm_dict: dict | None = None,
     ) -> tuple[Array, Array]:
         """Descriptor-level graph-native forward.
 
@@ -817,6 +835,12 @@
             (N,) flat LOCAL atom types where ``N = sum(n_node)``.
         type_embedding
             (ntypes_with_padding, tebd_dim) type-embedding table.
+        comm_dict
+            MPI communication metadata. Accepted for ABI parity with
+            :meth:`DescrptDPA2.call_graph` (uniform ``forward_atomic_graph``
+            call site), but UNUSED: a single se_atten descriptor has no
+            cross-rank message passing (``has_message_passing_across_ranks()``
+            is ``False``), so this is always ``None`` in practice.
 
         Returns
         -------
@@ -2207,13 +2231,13 @@
         self.temperature = temperature
         self.trainable_ln = trainable_ln
         self.ln_eps = ln_eps
         self.smooth = smooth
         self.precision = precision
         self.network_type = NeighborGatedAttentionLayer
 
         self.attention_layers = [
             NeighborGatedAttentionLayer(
                 nnei,
                 embed_dim,
                 hidden_dim,
                 dotr=dotr,
@@ -2352,13 +2376,13 @@
         self.normalize = normalize
         self.temperature = temperature
         self.trainable_ln = trainable_ln
         self.ln_eps = ln_eps
         self.precision = precision
         self.attention_layer = GatedAttentionLayer(
             nnei,
             embed_dim,
             hidden_dim,
             dotr=dotr,
             do_mask=do_mask,
             scaling_factor=scaling_factor,
             normalize=normalize,
@@ -2500,14 +2524,14 @@
         self.temperature = temperature
         self.precision = precision
         self.scaling = (
             (self.head_dim * scaling_factor) ** -0.5
             if temperature is None
             else temperature
         )
         self.normalize = normalize
         self.in_proj = NativeLayer(
             embed_dim,
             hidden_dim * 3,
             bias=bias,
             use_timestep=False,
             precision=precision,