fix charge_spin & limit torch to 2.11 for compile

Author: OutisLi

deepmd/pt/entrypoints/freeze_pt2.py

@@ -354,34 +354,21 @@ def _make_sample_inputs(
     )
     charge_spin = None
     if dim_chg_spin > 0:
-        default_chg_spin = model.get_default_chg_spin()
-        if default_chg_spin is None:
-            raise ValueError(
-                "SeZM .pt2 freeze requires default_chg_spin when charge/spin "
-                "conditioning is enabled; runtime charge_spin input is not exposed."
-            )
-        charge_spin = (
-            default_chg_spin.to(device=device, dtype=torch.float64)
-            .view(1, dim_chg_spin)
-            .expand(nframes, -1)
-            .contiguous()
+        charge_spin = torch.zeros(
+            nframes, dim_chg_spin, dtype=torch.float64, device=device
         )
     if has_spin:
-        if charge_spin is not None:
-            return (
-                ext_coord,
-                ext_atype,
-                ext_spin,
-                nlist_t,
-                mapping_t,
-                fparam,
-                aparam,
-                charge_spin,
-            )
-        return ext_coord, ext_atype, ext_spin, nlist_t, mapping_t, fparam, aparam
-    if charge_spin is not None:
-        return ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin
-    return ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam
+        return (
+            ext_coord,
+            ext_atype,
+            ext_spin,
+            nlist_t,
+            mapping_t,
+            fparam,
+            aparam,
+            charge_spin,
+        )
+    return ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin
 
 
 def _resolve_nframes(
@@ -446,6 +433,9 @@ def _build_dynamic_shapes(
     nloc_dim = torch.export.Dim("nloc", min=1)
     fparam = sample_inputs[5] if has_spin else sample_inputs[4]
     aparam = sample_inputs[6] if has_spin else sample_inputs[5]
+    charge_spin = None
+    if has_charge_spin:
+        charge_spin = sample_inputs[7] if has_spin else sample_inputs[6]
     if has_spin:
         shapes = (
             {0: nframes_dim, 1: nall_dim},  # extended_coord
@@ -457,7 +447,7 @@ def _build_dynamic_shapes(
             {0: nframes_dim, 1: nloc_dim} if aparam is not None else None,
         )
         if has_charge_spin:
-            shapes = (*shapes, {0: nframes_dim})
+            shapes = (*shapes, {0: nframes_dim} if charge_spin is not None else None)
         return shapes
     shapes = (
         {0: nframes_dim, 1: nall_dim},  # extended_coord: (nframes, nall, 3)
@@ -468,7 +458,7 @@ def _build_dynamic_shapes(
         {0: nframes_dim, 1: nloc_dim} if aparam is not None else None,
     )
     if has_charge_spin:
-        shapes = (*shapes, {0: nframes_dim})
+        shapes = (*shapes, {0: nframes_dim} if charge_spin is not None else None)
     return shapes
 
 
@@ -527,10 +517,48 @@ def freeze_sezm_to_pt2(
     # do_atomic_virial=True pulls every key that DeepPotPTExpt may read
     # (energy, energy_redu, energy_derv_r, energy_derv_c, energy_derv_c_redu)
     # into the traced graph.
-    traced = model.forward_common_lower_exportable(
-        *sample_inputs_cpu,
-        do_atomic_virial=True,
-    )
+    if is_spin:
+        (
+            ext_coord,
+            ext_atype,
+            ext_spin,
+            nlist_t,
+            mapping_t,
+            fparam,
+            aparam,
+            charge_spin,
+        ) = sample_inputs_cpu
+        traced = model.forward_common_lower_exportable(
+            ext_coord,
+            ext_atype,
+            ext_spin,
+            nlist_t,
+            mapping_t,
+            fparam=fparam,
+            aparam=aparam,
+            charge_spin=charge_spin,
+            do_atomic_virial=True,
+        )
+    else:
+        (
+            ext_coord,
+            ext_atype,
+            nlist_t,
+            mapping_t,
+            fparam,
+            aparam,
+            charge_spin,
+        ) = sample_inputs_cpu
+        traced = model.forward_common_lower_exportable(
+            ext_coord,
+            ext_atype,
+            nlist_t,
+            mapping_t,
+            fparam=fparam,
+            aparam=aparam,
+            charge_spin=charge_spin,
+            do_atomic_virial=True,
+        )
 
     # Output key order is taken from a concrete run; Python dict order
     # is stable and matches what DeepPotPTExpt::extract_outputs zips

deepmd/pt/model/model/sezm_model.py

@@ -376,6 +376,9 @@
 from einops import (
     rearrange,
 )
+from packaging.version import (
+    Version,
+)
 from torch.fx.experimental.proxy_tensor import (
     make_fx,
 )
@@ -475,6 +478,16 @@ def _parse_optional_env_bool(var_name: str) -> bool | None:
     )
 
 
+def _check_compile_torch_version() -> None:
+    """Fail fast when SeZM compile is requested on unsupported PyTorch."""
+    version = Version(torch.__version__).release
+    if len(version) < 2 or version[:2] != (2, 11):
+        raise RuntimeError(
+            "SeZM `use_compile` and `DP_COMPILE_INFER` require PyTorch 2.11.x; "
+            f"found torch {torch.__version__}."
+        )
+
+
 def _strip_saved_tensor_detach(gm: torch.fx.GraphModule) -> None:
     """Strip ``aten.detach`` nodes that ``make_fx`` inserts for saved tensors.
 
@@ -614,6 +627,8 @@ def __init__(
         self._env_use_compile_infer: bool | None = _parse_optional_env_bool(
             "DP_COMPILE_INFER"
         )
+        if self.use_compile or self._env_use_compile_infer is True:
+            _check_compile_torch_version()
 
         # === Bridging (optional short-range zone bridging) ===
         self.bridging_method: str = str(bridging_method).upper()
@@ -1828,8 +1843,9 @@ def forward_common_lower_exportable(
         mapping: torch.Tensor | None = None,
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
-        do_atomic_virial: bool = False,
         charge_spin: torch.Tensor | None = None,
+        *,
+        do_atomic_virial: bool = False,
     ) -> torch.nn.Module:
         """Trace ``forward_common_lower`` into an exportable FX ``GraphModule``.
 
@@ -1884,9 +1900,8 @@ def fn(
             mapping_: torch.Tensor | None,
             fparam_: torch.Tensor | None,
             aparam_: torch.Tensor | None,
-            *maybe_charge_spin: torch.Tensor | None,
+            charge_spin_: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
-            charge_spin_ = maybe_charge_spin[0] if maybe_charge_spin else None
             return lower_fn(
                 ext_coord,
                 ext_atype,
@@ -1905,7 +1920,7 @@ def fn(
                 dtype=extended_coord.dtype,
                 device=extended_coord.device,
             )
-            trace_inputs = (*trace_inputs, charge_spin)
+        trace_inputs = (*trace_inputs, charge_spin)
 
         return self._trace_lower_exportable(
             fn,

deepmd/pt/model/model/sezm_spin_model.py

@@ -301,8 +301,9 @@ def forward_common_lower_exportable(
         mapping: torch.Tensor | None = None,
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
-        do_atomic_virial: bool = False,
         charge_spin: torch.Tensor | None = None,
+        *,
+        do_atomic_virial: bool = False,
     ) -> torch.nn.Module:
         """Trace the spin lower interface into an exportable FX graph."""
         extra_sort = self.need_sorted_nlist_for_lower()
@@ -339,9 +340,8 @@ def fn(
             mapping_: torch.Tensor | None,
             fparam_: torch.Tensor | None,
             aparam_: torch.Tensor | None,
-            *maybe_charge_spin: torch.Tensor | None,
+            charge_spin_: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
-            charge_spin_ = maybe_charge_spin[0] if maybe_charge_spin else None
             return lower_fn(
                 ext_coord,
                 ext_atype,
@@ -369,7 +369,7 @@ def fn(
                 dtype=extended_coord.dtype,
                 device=extended_coord.device,
             )
-            trace_inputs = (*trace_inputs, charge_spin)
+        trace_inputs = (*trace_inputs, charge_spin)
 
         return self._trace_lower_exportable(
             fn,

deepmd/pt_expt/infer/deep_eval.py

@@ -1167,9 +1167,8 @@ def _eval_model(
             mapping_t,
             fparam_t,
             aparam_t,
+            charge_spin_t,
         )
-        if charge_spin_t is not None:
-            model_inputs = (*model_inputs, charge_spin_t)
         if self._is_pt2:
             # AOTInductor's __call__ unflattens output using stored out_spec,
             # returning a dict just like the .pte module.
@@ -1320,9 +1319,8 @@ def _eval_model_spin(
             mapping_t,
             fparam_t,
             aparam_t,
+            charge_spin_t,
         )
-        if charge_spin_t is not None:
-            model_inputs = (*model_inputs, charge_spin_t)
         if self._is_pt2:
             model_ret = self._pt2_runner(*model_inputs)
         else:

source/tests/pt/model/test_sezm_export.py

@@ -174,7 +174,7 @@ def _eager_forward(
     sample_inputs: tuple,
 ) -> dict[str, torch.Tensor]:
     """Mirror the trace closure: fresh leaf coord + ``requires_grad=True``."""
-    ext_coord, ext_atype, nlist, mapping, fparam, aparam = sample_inputs
+    ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin = sample_inputs
     eager_coord = ext_coord.detach().clone().requires_grad_(True)
     return model.forward_common_lower(
         eager_coord,
@@ -183,6 +183,7 @@ def _eager_forward(
         mapping=mapping,
         fparam=fparam,
         aparam=aparam,
+        charge_spin=charge_spin,
         do_atomic_virial=True,
         extra_nlist_sort=model.need_sorted_nlist_for_lower(),
     )