fix(pt): stabilize accelerated execution paths

Author: OutisLi

deepmd/pt/entrypoints/freeze_pt2.py

@@ -23,6 +23,7 @@
     annotations,
 )
 
+import contextlib
 import ctypes
 import json
 import logging
@@ -33,12 +34,18 @@
     deepcopy,
 )
 from typing import (
+    TYPE_CHECKING,
     Any,
 )
 
 import numpy as np
 import torch
 
+if TYPE_CHECKING:
+    from collections.abc import (
+        Iterator,
+    )
+
 from deepmd.dpmodel.utils.nlist import (
     build_neighbor_list,
     extend_coord_with_ghosts,
@@ -857,6 +864,7 @@ def _export_with_comm_artifact(
 # it defaults to exact float32.
 _FREEZE_KERNEL_LEVELS = {"DP_TRITON_INFER": "2", "DP_CUDA_INFER": "1"}
 _FREEZE_DISABLED_LEVELS = {"DP_CUTILE_INFER": "0", "DP_CUTE_INFER": "0"}
+_INFER_KERNEL_LEVELS = tuple(_FREEZE_KERNEL_LEVELS | _FREEZE_DISABLED_LEVELS)
 
 
 def _apply_kernel_level_defaults(target_device: torch.device) -> None:
@@ -869,12 +877,7 @@ def _apply_kernel_level_defaults(target_device: torch.device) -> None:
     Python-only eager backends and are disabled for every frozen archive.
     """
     if target_device.type != "cuda":
-        for name in (
-            "DP_TRITON_INFER",
-            "DP_CUDA_INFER",
-            "DP_CUTILE_INFER",
-            "DP_CUTE_INFER",
-        ):
+        for name in _INFER_KERNEL_LEVELS:
             os.environ[name] = "0"
         log.info("Freezing for CPU with accelerator-only DPA4 paths disabled.")
         return
@@ -892,6 +895,21 @@ def _apply_kernel_level_defaults(target_device: torch.device) -> None:
     )
 
 
+@contextlib.contextmanager
+def _kernel_level_defaults(target_device: torch.device) -> Iterator[None]:
+    """Apply freeze-time kernel levels without changing the caller's environment."""
+    saved = {name: os.environ.get(name) for name in _INFER_KERNEL_LEVELS}
+    try:
+        _apply_kernel_level_defaults(target_device)
+        yield
+    finally:
+        for name, value in saved.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value
+
+
 def freeze_sezm_to_pt2(
     ckpt_path: str,
     out_path: str,
@@ -929,14 +947,31 @@ def freeze_sezm_to_pt2(
     are ``DP_TRITON_INFER=2`` and ``DP_CUDA_INFER=1``, which is the fastest
     combination that keeps every operator in exact float32.
     """
+    target_device = device if device is not None else DEVICE
+    with _kernel_level_defaults(target_device):
+        _freeze_sezm_to_pt2(
+            ckpt_path,
+            out_path,
+            target_device=target_device,
+            head=head,
+            atomic_virial=atomic_virial,
+        )
+
+
+def _freeze_sezm_to_pt2(
+    ckpt_path: str,
+    out_path: str,
+    *,
+    target_device: torch.device,
+    head: str | None,
+    atomic_virial: bool,
+) -> None:
+    """Build one AOTInductor archive under an established kernel policy."""
     from torch._inductor import (
         aoti_compile_and_package,
     )
     from torch._inductor import config as inductor_config
 
-    target_device = device if device is not None else DEVICE
-    _apply_kernel_level_defaults(target_device)
-
     raw = torch.load(ckpt_path, map_location="cpu", weights_only=False)
     state_dict, params = _extract_state_and_params(raw)
     state_dict, params = _select_model_head(state_dict, params, head)

deepmd/pt/model/descriptor/env_mat.py

@@ -66,6 +66,7 @@ def prod_env_mat(
     radial_only: bool = False,
     protection: float = 0.0,
     use_exp_switch: bool = False,
+    training: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Generate smooth environment matrix from atom coordinates and other context.
 
@@ -79,19 +80,26 @@ def prod_env_mat(
     - radial_only: Whether to return a full description or a radial-only descriptor.
     - protection: Protection parameter to prevent division by zero errors during calculations.
     - use_exp_switch: Whether to use the exponential switch function.
+    - training: Whether the caller is in training mode. Training uses the eager
+      formulation because force losses require higher-order differentiation.
 
     Returns
     -------
     - env_mat: Shape is [nframes, natoms[1]*nnei*4].
     """
     # Opt-in inference (``DP_TRITON_INFER >= 1``, CUDA): the fused Triton kernel
     # forms the environment matrix in one node-parallel pass and carries a
-    # closed-form backward for the force path.  Training (level 0) and the CPU
-    # path keep the dense autograd chain below, which supports higher-order
+    # closed-form backward for the force path. Training and the CPU path keep
+    # the dense autograd chain below, which supports higher-order
     # differentiation.  The block is nested under ``torch.jit.is_scripting`` so
     # the whole (non-scriptable) branch is pruned under ``torch.jit.script``.
     if not torch.jit.is_scripting():
-        if TRITON_AVAILABLE and triton_infer_level() >= 1 and extended_coord.is_cuda:
+        if (
+            not training
+            and TRITON_AVAILABLE
+            and triton_infer_level() >= 1
+            and extended_coord.is_cuda
+        ):
             return _env_mat_triton(
                 extended_coord,
                 nlist,

deepmd/pt/model/descriptor/repflows.py

@@ -472,6 +472,7 @@ def forward(
             self.e_rcut_smth,
             protection=self.env_protection,
             use_exp_switch=self.use_exp_switch,
+            training=self.training,
         )
         nlist_mask = nlist != -1
         sw = torch.squeeze(sw, -1)
@@ -492,6 +493,7 @@ def forward(
             self.a_rcut_smth,
             protection=self.env_protection,
             use_exp_switch=self.use_exp_switch,
+            training=self.training,
         )
         a_nlist_mask = a_nlist != -1
         a_sw = torch.squeeze(a_sw, -1)

deepmd/pt/model/descriptor/repformers.py

@@ -429,6 +429,7 @@ def forward(
             self.rcut,
             self.rcut_smth,
             protection=self.env_protection,
+            training=self.training,
         )
         nlist_mask = nlist != -1
         sw = torch.squeeze(sw, -1)

deepmd/pt/model/descriptor/se_a.py

@@ -779,6 +779,7 @@ def forward(
             self.rcut,
             self.rcut_smth,
             protection=self.env_protection,
+            training=self.training,
         )
 
         dmatrix = dmatrix.view(-1, self.nnei, 4)

deepmd/pt/model/descriptor/se_atten.py

@@ -781,6 +781,7 @@ def forward(
             self.rcut,
             self.rcut_smth,
             protection=self.env_protection,
+            training=self.training,
         )
         # nb x nloc x nnei
         exclude_mask = self.emask(nlist, extended_atype)

deepmd/pt/model/descriptor/se_r.py

@@ -488,6 +488,7 @@ def forward(
             self.rcut_smth,
             True,
             protection=self.env_protection,
+            training=self.training,
         )
 
         assert self.filter_layers is not None

deepmd/pt/model/descriptor/se_t.py

@@ -840,6 +840,7 @@ def forward(
             self.rcut,
             self.rcut_smth,
             protection=self.env_protection,
+            training=self.training,
         )
         dmatrix = dmatrix.view(-1, self.nnei, 4)
         nfnl = dmatrix.shape[0]

deepmd/pt/model/descriptor/se_t_tebd.py

@@ -936,6 +936,7 @@ def forward(
             self.rcut,
             self.rcut_smth,
             protection=self.env_protection,
+            training=self.training,
         )
         # dmatrix: [1/r, dx/r^2, dy/r^2, dz/r^2], sw: distance weighting
         # nb x nloc x nnei

deepmd/pt/model/descriptor/sezm_nn/wignerd.py

@@ -22,6 +22,9 @@
 
 import torch
 import torch.nn as nn
+from packaging.version import (
+    Version,
+)
 
 from deepmd.pt.utils import (
     env,
@@ -39,6 +42,8 @@
     nvtx_range,
 )
 
+_TORCH_RELEASE = Version(torch.__version__).release[:2]
+
 
 class CaseCoefficients(nn.Module):
     """
@@ -418,6 +423,7 @@ def __init__(
         self.dtype = dtype
         self.device = env.DEVICE
         self.eps = float(eps)
+        self._materialize_inverse_rotation = _TORCH_RELEASE == (2, 11)
         self.dim_full = (self.lmax + 1) ** 2
         self.poly_lmin = 11
         self.poly_offset = self.poly_lmin * self.poly_lmin
@@ -627,6 +633,12 @@ def forward(
         # Consumers address the inverse rotation through explicit strides or
         # PyTorch strided operators, so the transpose can share D_full's storage.
         Dt_full = D_full.transpose(-1, -2)
+        if self._materialize_inverse_rotation:
+            # PyTorch 2.11 Inductor cannot safely lower this escaping transpose
+            # view after the slice assignments that assemble D_full. The
+            # materialized layout keeps the compiled graph semantically
+            # identical; later releases retain the shared-storage view.
+            Dt_full = Dt_full.contiguous()
         return D_full, Dt_full
 
     def forward_zonal(