feat(pt-expt): convert backend use auto to choose graph lower automatically

OutisLi · OutisLi · commit 6373f244eabe · 2026-07-09T17:51:32.000+08:00
diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py
@@ -36,19 +36,23 @@ def convert_backend(
     inp_hook = inp_backend.serialize_hook
     out_hook = out_backend.deserialize_hook
     data = inp_hook(INPUT)
-    # Forward atomic_virial to pt_expt deserialize_to_file if applicable;
-    # warn and skip the flag for backends that don't accept it so that
-    # scripts passing --atomic-virial indiscriminately don't break.
+    # Pass the optional flags a backend's deserialize hook accepts, probed by
+    # signature so scripts stay backend-agnostic. ``lower_kind="auto"`` lets the
+    # pt_expt hook pick the deploy-optimal lowering (the graph lower with its
+    # fused operator for a graph-eligible .pt2, the dense nlist lower otherwise);
+    # ``do_atomic_virial`` forwards ``--atomic-virial``.
     import inspect
 
     sig = inspect.signature(out_hook)
+    hook_kwargs: dict[str, Any] = {}
+    if "lower_kind" in sig.parameters:
+        hook_kwargs["lower_kind"] = "auto"
     if "do_atomic_virial" in sig.parameters:
-        out_hook(OUTPUT, data, do_atomic_virial=atomic_virial)
-    else:
-        if atomic_virial:
-            log.warning(
-                "--atomic-virial is only meaningful for pt_expt .pt2/.pte "
-                "outputs; ignoring it for output backend %s",
-                out_backend.name,
-            )
-        out_hook(OUTPUT, data)
+        hook_kwargs["do_atomic_virial"] = atomic_virial
+    elif atomic_virial:
+        log.warning(
+            "--atomic-virial is only meaningful for pt_expt .pt2/.pte "
+            "outputs; ignoring it for output backend %s",
+            out_backend.name,
+        )
+    out_hook(OUTPUT, data, **hook_kwargs)
diff --git a/deepmd/kernels/cuda/dpa1/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py
@@ -789,6 +789,19 @@ def ensure_registered() -> None:
     _cpu_library.impl("dpa1_graph_compress", _cpu_forward, "CPU")
     _cpu_library.impl("dpa1_graph_compress_backward", _cpu_backward, "CPU")
     if ef_op_available():
+        # The fused energy-force CPU reference (_ef_cpu, used for the freeze
+        # pipeline's trace-time sample evaluation) dispatches through the fitting
+        # and force / virial sub-operators, so their CPU implementations must be
+        # registered too.
+        from deepmd.kernels.cuda.edge_force_virial import (
+            ensure_registered as ensure_force_registered,
+        )
+        from deepmd.kernels.cuda.graph_fitting import (
+            ensure_registered as ensure_fitting_registered,
+        )
+
+        ensure_fitting_registered()
+        ensure_force_registered()
         torch.library.register_fake("deepmd::dpa1_graph_compress_energy_force")(
             _ef_fake
         )
diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py
@@ -1,12 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Compress a pt_expt model (.pte) by tabulating embedding nets."""
 
-import contextlib
 import logging
-import os
-from collections.abc import (
-    Iterator,
-)
 
 from deepmd.pt_expt.utils.serialization import (
     deserialize_to_file,
@@ -115,37 +110,25 @@ def enable_compression(
     #
     # A geometrically compressed graph-lower descriptor (DPA1 / se_atten strip,
     # attn_layer == 0) evaluates its tabulated embedding through fused CUDA
-    # operators, so its compressed graph exports directly and the ``.pt2`` runs
-    # the table kernel at inference. The trace forces ``DP_CUDA_INFER >= 2`` so
-    # the end-to-end ``deepmd::dpa1_graph_compress_energy_force`` operator is
-    # used: it returns the force as a value and stays opaque through
-    # ``torch.export``, whereas the level-1 autograd lower is decomposed to aten
-    # (the tabulated kernel is lost). Level 2 degrades to the level-1 fused table
-    # operator when the mega operator is unavailable or the descriptor width is
-    # not an instantiated kernel width; level 0 (the untraceable reference
-    # tabulation) is never used here.
-    #
-    # Every other compressed descriptor keeps tabulated operators that make_fx
-    # cannot trace: those export the UNCOMPRESSED graph and carry the compressed
-    # dict in ``model.json`` so ``deserialize()`` restores the compression state
-    # for the Python inference path.
+    # operators that make_fx can trace, so its compressed graph exports directly
+    # to a graph-lower ``.pt2``: ``deserialize_to_file`` bakes the end-to-end
+    # fused table operator and the mandatory per-atom virial (see its docstring).
+    # Every other compressed descriptor keeps tabulated operators make_fx cannot
+    # trace: those export the UNCOMPRESSED graph and carry the compressed dict in
+    # ``model.json`` so ``deserialize()`` restores the compression state for the
+    # Python inference path.
     from deepmd.pt_expt.train.training import (
         _model_uses_graph_lower,
     )
 
     model_def_script = model_dict.get("model_def_script")
     if output.endswith(".pt2") and _model_uses_graph_lower(model):
         log.info("Re-exporting compressed graph...")
-        with _cuda_infer_at_least_2():
-            deserialize_to_file(
-                output,
-                {"model": compressed_model_dict, "model_def_script": model_def_script},
-                lower_kind="graph",
-                # The per-atom virial is required by consumers that always read
-                # it (the Kokkos ``pair_style deepmd/kk``); match the uncompressed
-                # graph freeze, which also exports it.
-                do_atomic_virial=True,
-            )
+        deserialize_to_file(
+            output,
+            {"model": compressed_model_dict, "model_def_script": model_def_script},
+            lower_kind="graph",
+        )
     else:
         log.info("Re-exporting compressed model...")
         deserialize_to_file(
@@ -158,29 +141,3 @@ def enable_compression(
             },
         )
     log.info("Compressed model saved to %s", output)
-
-
-@contextlib.contextmanager
-def _cuda_infer_at_least_2() -> Iterator[None]:
-    """Pin ``DP_CUDA_INFER`` to at least 2 for the duration of a trace.
-
-    Level 2 selects the end-to-end fused table operator, which stays opaque
-    through ``torch.export``; the level-1 autograd lower is decomposed to aten
-    (losing the tabulated kernel), and level 0 selects the untraceable
-    reference tabulation. Level 2 degrades to level 1 internally when the mega
-    operator is unavailable or ineligible, so it is a safe floor.
-    """
-    from deepmd.kernels.utils import (
-        cuda_infer_level,
-    )
-
-    saved = os.environ.get("DP_CUDA_INFER")
-    if cuda_infer_level() < 2:
-        os.environ["DP_CUDA_INFER"] = "2"
-    try:
-        yield
-    finally:
-        if saved is None:
-            os.environ.pop("DP_CUDA_INFER", None)
-        else:
-            os.environ["DP_CUDA_INFER"] = saved
diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py
@@ -1,7 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import contextlib
 import ctypes
 import json
 import logging
+import os
+from collections.abc import (
+    Iterator,
+)
 from typing import (
     Any,
 )
@@ -807,6 +812,56 @@ def _serialize_from_file_pt2(model_file: str) -> dict:
     return model_dict
 
 
+@contextlib.contextmanager
+def _cuda_infer_at_least_2() -> Iterator[None]:
+    """Pin ``DP_CUDA_INFER`` to at least 2 for the duration of a trace.
+
+    Level 2 bakes the end-to-end fused operator, which stays opaque through
+    ``torch.export``; the level-1 autograd lower is decomposed to aten (losing
+    the hand-tuned kernel), and level 0 selects the untraceable reference
+    tabulation. Level 2 degrades to level 1 internally when the fused operator
+    is unavailable or ineligible, so it is a safe floor for any graph export.
+    """
+    from deepmd.kernels.utils import (
+        cuda_infer_level,
+    )
+
+    saved = os.environ.get("DP_CUDA_INFER")
+    if cuda_infer_level() < 2:
+        os.environ["DP_CUDA_INFER"] = "2"
+    try:
+        yield
+    finally:
+        if saved is None:
+            os.environ.pop("DP_CUDA_INFER", None)
+        else:
+            os.environ["DP_CUDA_INFER"] = saved
+
+
+def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
+    """Resolve ``lower_kind="auto"`` to a concrete lower-forward schema.
+
+    ``"auto"`` selects the graph lower for a graph-lower-eligible model exported
+    to ``.pt2`` -- the deploy-optimal path, whose fused mega operator the graph
+    export bakes -- and the dense nlist lower for everything else (a ``.pte``
+    target, a spin model, or a descriptor without a graph lower). An explicit
+    ``"nlist"`` / ``"graph"`` is returned unchanged.
+    """
+    if lower_kind != "auto":
+        return lower_kind
+    if not model_file.endswith(".pt2") or data["model"].get("type") == "spin_ener":
+        return "nlist"
+    from deepmd.pt_expt.model.model import (
+        BaseModel,
+    )
+    from deepmd.pt_expt.train.training import (
+        _model_uses_graph_lower,
+    )
+
+    model = BaseModel.deserialize(data["model"])
+    return "graph" if _model_uses_graph_lower(model) else "nlist"
+
+
 def deserialize_to_file(
     model_file: str,
     data: dict,
@@ -835,25 +890,42 @@ def deserialize_to_file(
         tracing the uncompressed model (make_fx cannot trace custom ops).
     do_atomic_virial : bool
         If True, export with per-atom virial correction (3 extra backward
-        passes, ~2.5x slower).  Default False for best performance.
+        passes, ~2.5x slower).  Default False for best performance. Forced True
+        for a graph lower, whose LAMMPS Kokkos consumer always reads it.
     lower_kind : str
         Which lower-forward schema the compiled AOTI graph consumes:
         ``"nlist"`` (default) traces the dense quartet
         (``extended_coord``/``extended_atype``/``nlist``/``mapping``);
         ``"graph"`` traces the NeighborGraph schema
         (``atype``/``n_node``/``edge_index``/``edge_vec``/``edge_mask``) with a
         DYNAMIC edge axis ``E`` (``Dim("nedge", min=2)``), so the artifact
-        accepts any system size.  The selected schema is recorded as
-        ``lower_input_kind`` in ``metadata.json``.
+        accepts any system size.  ``"auto"`` (used by ``convert-backend``)
+        resolves to ``"graph"`` for a graph-eligible ``.pt2`` and ``"nlist"``
+        otherwise (see :func:`_resolve_lower_kind`). A graph lower always bakes
+        the fused mega operator (``DP_CUDA_INFER >= 2``) and the per-atom virial.
+        The selected schema is recorded as ``lower_input_kind`` in
+        ``metadata.json``.
     """
-    if model_file.endswith(".pt2"):
-        _deserialize_to_file_pt2(
-            model_file, data, model_json_override, do_atomic_virial, lower_kind
-        )
+    lower_kind = _resolve_lower_kind(model_file, data, lower_kind)
+    # A graph lower deploys the fused mega operator: the trace must run at
+    # DP_CUDA_INFER >= 2 (level 1 decomposes to aten under torch.export), and the
+    # per-atom virial is mandatory for its LAMMPS Kokkos consumer. Enforcing both
+    # here gives every producer -- freeze, compress, convert-backend -- the same
+    # deployable artifact.
+    if lower_kind == "graph":
+        do_atomic_virial = True
+        ctx: contextlib.AbstractContextManager = _cuda_infer_at_least_2()
     else:
-        _deserialize_to_file_pte(
-            model_file, data, model_json_override, do_atomic_virial, lower_kind
-        )
+        ctx = contextlib.nullcontext()
+    with ctx:
+        if model_file.endswith(".pt2"):
+            _deserialize_to_file_pt2(
+                model_file, data, model_json_override, do_atomic_virial, lower_kind
+            )
+        else:
+            _deserialize_to_file_pte(
+                model_file, data, model_json_override, do_atomic_virial, lower_kind
+            )
 
 
 def _trace_and_export(
