perf(pt2): optimize .pt2 C++ inference path

deepmd/entrypoints/convert_backend.py

@@ -12,6 +12,7 @@ def convert_backend(
     *,  # Enforce keyword-only arguments
     INPUT: str,
     OUTPUT: str,
+    atomic_virial: bool = False,
     **kwargs: Any,
 ) -> None:
     """Convert a model file from one backend to another.
@@ -20,12 +21,26 @@ def convert_backend(
     ----------
     INPUT : str
         The input model file.
-    INPUT : str
+    OUTPUT : str
         The output model file.
+    atomic_virial : bool
+        If True, export .pt2/.pte models with per-atom virial correction.
+        This adds ~2.5x inference cost.  Default False.
     """
     inp_backend: Backend = Backend.detect_backend_by_model(INPUT)()
     out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)()
     inp_hook = inp_backend.serialize_hook
     out_hook = out_backend.deserialize_hook
     data = inp_hook(INPUT)
-    out_hook(OUTPUT, data)
+    # Forward atomic_virial to pt_expt deserialize_to_file if applicable
+    import inspect
+
+    sig = inspect.signature(out_hook)
+    if "do_atomic_virial" in sig.parameters:
+        out_hook(OUTPUT, data, do_atomic_virial=atomic_virial)
+    else:
+        if atomic_virial:
+            raise ValueError(
+                "--atomic-virial is only supported for pt_expt .pt2/.pte outputs"
+            )
+        out_hook(OUTPUT, data)

deepmd/main.py

@@ -915,6 +915,14 @@ def main_parser() -> argparse.ArgumentParser:
     )
     parser_convert_backend.add_argument("INPUT", help="The input model file.")
     parser_convert_backend.add_argument("OUTPUT", help="The output model file.")
+    parser_convert_backend.add_argument(
+        "--atomic-virial",
+        action="store_true",
+        default=False,
+        help="Export .pt2/.pte models with per-atom virial correction. "
+        "This adds ~2.5x inference cost but is required for "
+        "LAMMPS compute/atom virial output.",
+    )
 
     # * show model ******************************************************************
     parser_show = subparsers.add_parser(

deepmd/pt_expt/model/make_model.py

@@ -346,6 +346,21 @@ def fn(
                 aparam: torch.Tensor | None,
             ) -> dict[str, torch.Tensor]:
                 extended_coord = extended_coord.detach().requires_grad_(True)
+                # Pad nlist with one extra -1 column inside the traced function.
+                # This ensures n_nnei > sum(sel), forcing the sort branch in
+                # _format_nlist.  The padding becomes part of the compiled graph,
+                # so callers never need to pad externally.
+                nlist = torch.cat(
+                    [
+                        nlist,
+                        -torch.ones(
+                            (*nlist.shape[:2], 1),
+                            dtype=nlist.dtype,
+                            device=nlist.device,
+                        ),
+                    ],
+                    dim=-1,
+                )
                 return model.forward_common_lower(
                     extended_coord,
                     extended_atype,
@@ -356,13 +371,19 @@ def fn(
                     do_atomic_virial=do_atomic_virial,
                 )
 
-            return make_fx(fn, **make_fx_kwargs)(
-                extended_coord,
-                extended_atype,
-                nlist,
-                mapping,
-                fparam,
-                aparam,
-            )
+            # Force format_nlist to always use the sort branch during tracing.
+            model.need_sorted_nlist_for_lower = lambda: True
+            try:
+                traced = make_fx(fn, **make_fx_kwargs)(
+                    extended_coord,
+                    extended_atype,
+                    nlist,
+                    mapping,
+                    fparam,
+                    aparam,
+                )
+            finally:
+                del model.need_sorted_nlist_for_lower
+            return traced
 
     return CM

deepmd/pt_expt/model/spin_model.py

@@ -96,6 +96,18 @@ def fn(
             aparam: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
             extended_coord = extended_coord.detach().requires_grad_(True)
+            # Pad nlist inside traced function (see make_model.py for rationale).
+            nlist = torch.cat(
+                [
+                    nlist,
+                    -torch.ones(
+                        (*nlist.shape[:2], 1),
+                        dtype=nlist.dtype,
+                        device=nlist.device,
+                    ),
+                ],
+                dim=-1,
+            )
             return model.forward_common_lower(
                 extended_coord,
                 extended_atype,
@@ -107,15 +119,22 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn, **make_fx_kwargs)(
-            extended_coord,
-            extended_atype,
-            extended_spin,
-            nlist,
-            mapping,
-            fparam,
-            aparam,
-        )
+        # Force format_nlist to always use the sort branch during tracing.
+        backbone = model.backbone_model
+        backbone.need_sorted_nlist_for_lower = lambda: True
+        try:
+            traced = make_fx(fn, **make_fx_kwargs)(
+                extended_coord,
+                extended_atype,
+                extended_spin,
+                nlist,
+                mapping,
+                fparam,
+                aparam,
+            )
+        finally:
+            del backbone.need_sorted_nlist_for_lower
+        return traced
 
     def forward_common_lower(
         self, *args: Any, **kwargs: Any

deepmd/pt_expt/train/training.py

@@ -945,7 +945,6 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None:
                         n_attn,
                         task_key,
                     )
-
             inp, _ = self.get_data(is_train=True, task_key=task_key)
             coord = inp["coord"].detach()
             atype = inp["atype"].detach()