fix(pt_expt): derive graph trace edge_capacity from real edge count, not sel

The carry-all graph builder is sel-free (edges are cutoff-determined; sel
is only a normalization constant), so the sel-derived static trace
capacity ceil(1.25*nloc*sum(sel)) overflowed whenever the synthetic trace
system's actual degree exceeded sel: compiled training with
repinit/repformer nsel=10/6 raised 'edge overflow: 106 real edges >
edge_capacity 89'; graph .pt2 export with sel=2 raised 'edge overflow: 36
real edges > edge_capacity 18'.

The capacity now derives from the actual unpadded synthetic graph:
build_synthetic_graph_inputs accepts e_max=None (the dynamic probe
layout) and the new count_synthetic_graph_edges counts the real edges;
export pads 25% (min +2) and keeps the concrete edge length
duck-sizing-distinct from nframes/N, compiled training prime-pads it
collision-free as before. The value remains trace-sample-only: the
exported edge axis stays fully dynamic (Dim(nedge, min=2)) and compiled
training builds run-time graphs with no capacity.

Regressions: dpa2 nsel=10/6 compiled training; dpa1 sel=2 graph export.

Author: Han Wang

deepmd/pt_expt/train/training.py

@@ -679,8 +679,6 @@ def _trace_and_compile_graph(
         Per-task buffers promoted to FX placeholders (see
         :func:`_detect_task_buffers`).
     """
-    import math
-
     from torch._decomp import (
         get_decompositions,
     )
@@ -740,20 +738,30 @@ def _trace_and_compile_graph(
     while (trace_nf * nloc_trace) in (_forbidden | {trace_nf}):
         nloc_trace += 1
     trace_N = trace_nf * nloc_trace
-    # Static edge capacity, prime-padded to stay distinct from nf and N.
-    nnei = sum(model.get_sel())
-    e_max_base = max(math.ceil(1.25 * nloc_trace * nnei), 7)
-    e_max = _next_safe_prime(e_max_base, _forbidden | {trace_nf, trace_N})
-
     # Shared with the .pt2 export trace (serialization.py) so the two graph
     # traces can never desync on the input schema.  Training uses the run-time
     # float precision and device; optional tensors match the actual call.
     from deepmd.pt_expt.utils.serialization import (
         build_synthetic_graph_inputs,
         check_graph_trace_torch_version,
+        count_synthetic_graph_edges,
     )
 
     check_graph_trace_torch_version(model)
+
+    # Static edge capacity: derived from the ACTUAL edge count of the
+    # synthetic trace system (the carry-all builder is sel-free; a
+    # 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.
+    e_real = count_synthetic_graph_edges(
+        model,
+        nframes=trace_nf,
+        nloc=nloc_trace,
+        dtype=GLOBAL_PT_FLOAT_PRECISION,
+        device=_model_trace_device(model),
+    )
+    e_max = _next_safe_prime(e_real + 2, _forbidden | {trace_nf, trace_N})
     sample = build_synthetic_graph_inputs(
         model,
         e_max=e_max,

deepmd/pt_expt/utils/serialization.py

@@ -362,7 +362,7 @@ def _make_sample_inputs(
 
 def build_synthetic_graph_inputs(
     model: torch.nn.Module,
-    e_max: int,
+    e_max: int | None,
     nframes: int = 2,
     nloc: int = 7,
     *,
@@ -394,8 +394,13 @@ def build_synthetic_graph_inputs(
     ----------
     model : torch.nn.Module
         The pt_expt energy model (must expose ``get_rcut``/``get_type_map``/...).
-    e_max : int
-        Static edge capacity ``E`` to pad the (masked) edge axis to.
+    e_max : int or None
+        Static edge capacity ``E`` to pad the (masked) edge axis to.  Must be
+        at least the system's real edge count (the carry-all builder raises
+        ``edge overflow`` otherwise) — derive it from
+        :func:`count_synthetic_graph_edges`, never from ``sel`` (the builder
+        is sel-free).  ``None`` selects the dynamic layout (real edges plus
+        ``min_edges`` guard rows); used by the edge-count probe itself.
     nframes : int
         Number of frames in the sample system.
     nloc : int
@@ -471,6 +476,56 @@ def build_synthetic_graph_inputs(
     )
 
 
+def count_synthetic_graph_edges(
+    model: torch.nn.Module,
+    nframes: int,
+    nloc: int,
+    *,
+    dtype: torch.dtype,
+    device: torch.device | None = None,
+) -> int:
+    """Count the real (unpadded) edges of the synthetic trace system.
+
+    Probes :func:`build_synthetic_graph_inputs` with ``e_max=None`` (the
+    dynamic carry-all layout, whose edge axis is the real edge count plus
+    the ``min_edges`` guard rows) and counts the ``edge_mask`` real prefix.
+    The carry-all builder is sel-free — edges are cutoff-determined, ``sel``
+    is only a normalization constant — so the static trace ``edge_capacity``
+    must derive from this geometry-determined count; a sel-based estimate
+    overflows whenever the synthetic system's real degree exceeds ``sel``
+    (small-``sel`` models).
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        The pt_expt energy model (must expose ``get_rcut``/``get_type_map``).
+    nframes : int
+        Number of frames of the synthetic system; must match the subsequent
+        :func:`build_synthetic_graph_inputs` call.
+    nloc : int
+        Local atoms per frame; must match the subsequent call.
+    dtype : torch.dtype
+        Float precision of the probe coordinates; must match the subsequent
+        call (the edge count is cutoff-thresholded).
+    device : torch.device, optional
+        Probe device; must match the subsequent call.
+
+    Returns
+    -------
+    int
+        Number of real edges of the synthetic system at the model cutoff.
+    """
+    edge_mask = build_synthetic_graph_inputs(
+        model,
+        e_max=None,
+        nframes=nframes,
+        nloc=nloc,
+        dtype=dtype,
+        device=device,
+    )[4]
+    return int(edge_mask.sum())
+
+
 def _build_graph_dynamic_shapes(
     *sample_inputs: torch.Tensor | None,
 ) -> tuple:
@@ -1046,11 +1101,26 @@ def _trace_and_export(
         # The edge axis is DYNAMIC (B2.0): the AOTI artifact accepts any edge
         # count, so there is no capacity to bake. The trace sample is built at a
         # concrete, padded edge size only to keep the trace tensors distinct
-        # from the other dynamic dims (nframes=2, N=14) under torch.export's
-        # duck-sizing; the value itself does NOT constrain runtime.
+        # from the other dynamic dims (nframes, N) under torch.export's
+        # duck-sizing; the value itself does NOT constrain runtime.  The
+        # capacity derives from the ACTUAL edge count of the synthetic system
+        # (the carry-all builder is sel-free; a sel-derived estimate overflows
+        # for small-sel models): 25% headroom keeps the masked padded tail
+        # genuinely traced, the ``+ 2`` floor guarantees it even for tiny edge
+        # counts, and the final bump keeps the concrete edge length distinct
+        # from the other trace dims under duck-sizing.
         nloc_sample = 7
-        nnei = sum(model.get_sel())
-        e_sample = math.ceil(1.25 * nloc_sample * nnei)
+        nframes_sample = 1 if with_comm_dict else 2
+        e_real = count_synthetic_graph_edges(
+            model,
+            nframes=nframes_sample,
+            nloc=nloc_sample,
+            dtype=torch.float64,
+            device=torch.device("cpu"),
+        )
+        e_sample = max(math.ceil(1.25 * e_real), e_real + 2)
+        while e_sample in (nframes_sample, nframes_sample * nloc_sample):
+            e_sample += 1
 
         if with_comm_dict:
             # Load libdeepmd_op_pt.so and register border_op fake/autograd

source/tests/pt_expt/model/test_dpa2_graph_lower.py

@@ -552,6 +552,64 @@ def test_compiled_training_graph_smoke(self) -> None:
             **tol,
         )
 
+    def test_compiled_training_graph_small_sel(self) -> None:
+        """The compiled-training trace capacity derives from the synthetic
+        system's REAL edge count, not from ``sel``.
+
+        The carry-all graph builder is sel-free (sel = normalization
+        constant only), so a sel-derived static trace capacity
+        (``ceil(1.25 * nloc * sum(sel))``) overflows whenever the synthetic
+        trace system's actual degree exceeds ``sel``: with repinit/repformer
+        ``nsel=10/6`` the trace used to raise ``edge overflow: 106 real
+        edges > edge_capacity 89``.  The capacity is now probed from the
+        actual unpadded synthetic graph, so the trace must succeed and the
+        compiled lower must match the eager graph lower (compiled and eager
+        are the SAME route, so parity holds even at binding sel).
+        """
+        from deepmd.pt_expt.train.training import (
+            _trace_and_compile_graph,
+        )
+        from deepmd.pt_expt.utils.serialization import (
+            build_synthetic_graph_inputs,
+        )
+
+        model = self._make_model(repinit_nsel=10, repformer_nsel=6).to("cpu")
+        model.eval()
+
+        compiled_lower, _ = _trace_and_compile_graph(model, None, None, None)
+
+        sample = build_synthetic_graph_inputs(
+            model,
+            e_max=97,
+            nframes=3,
+            nloc=5,
+            dtype=torch.float64,
+            device=torch.device("cpu"),
+            want_fparam=False,
+            want_aparam=False,
+            want_charge_spin=False,
+        )
+        atype, n_node, ei, ev, em, fp, ap, cs = sample
+        compiled_out = compiled_lower(atype, n_node, ei, ev, em, fp, ap, cs)
+        eager = model.forward_common_lower_graph(
+            atype,
+            n_node,
+            ei,
+            ev,
+            em,
+            do_atomic_virial=False,
+            fparam=fp,
+            aparam=ap,
+            charge_spin=cs,
+        )
+        tol = {"rtol": 1e-10, "atol": 1e-10}
+        torch.testing.assert_close(compiled_out["energy"], eager["energy_redu"], **tol)
+        torch.testing.assert_close(
+            compiled_out["force"],
+            eager["energy_derv_r"].reshape(compiled_out["force"].shape),
+            **tol,
+        )
+
     def test_graph_lower_fparam_symbolic_trace_and_compile(self) -> None:
         """A graph-eligible DPA2 model with ``numb_fparam > 0`` must export
         (``make_fx`` symbolic) AND inductor-compile.

source/tests/pt_expt/utils/test_graph_pt2_metadata.py

@@ -49,16 +49,24 @@
 }
 
 
-def _build_dpa1_data() -> dict:
-    """Build a serialized dpmodel data dict for a dpa1(attn_layer=0) energy model."""
+def _build_dpa1_data(config: dict | None = None) -> dict:
+    """Build a serialized dpmodel data dict for a dpa1(attn_layer=0) energy model.
+
+    Parameters
+    ----------
+    config : dict, optional
+        Model config to build from.  Defaults to ``DPA1_CONFIG``.
+    """
     from deepmd.dpmodel.model.model import (
         get_model,
     )
 
-    model = get_model(copy.deepcopy(DPA1_CONFIG))
+    if config is None:
+        config = DPA1_CONFIG
+    model = get_model(copy.deepcopy(config))
     return {
         "model": model.serialize(),
-        "model_def_script": copy.deepcopy(DPA1_CONFIG),
+        "model_def_script": copy.deepcopy(config),
         "backend": "dpmodel",
         "software": "deepmd-kit",
         "version": "3.0.0",
@@ -94,6 +102,30 @@ def test_graph_pt2_has_lower_input_kind_graph(dpa1_dpmodel_data) -> None:
     assert "edge_capacity" not in meta
 
 
+def test_graph_pt2_small_sel_exports() -> None:
+    """Graph-form ``.pt2`` export succeeds for a small-``sel`` model.
+
+    The graph trace capacity derives from the synthetic trace system's
+    REAL edge count; the former sel-derived estimate
+    (``ceil(1.25 * nloc * sum(sel))``) overflowed the sel-free carry-all
+    builder whenever the actual degree exceeded ``sel`` (``edge overflow:
+    36 real edges > edge_capacity 18`` at ``sel=2``).
+    """
+    cfg = copy.deepcopy(DPA1_CONFIG)
+    cfg["descriptor"]["sel"] = 2
+    data = _build_dpa1_data(cfg)
+    with tempfile.TemporaryDirectory() as d:
+        p = os.path.join(d, "m_graph_small_sel.pt2")
+        deserialize_to_file(
+            p,
+            data,
+            do_atomic_virial=True,
+            lower_kind="graph",
+        )
+        meta = _read_metadata(p)
+    assert meta["lower_input_kind"] == "graph"
+
+
 def test_dense_pt2_has_lower_input_kind_nlist(dpa1_dpmodel_data) -> None:
     """Default (``lower_kind="nlist"``) -> metadata ``lower_input_kind == "nlist"``."""
     with tempfile.TemporaryDirectory() as d: